# Convolutional neural network (two convolutional layers)classConvNet(nn.Module):def__init__(self, num_classes=10):super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1,16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(16,32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
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)# 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+1, num_epochs, i+1, total_step, loss.item()))
# 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))
# Save the model checkpoint
torch.save(model.state_dict(),'model.ckpt')
from __future__ import print_functionimport torch import torch.nn as nnimport torchvisionimport torchvision.transforms as transforms# Device configurationdevice = torch.device('cuda:0' if torc...