一文搞定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