import torch from torch.utils.data import DataLoader from torchvision import datasets from torchvision import transforms from torch import nn, optim from torch.nn import functional as F EPOCH = 1000 BATCH_SIZE = 128 LR = 0.001 DOWNLOAD_MNIST = False train_data = datasets.MNIST( root='./mnist', train=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]),#0-255 -> 0-1 download=DOWNLOAD_MNIST ) #plot one example print(train_data.train_data.size()) print(train_data.train_labels.size()) plt.imshow(train_data.train_data[0].numpy(), cmap='gray') plt.title('%i' % train_data.train_labels[0]) plt.show() train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE,\ shuffle=True, num_workers=2) test_data = datasets.MNIST( root='./mnist', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]),#0-255 -> 0-1 download=DOWNLOAD_MNIST ) test_loader = DataLoader(dataset=test_data, batch_size=BATCH_SIZE,\ shuffle=True, num_workers=2) x, label = iter(test_loader).next() #这个iter能把一个batch_size提取出来 print("x:",x.shape, 'label:',label.shape) class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Sequential( # input shape (1, 28, 28) nn.Conv2d( in_channels=1, # input height out_channels=10, # n_filters kernel_size=5, # filter size stride=1, # filter movement/step padding=2, # if want same width and length of this image after Conv2d, padding=(kernel_size-1)/2 if stride=1 ), # output shape (16, 28, 28) nn.ReLU(), # activation nn.MaxPool2d(kernel_size=2), # choose max value in 2x2 area, output shape (16, 14, 14) ) self.conv2 = nn.Sequential( # input shape (16, 14, 14) nn.Conv2d(10, 20, 5, 1, 2), # output shape (32, 14, 14) nn.ReLU(), # activation nn.MaxPool2d(2), # output shape (32, 7, 7) ) self.out1 = nn.Linear(20 * 7 * 7, 512) # fully connected layer, output 10 classes self.out2 = nn.Linear(512, 10) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = x.view(x.size(0), -1) # flatten the output of conv2 to (batch_size, 32 * 7 * 7) output = self.out1(x) output = F.relu(output) output = self.out2(output) return output # return x for visualization DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") cnn = CNN().to(DEVICE) print(cnn) # net architecture optimizer = optim.Adam(cnn.parameters(), lr=LR) # optimize all cnn parameters criteon = nn.CrossEntropyLoss().to(DEVICE) # the target label is not one-hotted def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() if(batch_idx+1)%30 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) def test(model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) pred = output.argmax(dim=1) correct += torch.eq(pred, target).float().sum().item() test_loss += criterion(output, target) # test_loss += F.nll_loss(output, target, reduction='sum').item() # 将一批的损失相加 # pred = output.max(1, keepdim=True)[1] # 找到概率最大的下标 # correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) # training and testing for epoch in range(EPOCH): train(cnn, DEVICE, train_loader, optimizer, epoch) test(cnn, DEVICE, test_loader)
最后能得到99%的准确率
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:Pytorch搭建卷积神经网络用于MNIST分类 - Python技术站