|
|
- import torch
- import torch.nn as nn
- import torch.optim as optim
- from torch.utils.data import DataLoader, Dataset
- from PIL import Image
- import os
- import torch.nn.functional as F
- from torchvision import transforms
- # 设置训练超参数
- batch_size = 32#每一轮输入图像的张数。
- learning_rate = 0.001#学习率，每一次梯度下降的程度，学习率太大容易找不到最优价。
- num_epochs = 10#最大epoch数量，这里训练10轮，
- # 定义模型类
- class DigitNet(nn.Module):#从torch的Module模块继承并构建新的网络DigitNet。
- def __init__(self):
- super().__init__()#继承Module的初始化方法。
- self.conv1 = nn.Conv2d(1, 10, 5)#卷积层1，3个参数分别为输入通道数，输出通道数，卷积核大小
- self.conv2 = nn.Conv2d(10, 20, 3)#卷积层2，同上
- self.fc1 = nn.Linear(20 * 10 * 10, 500)#这里的20是conv2输出的通道数，10是由于输入图像大小为28，经过conv1大小变为24，通过max_pool2d大小减半变为12，再通过conv2变为10，因此大小是10*10，输出500维的特征
- self.fc2 = nn.Linear(500, 16)#通过500维的特征输出16个分类
- def forward(self, x):#前向传递过程
- input_size = x.size(0)
- x = self.conv1(x)#卷积层1
- x = F.relu(x)#激活函数，不改变大小
- x = F.max_pool2d(x, 2, 2)#池化，大小减半
- x = self.conv2(x)#卷积层2
- x = F.relu(x)#激活函数
- x = x.view(input_size, -1)#展平操作，用于构建fc层输入
- x = self.fc1(x)#全连接层1
- x = F.relu(x)#激活函数
- x = self.fc2(x)#全连接层2
- output = F.log_softmax(x, dim=1)#softmax分类
- return output#返回输出类别
- # 定义自定义数据集类
- class CustomDataset(Dataset):
- def __init__(self, root_dir, transform=None):
- self.root_dir = root_dir # 根目录路径
- self.classes = os.listdir(root_dir) # 获取根目录下的所有类别（子文件夹）名字
- self.data = [] # 存储数据文件路径
- self.targets = [] # 存储数据对应的标签
- self.transform = transform # 数据预处理的转换操作
- for i, class_name in enumerate(self.classes):
- class_dir = os.path.join(self.root_dir, class_name) # 每个类别的文件夹路径
- file_names = os.listdir(class_dir) # 获取当前类别文件夹下的所有文件名
- for file_name in file_names:
- file_path = os.path.join(class_dir, file_name) # 每个文件的完整路径
- self.data.append(file_path) # 将文件路径添加到data列表中
- self.targets.append(i % 16) # 将类别的索引添加到targets列表中，取余是为了使标签在0-15之间循环
- def __len__(self):
- return len(self.data) # 返回数据集的样本数量
- def __getitem__(self, idx):
- image_path = self.data[idx] # 获取指定索引处的图像路径
- image = Image.open(image_path).convert('L') # 使用PIL库打开图像，并将其转换为灰度图像
- if self.transform is not None:
- image = self.transform(image) # 对图像进行预处理转换
- target = self.targets[idx] # 获取指定索引处的标签
- return image, target # 返回图像和对应的标签
- # 设置数据预处理和转换
- data_transform = transforms.Compose([
- transforms.Resize((28, 28)), # 调整图像大小为 28x28
- transforms.ToTensor(), # 将图像转换为张量
- transforms.Normalize((0.5,), (0.5,)) # 归一化处理
- ])
- # 创建数据加载器
- dataset = CustomDataset('D:/data', transform=data_transform) # 创建自定义数据集实例，并应用数据预处理的转换操作
- dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) # 创建数据加载器，指定批量大小和是否随机打乱数据
- # 创建模型实例
- model = DigitNet() # 创建数字识别模型实例
- # 定义损失函数和优化器
- criterion = nn.CrossEntropyLoss() # 交叉熵损失函数，用于多分类问题
- optimizer = optim.Adam(model.parameters(), lr=learning_rate) # Adam优化器，用于参数优化
- # 设置模型为训练模式
- model.train()
- # 开始训练
- for epoch in range(num_epochs): # 遍历每个epoch
- for images, labels in dataloader: # 遍历每个batch的图像和标签
- # 前向传播
- outputs = model(images) # 将图像输入模型，获取预测结果
- # 计算损失
- loss = criterion(outputs, labels) # 计算预测结果与真实标签之间的损失
- # 反向传播和优化
- optimizer.zero_grad() # 清空梯度
- loss.backward() # 反向传播计算梯度
- optimizer.step() # 更新模型参数
- # 每个epoch结束后打印损失
- print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
- # 保存模型
- torch.save(model.state_dict(), 'digit_model.pth') # 保存模型参数到文件
|
|
发表于 2023-9-8 12:01:58
楼主