根据题目先写读取数据集的代码
# 图像转换
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
])
# 读取数据集
train_dataset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
然后定义网络结构，这里我完全按照题目来设计的网络，建议在用代码实现这部分时，先手画出网络整个流程，下面是我在纸上写的从input到output的整个流程
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()self.net = nn.Sequential(nn.Conv2d(3, 32, kernel_size=5, padding=2,stride=1),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2, padding=0),nn.Conv2d(32, 32, kernel_size=5, padding=2, stride=1),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2, padding=0),nn.Conv2d(32, 64, kernel_size=5, padding=2, stride=1),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2, padding=0))self.fc = nn.Sequential(nn.Linear(64 * 4 * 4, 1024),nn.Linear(1024, 10))def forward(self, x):x = self.net(x)x = x.view(-1, 64 * 4 * 4)x = self.fc(x)return x
# x = torch.rand(128, 3, 32, 32)
model = CNN()
接下来是定义损失函数和优化器
# 损失函数
loss = torch.nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 训练过程
num_epochs = 100
最后流程是训练模型，输出训练损失，测试损失，准确率以及每一批次的训练时间
writer = SummaryWriter(log_dir='logs/fit')
for epoch in range(num_epochs):model.train()model= model.to(device)running_loss, test_loss, correct, start_time = 0.0, 0.0, 0, time.time()for images, labels in train_loader:images = images.to(device)labels = labels.to(device)optimizer.zero_grad()outputs = model(images)l = loss(outputs, labels)# 反向传播计算梯度l.backward()# 更新模型参数optimizer.step()running_loss += l.item()model.eval()with torch.no_grad():for images, labels in test_loader:images = images.to(device)labels = labels.to(device)outputs = model(images)l = loss(outputs,labels)# 测试集损失函数总量test_loss += l.item()_, predicted = torch.max(outputs, 1)correct += (predicted == labels).float().sum().item()avg_loss = running_loss / len(train_loader)avg_test_loss = test_loss / len(test_loader)end_time = time.time()accuracy = correct/len(test_dataset)writer.add_scalar('Loss/train',avg_loss, epoch)writer.add_scalar('Loss/test',avg_test_loss, epoch)writer.add_scalar('Accuracy', accuracy, epoch)print(f'Epoch:{epoch+1}, 'f'Train Loss:{avg_loss:.4f}, 'f'Test Loss:{avg_test_loss:.4f}, 'f'Accuracy:{accuracy:.4f}, 'f'Time:{end_time-start_time:.2f}seconds')
writer.close()
训练结果
tensorboard效果展示
总结
最重要的部分是模型设计部分，建议先在纸上计算一个批次数据通过卷积、池化等部分变化的整个流程，然后再动手写代码会快很多。