在上节中我们用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