PyTorch is an open source Python machine learning library that is widely used in the field of deep learning. This article will introduce how to use PyTorch to implement a simple neural network, and train and test it on the MNIST dataset.
- Environmental preparation
First, you need to install PyTorch and related dependencies. PyTorch can be installed with the following command:
pip install torch torchvision
- Dataset preparation
We will use the MNIST handwritten digits dataset, which is a very classic dataset with 60,000 training samples and 10,000 test samples. You can use the following code to download the data set: Follow v❤Public H: Ai Technology Planet Reply (123) must get pytorch deep learning materials
- data preprocessing
Before using the dataset, it needs to be preprocessed. In this example, we will scale each digital image to 28x28 size and normalize the pixel values between 0 and 1. Preprocessing can be done with the following code:
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.Resize((28, 28)),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_dataset.transform = transform
test_dataset.transform = transform
- Define the neural network model
We will use a simple multi-layer perceptron (MLP) model to classify the MNIST dataset. The model has two hidden layers, each with 128 neurons, and uses the ReLU activation function. The output layer has 10 neurons corresponding to the digits from 0 to 9. A model can be defined with the following code:
import torch.nn as nn
class MLP(nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 128)
self.fc3 = nn.Linear(128, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = x.view(-1, 28*28)
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.fc3(x)
return x
- Define loss function and optimizer
We will use a cross-entropy loss function and a stochastic gradient descent (SGD) optimizer. Loss functions and optimizers can be defined using the following code:
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
- training model
Now we can start training the model. The model can be trained using the following code:
num_epochs = 10
After you have finished building the neural network, you need to train it with data. Here is a simple training example:
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 训练网络
for epoch in range(10): # 循环遍历数据集多次
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# 获取输入
inputs, labels = data
# 梯度清零
optimizer.zero_grad()
# 正向传播、反向传播、优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 统计损失
running_loss += loss.item()
if i % 2000 == 1999: # 每2000批次打印一次平均损失
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
In this example, we use the cross-entropy loss function and the stochastic gradient descent (SGD) optimizer for training. We loop through the dataset 10 times and do forward pass, back pass and optimize in each batch. Every 2000 batches, we print the average loss. When training is complete, we print out the "Finished Training" message.
Finally, we can use the test set to evaluate our network:
# 在测试集上测试网络
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (
100 * correct / total))
In this example, we loop through the test set and run our network on each image. We use the torch.max() function to get the maximum value and corresponding index for each output tensor. Then, we count the number of correct classifications and calculate the accuracy. Finally, we print out the accuracy.
This is an example of a basic implementation of a neural network in PyTorch. As you learn deeper, you can explore more complex network structures, different optimizers and loss functions, and more advanced features such as automatic differentiation and parallelized training.