LeNet-5
- 导包
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms- 定义自定义图像数据集
class CustomImageDataset(Dataset):
def __init__(self, main_dir, transform=None):
self.main_dir = main_dir
self.transform = transform
self.files = []
self.labels = []
self.label_to_index = {}定义一个自定义的图像数据集类，继承自Dataset，初始化方法，接收主目录和转换方法，主目录，包含多个子目录，每个子目录包含同一类别的图像，图像转换方法，用于对图像进行预处理，存储所有图像文件的路径，存储所有图像的标签，创建一个字典，用于将标签映射到索引
for index, label in enumerate(os.listdir(main_dir)):
self.label_to_index[label] = index
label_dir = os.path.join(main_dir, label)
if os.path.isdir(label_dir):
for file in os.listdir(label_dir):
self.files.append(os.path.join(label_dir, file))
self.labels.append(label)遍历主目录中的所有子目录，将标签映射到索引，构建标签子目录的路径，如果是目录，遍历目录中的所有文件，将文件路径添加到列表，将标签添加到列表
def __len__(self):
return len(self.files)定义数据集的长度，返回文件列表的长度
def __getitem__(self, idx):
image = Image.open(self.files[idx])
label = self.labels[idx]
if self.transform:
image = self.transform(image)
return image, self.label_to_index[label]定义获取数据集中单个样本的方法，打开图像文件，获取图像的标签，如果有转换方法，对图像进行转换，返回图像和对应的标签索引
- 定义数据转换
transform = transforms.Compose([
transforms.Resize((32, 32)), # LeNet-5的输入图像大小
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5]), # LeNet-5的标准化
])- 创建数据集
dataset = CustomImageDataset(main_dir="D:\图像处理、深度学习\cat.dog", transform=transform)- 创建数据加载器
train_data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
test_data_loader = DataLoader(dataset, batch_size=32, shuffle=False)- 定义LeNet-5模型
class LeNet5(nn.Module):
def __init__(self, num_classes):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, num_classes)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x- 实例化模型
num_classes = len(dataset.label_to_index)
lenet_model = LeNet5(num_classes)- 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(lenet_model.parameters(), lr=0.001)- 如果有多个GPU，可以使用nn.DataParallel来并行化模型
if torch.cuda.device_count() > 1:
lenet_model = nn.DataParallel(lenet_model)- 将模型发送到GPU（如果可用）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
lenet_model.to(device)- 训练模型
num_epochs = 10
for epoch in range(num_epochs):
lenet_model.train()
running_loss = 0.0
for images, labels in train_data_loader:
images, labels = images.to(device), labels.to(device)前向传播
反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_data_loader):.4f}')在每个周期结束时评估模型
lenet_model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_data_loader:
images, labels = images.to(device), labels.to(device)
outputs = lenet_model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')- 保存训练好的模型
torch.save(lenet_model.state_dict(), "D:\图像处理、深度学习\训练保存\lenet_model.pth")- 导包
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms- 定义LeNet-5模型
class LeNet5(nn.Module):
def __init__(self, num_classes):
super(LeNet5, self).__init__()
self.convnet = nn.Sequential(
nn.Conv2d(3, 6, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(6, 16, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
self.fc = nn.Sequential(
nn.Linear(16 * 5 * 5, 120),
nn.ReLU(),
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, num_classes)
)
def forward(self, x):
x = self.convnet(x)
x = x.view(x.size(0), -1) # 展平多维卷积层输出
x = self.fc(x)
return x- 定义数据转换
transform = transforms.Compose([
transforms.Resize((32, 32)), # LeNet-5的输入图像大小
transforms.ToTensor(),
])- 假设您的数据集是一个ImageFolder格式，并且路径为 "path_to_your_dataset"
dataset = ImageFolder(root="D:\图像处理、深度学习\cat.dog", transform=transform)- 创建数据加载器
data_loader = DataLoader(dataset, batch_size=32, shuffle=True)- 获取数据集中的类别数
num_classes = len(dataset.classes)- 创建LeNet-5模型实例
lenet_model = LeNet5(num_classes=num_classes)- 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(lenet_model.parameters(), lr=0.001)- 将模型发送到GPU（如果可用）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
lenet_model.to(device)
- 训练模型
num_epochs = 10
for epoch in range(num_epochs):
lenet_model.train()
running_loss = 0.0
correct = 0
total = 0
for images, labels in data_loader:
images, labels = images.to(device), labels.to(device)
# 前向传播
outputs = lenet_model(images)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(data_loader):.4f}, Accuracy: {100 * correct / total:.2f}%')
