pytorch小例子mnist

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torch.autograd import Variable

from torchvision import transforms, datasets

batch_size = 64

train_data = datasets.MNIST(root='./data/', train=True, transform=transforms.ToTensor(), download= True)

test_data = datasets.MNIST(root= './data/', train=True, transform=transforms.ToTensor(), download= True)

train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size= batch_size, shuffle= True)

test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size= batch_size, shuffle=True)

class Net(nn.Module):

    def __init__(self):

        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)

        self.conv2 = nn.Conv2d(10,20,5)

        self.conv3 = nn.Conv2d(20,40,3)

        self.maxpool = nn.MaxPool2d(2)

        self.fc = nn.Linear(40,10)

    def forward(self,x):

        input_channel=x.size(0)

        x = F.relu(self.maxpool(self.conv1(x)))

        x = F.relu(self.maxpool(self.conv2(x)))

        x = F.relu(self.maxpool(self.conv3(x)))

        x = x.view(input_channel, -1)

        x = self.fc(x)

        return F.log_softmax(x, dim=1)

model= Net()

print(model)

optimizer = optim.SGD(model.parameters(), lr= 0.01, momentum= 0.5)

def train(epoch):

    for index, (data, target) in enumerate (train_loader):   # data:[64,1,1,28] #target:[64]

        output= model(data)

        loss= F.nll_loss(output, target)

        if index % 200 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, index * len(data), len(train_loader.dataset), 100. * index/len(train_loader), loss.item()))

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

for epoch in range(1, 10):

    train(epoch)

    print("done")