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: