一文搞定Pytorch快速入门

1 简介

Pytorch 是 Facebook 开发的 DL 开源框架，目前学术界运用广泛，本文将通过 MNIST 手写数字识别案例分享 Pytorch 的常用操作。

2 CNN

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision  
import matplotlib.pyplot as plt

torch.manual_seed(1)    

# 超参数
EPOCH = 1           
BATCH_SIZE = 50
LR = 0.001     
DOWNLOAD_MNIST = True  

# 训练集
train_data = torchvision.datasets.MNIST(
    root='./mnist/',    # 文件路径
    train=True,  
    transform=torchvision.transforms.ToTensor(),    # 转换数据格式：C x H x W, 训练时归一化
    download=DOWNLOAD_MNIST,    
)

# 测试集
test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)

# 批训练：(50, 1, 28, 28)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

# 测试前2000个样本
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000]/255.  
test_y = test_data.test_labels[:2000]

# 模型
class CNN(nn.Module):
    def __init__(self):
    
        super(CNN, self).__init__()
        
        self.conv1 = nn.Sequential(  # 输入：(1, 28, 28)
        	# 卷积层
            nn.Conv2d(
                in_channels=1,      # 输入通道数
                out_channels=16,    # 卷积核数目
                kernel_size=5,      # 卷积核大小
                stride=1,          
                padding=2,   
            ),     
            # **层
            nn.ReLU(),  
            # 池化层
            nn.MaxPool2d(kernel_size=2),   
        ) # 输出：(16, 14, 14)
        
        self.conv2 = nn.Sequential(  # 输入：(16, 14, 14)
            nn.Conv2d(16, 32, 5, 1, 2), 
            nn.ReLU(),  
            nn.MaxPool2d(2),  
        ) # 输出：(32, 7, 7)
        
        self.out = nn.Linear(32 * 7 * 7, 10)  # 全连接层，输出：10

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)   # 展平成 (batch_size, 32 * 7 * 7=1568)
        output = self.out(x)
        return output

cnn = CNN()

# 训练模型
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)   # Adam优化器
loss_func = nn.CrossEntropyLoss()   # 交叉熵损失
for epoch in range(EPOCH):
    for step, (b_x, b_y) in enumerate(train_loader):   
        output = cnn(b_x)              
        loss = loss_func(output, b_y)  
        optimizer.zero_grad()          
        loss.backward()                
        optimizer.step()              

Epoch: 0 | train loss: 0.0306 | test accuracy: 0.97
Epoch: 0 | train loss: 0.0147 | test accuracy: 0.98
Epoch: 0 | train loss: 0.0427 | test accuracy: 0.98
Epoch: 0 | train loss: 0.0078 | test accuracy: 0.98

3 RNN

import torch
from torch import nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

torch.manual_seed(1)    

EPOCH = 1      
BATCH_SIZE = 64
TIME_STEP = 28      # 图片高度
INPUT_SIZE = 28     # 图片宽度
LR = 0.01           
DOWNLOAD_MNIST = True  

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

# 测试集
test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)

# 批训练：(50, 1, 28, 28)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

# 测试前2000个样本
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000]/255.  
test_y = test_data.test_labels[:2000]

# 模型
class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.LSTM(     
            input_size=28,      
            hidden_size=64,     # 隐藏单元
            num_layers=1,     
            batch_first=True,   
        )

        self.out = nn.Linear(64, 10)    # 输出层

    def forward(self, x):
        # 输入：(batch, time_step, input_size)
        # 输出：(batch, time_step, output_size)
        # 隐藏分状态h_n ：(n_layers, batch, hidden_size) 
        # 隐藏主状态h_c ：(n_layers, batch, hidden_size)
        r_out, (h_n, h_c) = self.rnn(x, None)  
        out = self.out(r_out[:, -1, :])
        return out

rnn = RNN()

# 训练模型
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   
loss_func = nn.CrossEntropyLoss()   
for epoch in range(EPOCH):
    for step, (x, b_y) in enumerate(train_loader):   
        b_x = x.view(-1, 28, 28)   # 数据转格式：(batch, time_step, input_size)

        output = rnn(b_x)             
        loss = loss_func(output, b_y)   
        optimizer.zero_grad()          
        loss.backward()                
        optimizer.step()             

Epoch: 0 | train loss: 0.0945 | test accuracy: 0.94
Epoch: 0 | train loss: 0.0984 | test accuracy: 0.94
Epoch: 0 | train loss: 0.0332 | test accuracy: 0.95
Epoch: 0 | train loss: 0.1868 | test accuracy: 0.96