import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms # Device configuration device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Hyper parameters num_epochs = 5 num_classes = 10 batch_size = 100 learning_rate = 0.001 # MNIST dataset train_dataset = torchvision.datasets.MNIST(root='../../data/', train=True, transform=transforms.ToTensor(), download=True) test_dataset = torchvision.datasets.MNIST(root='../../data/', train=False, transform=transforms.ToTensor()) # Data loader #dataloader 已经在一个数据前又加了一个维度，batch_size维度，表示每批训练用的照片数目，乘上批次数目等于data总数. #而epoch指的是迭代数，一共迭代5次，每批处理100个图片，每迭代一次都要讲train_data 走完，所以计算的批次数 train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # Convolutional neural network (two convolutional layers) class ConvNet(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), #1代表只有一个通道数，16代表有16个卷积神经元，每一个卷积神经元由一个kernel组成，也就是说有16个kernel， #每个kernel的大小是5， 步长为1， padding = 2说明每一侧有两条边补0，(M-F+1)/S表示每一维有多少个0，这里计算得到4 #所以左边2右边2，上边2下边2，所以padding要设为2 #1个通道的图片通过16个kernel卷积变换后得到16个通道的图片，这里的1可以指二维数据的input_size,第二个就相当于二维的output_size, #只是每一个二维数据都是一个二维张量，整体看待 nn.BatchNorm2d(16),#对卷积得到的16个二维张量进行一个BatchNormalize nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) #图片长宽缩小一半 14*14 self.layer2 = nn.Sequential( nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),#16个输入神经元，32个输出神经元 nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) #7*7，但通道数变成了32 self.fc = nn.Linear(7*7*32, num_classes) #最表层fc,输出一个长度为10的值，可以代表属于这十个类别的概率，输入值就是将这个 def forward(self, x): out = self.layer1(x) out = self.layer2(out) #print('out.size:{}'.format(out.size())) #原来out的size为100，32，7，7 out = out.reshape(out.size(0), -1) #第一个维度长其实就等于batch_size的大小，将输出转为batch_size个链数据， #相当于把每一张32通道的输出图片都摊平，形状变化对卷积会产生影响，但是对回归不产生影响，之后再加上一个线性回归层 out = self.fc(out) return out model = ConvNet(num_classes).to(device) # Loss and optimizer criterion = nn.CrossEntropyLoss() #交叉熵：先对out分布归一化(转化为概率分布，一般用softmax函数)，根据实际值计算交叉熵 #交叉熵其实表示的 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # Train the model total_step = len(train_loader) for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) # Forward pass outputs = model(images) #if i == 1: # print('outputs.size:{},labels.size:{}'.format(outputs.size(),labels.size())) loss = criterion(outputs, labels) #交叉熵需要预测为概率值（每一张图片的概率分布）一批有100张，因此对于每一张图片要有num_classes个表示每一类概率的值 #outputs为概率分布，outputs.size: 100x10 ; labels.size: 10 #与 # 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(): #因为没有优化器，所以得自己将grad设为无，也就是梯度参数不保存，不然每次调用model函数都会产生grad参数，浪费内存 correct = 0 total = 0 for 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')
学习这段代码对卷积神经网络的建立很有帮助。