Pytorch standardized training process (Mnist classification task)

Preface

I followed my brother to do some tasks in the previous summer vacation, and found myself as a code porter when the school started. For example, the dataloader encapsulation link, I still can't write it myself, so I want to learn pytorch again with the opportunity of pytorch comprehensive practice.
The article refers to the tutorial WHAT IS TORCH.NN REALLY on the pytorch official website ? , Integrate the final mnist classification task process.

text

Step 1: Download the mnist data set

from pathlib import Path
import requests
import pickle
import gzip
DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"
PATH.mkdir(parents = True, exist_ok = True)  # 联级创建

URL = "http://deeplearning.net/data/mnist/"
FILENAME = "mnist.pkl.gz"

if not (PATH/ FILENAME).exists():
    content = requests.get(URL + FILENAME).content
    (PATH / FILENAME).open("wb").write(content)
with gzip.open((PATH / FILENAME).as_posix(),"rb" ) as f:
    ((x_train,y_train),(x_valid,y_valid),_) = pickle.load(f, encoding = "latin-1")

Step 2: Package into dataloader

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from torch import optim
###############
#numpy.array->torch.tensor->tensordataset->dataloader
###############
x_train, y_train, x_valid, y_valid = map(torch.tensor, (x_train, y_train, x_valid, y_valid))
train_ds = TensorDataset(x_train, y_train)
valid_ds = TensorDataset(x_valid, y_valid)
train_dl = DataLoader(train_ds, bs,shuffle = True)
valid_dl = DataLoader(valid_ds, bs*2)
###############
#在dataloader中加入预处理
###############
class WrappedDataLoader:
    def __init__(self,dl):
        self.dl  = dl
#         self.func = func
    
    def preprocess(self,x,y):
         return x.view(-1,1,28,28),y
    
    def __len__(self):
        return len(self.dl)
    
    def __iter__(self):
        batches = iter(self.dl)
        for b in batches:
            yield (self.preprocess(*b))

train_dl, valid_dl = get_data(train_ds, valid_ds, bs = 64)
train_dl = WrappedDataLoader(train_dl)
valid_dl = WrappedDataLoader(valid_dl)

Step 3: Define the model and loss function

###############
# 模型类
###############
class Cnn_model(nn.Module):
    
    def __init__(self):
        super(Cnn_model,self).__init__()
        
#         self.sequence = nn.Sequential(
        self.sequence = nn.Sequential(
        nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
        nn.ReLU(),
        nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
        nn.ReLU(),
        nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
        nn.ReLU(),
        nn.AdaptiveAvgPool2d(1), 
        )
    
    def forward(self,xb):
        xb  = self.sequence(xb)
        xb = xb.view(xb.size(0),-1)
        return xb
###############
# 返回模型，和optim
###############
def get_model():
    model = Cnn_model()
    optimizer = optim.SGD(model.parameters(),lr = 1e-2)
    return model, optimizer
model, optimizer = get_model()
loss_func = F.cross_entropy  # loss_func

Step 4: define the fit process

###############
# 结果的计算/loss的计算/是否要反向传播
###############
def loss_batch(model,loss_func,xb,yb, optimizer= None):
    y_pred = model(xb)
    loss = loss_func(y_pred, yb)
    if optimizer != None:
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    return loss, len(xb) # len(xb)用于计算valid中loss平均值
###############
# fit 
###############
def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)
        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)
        print(epoch, val_loss)

In the end just need:

fit(epochs=10, model=model, loss_func=loss_func, opt=optimizer, train_dl=train_dl, valid_dl=valid_dl)

You can run the result:

to sum up

General techniques used:

• pathlib\requests\pickle\gzip库
• Tensordataset\dataloader, and custom dataloader class (need to define: init\len\iter)
• nn.Sequential and nn.AdaptiveAvgPool2d (parameters depend on the dimension you need to output)
• Define the get_model function, return model and optim, additional definition of loss_fn
• Define the fit function to simplify the training process
• Write the result calculation, loss calculation, and optimization into the loss_batch class, and use the opt parameter to make the train and valid processes available

Next:
Text Segmentation as a Supervised Learning Task model reproduction