pytorch 完整搭建一个模型的套路

目录

1、导入必要的模块

 2、加载数据集

 3、搭建模型

计算卷积层padding的方法

 4、定义模型运行的设备

5、创建网络模型、损失函数、优化器

6、训练/测试模型

以CIFAR10为数据集，搭建如下结构的神经网络模型

1、导入必要的模块

# 导入必要的模块
import torch
from torch import nn
from torch.utils.data import DataLoader
import torchvision
from torch.utils.tensorboard import SummaryWriter

 2、加载数据集

## 加载数据集
# 训练集
train_set = torchvision.datasets.CIFAR10(
    root='./dataset'# 保存路径
    ,train=True# 选择训练集
    ,transform=torchvision.transforms.ToTensor()# 预处理为tensor类
)
# 测试集
test_set = torchvision.datasets.CIFAR10(
    root='./dataset'# 保存路径
    ,train=False# 选择测试集
    ,transform=torchvision.transforms.ToTensor()# 预处理为tensor类
)

#查看数据集的长度
train_set_len = len(train_set)
test_set_len = len(test_set)

print(f'训练集的长度为{train_set_len}')
print(f'测试集的长度为{test_set_len}')

# 利用DataLoader加载数据集
train_dataloader = DataLoader(dataset=train_set, batch_size=64)
test_dataloader = DataLoader(dataset=test_set, batch_size=64)

 3、搭建模型

# 搭建神经网络
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1 = nn.Sequential(
            # Convolution1:
            # [inputs](3@32,32)=> 5*5 kernel=> (32@32,32)
            # padding = calc_padding(32,32,5) # 2 
            # 不可以将calc_padding函数直接放在Conv2d里面计算
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            
            # Max-pooling1:
            # (32@32,32)=> 2*2 kernel=> (32@16,16)
            nn.MaxPool2d(kernel_size=2),

            # Convolution2:
            # (32@16,16)=> 5*5 kernel=> (32@16,16)
            # padding = calc_padding(16,16,5) # 2
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),

            # Max-pooling2:
            # (32@16,16)=> 2*2 kernel=> (32@8,8)
            nn.MaxPool2d(kernel_size=2),

            # Convolution3:
            # (32@8,8)=> 5*5 kernel=> (64@8,8)
            # padding = calc_padding(8,8,5) # 2
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),

            # Max-pooling3:
            # (64@8,8)=> 2*2 kernel=> (64@4,4)
            nn.MaxPool2d(kernel_size=2),

            # Flatten:
            # (64@4,4)=> Flatten=> 64*4*4
            nn.Flatten(),
            
            # Linear1:
            # 1024 =>64
            nn.Linear(in_features=64*4*4, out_features=64),
            
            # Linear2:
            # 64 =>10 [outputs]
            nn.Linear(in_features=64, out_features=10)          
        )
    
    def forward(self,input):
        output = self.model1(input)
        return output

# 验证网络的准确性
if __name__ == '__main__':
    # 实例化模型
    mm = MyModel()
    # 新建一个与模型输入形状一样的数据
    input = torch.ones(64,3,32,32)
    output = mm(input)
    print(output.shape)# 看看输出形状与模型设计的输出是否一致

计算卷积层padding的方法

# 计算padding的方法
'''
    args:
        input_size(int): Number of size in the input image
        output_size(int): Number of size in the output image
        kernel_size (int or tuple): Size of the convolving kernel
        stride (int or tuple, optional): Stride of the convolution. Default: 1
        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
'''
def calc_padding(input_size, output_size, kernel_size, stride=1, dilation=1):
    # output_size= (input_size+ 2* padding- dilation* (kernel_size- 1)- 1)/ stride+ 1
    padding = ((output_size-1)* stride+ dilation* (kernel_size -1)- input_size+ 1) / 2
    return padding

 4、定义模型运行的设备

device = torch.device("cpu")
# 如果可以用cuda
if torch.cuda.is_available():
    device = torch.device("cuda")

## 语法糖，与上述表达式等价
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

5、创建网络模型、损失函数、优化器

## 创建网络模型
mm = MyModel()
mm = mm.to(device)# 使用定义的设备运行
# 这里可以不接收返回值、但是在数据中需要接收返回值（imgs = imgs.to(device)）
# mm.to(device)

## 创建损失函数
# 分类问题使用交叉熵作为损失函数
ce_loss = nn.CrossEntropyLoss()
ce_loss = ce_loss.to(device)# 使用定义的设备运行
# 这里也可以不接收返回值
# ce_loss.to(device)

## 创建优化器
# 学习速率
learining_rate = 1e-2
# 选取随机梯度下降优化器，优化器需要根据实际应用场景选取，这里只是做一个例子
'''
    params：模型的参数
'''
optimizer = torch.optim.SGD(params=mm.parameters(), lr=learing_rate)

6、训练/测试模型

# 创建一个SummaryWriter将训练结果和测试结果写到tensorboard
writer = SummaryWriter('logs_prog')

## 训练模型
epoch = 10 # 训练轮数
total_train_step = 0 # 训练次数
total_test_step = 0 # 测试次数

for i in range(epoch):
    print('-'*15+f'第{i+1}轮训练开始'+'-'*15)
    # 部分模型需要调用train() 启动训练模式
    # mm.train()
    for imgs,targets in train_dataloader:
        # 设置定义的设备，数据需要接收返回值
        imgs = imgs.to(device)
        targets = targets.to(device)

        # 输入数据到模型并接收输出值
        outputs = mm(imgs)
        
        # 计算输出与目标的损失
        result_loss = ce_loss(outputs,targets)

        # 使用优化器优化模型
        optimizer.zero_grad() # 1、梯度清零
        result_loss.backward() # 2、反向传播
        optimizer.step() # 调整优化

        # 记录训练次数
        total_train_step += 1
        
        # 输出训练结果
        if total_train_step % 100 == 0:#每一百次输出一次
            print(f'训练次数：{}，Loss：{result_loss.item()}')
            # 记录到tensorboard里
            writer.add_scalar('train_loss',result_loss.item(),total_train_step)

    # 训练完一轮后进行一次测试
    # 部分模型需要调用eval()启动验证模式
    # mm.eval()
    # 记录此次测试的总损失
    total_test_loss = 0
    # 记录此次测试的总准确数
    total_accuracy = 0
    # 没有梯度的情况下进行测试
    with torch.no_grad():
        for imgs,targets in test_dataloader:
            # 设置定义的设备
            imgs = imgs.to(device)
            targets = targets.to(device)

            # 输入到模型并接收输出值
            outputs = mm(imgs)

            # 计算输出与目标的损失
            result_loss = ce_loss(outputs,targets)
            # 记录总损失
            total_test_loss += result_loss.item()
            # 准确数为输出每一行中的最大值的索引与目标匹配的总和
            accuracy = (outputs.argmax(1) == targets).sum()
            # 记录总准确数
            total_accuracy+= accuracy
        # 记录测试次数
        total_test_step += 1
    print('-'*8+f'第{i+1}轮训练结束，整体测试集上的Loss：{total_test_loss}，准确率：{total_accuracy/test_set_len}'+'-'*8)# 准确率=准确数÷测试集数据长度
    # 记录测试结果到tensorboard
    writer.add_scalar('test_loss',total_test_loss,total_test_step)

    # 训练完一次之后保存模型，或者设置一个条件，满足条件再保存
    # if total_test_loss < @goal_loss:
    torch.save(mm.state_dict(),f'SaveModel/MyModel_epoch{i+1}.pth')

# 关闭writer
writer.close()