classConvNet(nn.Module):def__init__(self, num_classes=10):super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(#输入数据为1*28*28
nn.Conv2d(1,16, kernel_size=5, stride=1, padding=2),# 输入有1个通道,输出有16个通道(16个卷积核进行卷积),padding之后的数据大小为1*(28+2+2)*(28+2+2)=1*32*32,每个卷积核之后的大小为1*(32-5+1)*(32-5+1)=1*28*28,最终得到16*28*28
nn.BatchNorm2d(16),##BatchNorm2d的参数为上一个网络层的channel数:16
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))#16*14*14
self.layer2 = nn.Sequential(
nn.Conv2d(16,32, kernel_size=5, stride=1, padding=2),#padding之后的大小18*18,卷积之后的大小14*14,输出为32*14*14
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))#32*7*7
self.fc = nn.Linear(7*7*32, num_classes)#全连接defforward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0),-1)
out = self.fc(out)return out
model = ConvNet(num_classes).to(device)# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
训练
# Train the model
total_step =len(train_loader)for epoch inrange(num_epochs):for i,(images, labels)inenumerate(train_loader):
images = images.to(device)
labels = labels.to(device)print(images.size())#torch.Size([100(batch_size), 1(channel), 28, 28])# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()if(i +1)%100==0:print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch
测试
Test the model
model.eval()# eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)with torch.no_grad():
correct =0
total =0for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data,1)
total += labels.size(0)
correct +=(predicted == labels).sum().item()print('Test Accuracy of the model on the 10000 test images: {} %'.format(100* correct / total))