用PyTorch构建基于卷积神经网络的手写数字识别模型,包括卷积层、激活函数、损失函数等等,使用GPU进行训练,采用精确率、召回率、F1值评测模型表现。给出详细的代码注释。
时间: 2023-12-16 22:05:28 浏览: 41
好的,下面是代码和注释:
```python
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader
from sklearn.metrics import precision_score, recall_score, f1_score
# 定义超参数
batch_size = 64
learning_rate = 0.01
num_epochs = 20
# 判断是否有GPU加速
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载数据集并进行预处理
transform = transforms.Compose([
transforms.ToTensor(), # 将图片转换为Tensor
transforms.Normalize((0.1307,), (0.3081,)) # 标准化
])
# 加载训练集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# 加载测试集
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 定义卷积神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 第一层卷积层,输入通道为1,输出通道为32,卷积核大小为3
self.conv2 = nn.Conv2d(32, 64, kernel_size=3) # 第二层卷积层,输入通道为32,输出通道为64,卷积核大小为3
self.dropout1 = nn.Dropout2d(0.25) # 第一层dropout,丢弃比例为0.25
self.dropout2 = nn.Dropout2d(0.5) # 第二层dropout,丢弃比例为0.5
self.fc1 = nn.Linear(9216, 128) # 第一个全连接层,输入大小为9216,输出大小为128
self.fc2 = nn.Linear(128, 10) # 第二个全连接层,输入大小为128,输出大小为10
def forward(self, x):
x = F.relu(self.conv1(x)) # 第一层卷积层
x = F.relu(self.conv2(x)) # 第二层卷积层
x = F.max_pool2d(x, 2) # 最大池化层,大小为2
x = self.dropout1(x) # 第一层dropout
x = torch.flatten(x, 1) # 将特征展成一维
x = F.relu(self.fc1(x)) # 第一个全连接层
x = self.dropout2(x) # 第二层dropout
x = self.fc2(x) # 第二个全连接层
return x
# 实例化模型并将其移动到GPU上
model = Net().to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9) # 随机梯度下降优化器
# 训练模型
for epoch in range(num_epochs):
train_loss = 0.0
train_correct = 0.0
train_total = 0.0
for i, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device) # 将输入和标签移动到GPU上
optimizer.zero_grad() # 梯度清零
outputs = model(inputs) # 前向传播
loss = criterion(outputs, labels) # 计算损失函数
loss.backward() # 反向传播
optimizer.step() # 更新参数
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
train_total += labels.size(0)
train_correct += (predicted == labels).sum().item()
# 在测试集上评测模型表现
test_loss = 0.0
test_correct = 0.0
test_total = 0.0
y_true = []
y_pred = []
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device) # 将输入和标签移动到GPU上
outputs = model(inputs) # 前向传播
loss = criterion(outputs, labels) # 计算损失函数
test_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
test_total += labels.size(0)
test_correct += (predicted == labels).sum().item()
y_true.extend(labels.cpu().numpy()) # 将标签转换为numpy格式并添加到y_true中
y_pred.extend(predicted.cpu().numpy()) # 将预测结果转换为numpy格式并添加到y_pred中
# 计算精确率、召回率和F1值
precision = precision_score(y_true, y_pred, average='macro')
recall = recall_score(y_true, y_pred, average='macro')
f1 = f1_score(y_true, y_pred, average='macro')
# 打印训练和测试结果
print('Epoch [{}/{}], Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}, Precision: {:.4f}, Recall: {:.4f}, F1: {:.4f}'
.format(epoch + 1, num_epochs, train_loss / len(train_loader), train_correct / train_total, test_loss / len(test_loader), test_correct / test_total, precision, recall, f1))
```
代码注释如下:
1. 导入所需的库和模块。
2. 定义超参数。
3. 判断是否有GPU加速。
4. 加载数据集并进行预处理。
5. 加载训练集。
6. 加载测试集。
7. 定义卷积神经网络模型。
8. 实例化模型并将其移动到GPU上。
9. 定义损失函数和优化器。
10. 训练模型。
11. 在测试集上评测模型表现。
12. 计算精确率、召回率和F1值。
13. 打印训练和测试结果。