import torch import torch.nn as nn import torchvision.datasets as normal_datasets import torchvision.transforms as transforms from torch.autograd import Variablenum_epochs = 1 batch_size = 100 learning_rate = 0.001# 將數據處理成Variable, 如果有GPU, 可以轉成cuda形式 def get_variable(x):x = Variable(x)return x.cuda() if torch.cuda.is_available() else x# 從torchvision.datasets中加載一些常用數據集 train_dataset = normal_datasets.MNIST(root='./mnist/', # 數據集保存路徑train=True, # 是否作為訓練集transform=transforms.ToTensor(), # 數據如何處理, 可以自己自定義download=True) # 路徑下沒有的話, 可以下載# 見數據加載器和batch test_dataset = normal_datasets.MNIST(root='./mnist/',train=False,transform=transforms.ToTensor())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)# 兩層卷積 class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()# 使用序列工具快速構建self.conv1 = nn.Sequential(nn.Conv2d(1, 16, kernel_size=5, padding=2),nn.BatchNorm2d(16),nn.ReLU(),nn.MaxPool2d(2))self.conv2 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=5, padding=2),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2))self.fc = nn.Linear(7 * 7 * 32, 10)def forward(self, x):out = self.conv1(x)out = self.conv2(out)out = out.view(out.size(0), -1) # reshapeout = self.fc(out)return outcnn = CNN() if torch.cuda.is_available():cnn = cnn.cuda()# 選擇損失函數和優化方法 loss_func = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)for epoch in range(num_epochs):for i, (images, labels) in enumerate(train_loader):images = get_variable(images)labels = get_variable(labels)outputs = cnn(images)loss = loss_func(outputs, labels)optimizer.zero_grad()loss.backward()optimizer.step()if (i + 1) % 100 == 0:print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f'% (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size, loss.data[0]))# 測試模型 cnn.eval() # 改成測試形態, 應用場景如: dropout correct = 0 total = 0 for images, labels in test_loader:images = get_variable(images)labels = get_variable(labels)outputs = cnn(images)_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels.data).sum()print(' 測試 準確率: %d %%' % (100 * correct / total))# Save the Trained Model torch.save(cnn.state_dict(), 'cnn.pkl')

總結

以上是生活随笔為你收集整理的pytorch 入门（二） cnn 手写数字识别的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。