Pytorch完整训练自己的数据集

在上节中我们用Pytorch读取自己的数据集：https://blog.csdn.net/lifei1229/article/details/105519924。

在这节中，我们继续完成模型构建及其训练。可参考我之前写的：Pytorch训练模板

1 构建模型

本次实验采用ResNet模型。代码很简单，网上找下有很多，也可以用torchvision里面官方写的。

import torch
from torch import nn
from torch.nn import functional as F

class ResBlk(nn.Module):
    """
    resnet block
    """
    def __init__(self, ch_in, ch_out, stride=1):
        """
        :param ch_in:
        :param ch_out:
        """
        super(ResBlk, self).__init__() #下面是瓶颈层
        self.conv1 = nn.Conv2d( ch_in,ch_out, kernel_size=3,stride=stride,padding=1) #padding=1保证输出大小一致
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)
        self.extra = nn.Sequential()
        if ch_out != ch_in:   #保存通道一致，才可以相加
            # [b, ch_in, h, w] => [b, ch_out, h, w]
            self.extra = nn.Sequential(
						                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),
						                nn.BatchNorm2d(ch_out))
    def forward(self, x):
        """
        :param x: [b, ch, h, w]
        :return:
        """
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        # short cut.
        # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
        # element-wise add:
        out = self.extra(x) + out
        out = F.relu(out)
        return out

class ResNet18