Preface Recently, I need to use pyTorch framework in my learning process. I briefly studied it and wrote a simple case to record the building of an identification network foundation in PyTorch. Change the pyTorch framework to identify the MNIST dataset of TensorFlow written by a blogger, and you can also see in detail the general differences between the two frameworks.

Tensorflow version reproduced source (” 1024 “bugs bunny CSDN blogger) : blog.csdn.net/lzx159951/a…

Pytorch combat mnIST handwritten number recognition

Package to import
import torch
import torch.nn as nn# used to build the network layer
import torch.optim as optim# import optimizer
from torch.utils.data import DataLoaderLoad iterators for the dataset
from torchvision import datasets, transforms# used to load the MNSIT dataset

# Download the dataset

train_set = datasets.MNIST('./data', train=True, download=True,transform = transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1037,), (0.3081,))
              ]))
test_set = datasets.MNIST('./data', train=False, download=True,transform = transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1037,), (0.3081,))))# Build networks (network structure corresponds to the tensorflow article)

class Net(nn.Module):

    def __init__(self, num_classes=10):
        super(Net, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1.32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(32.64, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2,stride=2),

        )
        self.classifier = nn.Sequential(
            nn.Linear(3136.7*7*64),
            nn.Linear(3136, num_classes),

        )

    def forward(self,x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)

        return x
net=Net()
net.cuda()# Run on GPU

# Calculate error, use Adam optimizer to optimize error
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), 1e-2)

train_data = DataLoader(train_set, batch_size=128, shuffle=True)
test_data = DataLoader(test_set, batch_size=128, shuffle=False)



# Training process
for epoch in range(1):
    net.train()  Train (); eval();
    batch = 0

    for batch_images, batch_labels in train_data:

        average_loss = 0
        train_acc = 0

        ## After PyTorch0.4 Variable was combined with tensor, so you don't need Variable encapsulation here
        if torch.cuda.is_available():
            batch_images, batch_labels = batch_images.cuda(),batch_labels.cuda()

        # forward propagation
        out = net(batch_images)
        loss = criterion(out,batch_labels)


        average_loss = loss
        prediction = torch.max(out,1) [1]
        # print(prediction)

        train_correct = (prediction == batch_labels).sum()
        Train_correct is a longtensor and needs to be converted to float

        train_acc = (train_correct.float())  / 128

        optimizer.zero_grad() Empty the gradient information, otherwise it will add up every time you propagate back
        loss.backward()  # Loss Backpropagation
        optimizer.step()  ## Gradient update

        batch+=1
        print("Epoch: %d/%d || batch:%d/%d average_loss: %.3f || train_acc: %.2f"
              %(epoch, 20, batch, float(int(50000/128)), average_loss, train_acc))

Check the effect on the test set
net.eval()  # Change the model to a predictive model
for idx,(im1, label1) in enumerate(test_data):
    if torch.cuda.is_available():
        im, label = im1.cuda(),label1.cuda()
    out = net(im)
    loss = criterion(out, label)

    eval_loss = loss

    pred = torch.max(out,1) [1]
    num_correct = (pred == label).sum()
    acc = (num_correct.float())/ 128
    eval_acc = acc

    print('EVA_Batch:{}, Eval Loss: {:.6f}, Eval Acc: {:.6f}'
      .format(idx,eval_loss , eval_acc))

Copy the code

Running results: