众所周知，神经网络搭建常用基础模块有卷积，池化，归一，激活，全连接等等。如果使用Pytorch进行网络的搭建时，除了需要掌握这些基础模块外，还需要熟悉模型容器。
Pytorch.nn 的容器containers
Pytorch搭建模型常用的3种方式如下：
- 使用 nn.Sequential 直接按顺序构建的模型;
- 继承 nn.Module 基类构建自定义模型;
- 继承 nn.Module 基类，通过辅助模型容器(Sequential, ModuleList, ModuleDict)来构建。
辅助模型容器如下：
- nn.Sequential：顺序性，各网络层之间严格按顺序执行，用于固定block构建；
- nn.ModuleList：迭代性，list形式，用于大量重复网络层构建，通过for循环实现重复构建；
- nn.ModuleDict：索引性，dict形式，用于可选择的网络层，key+value调用方式；
部分基础模块(其他可参考 >> Pytorch.nn)：
- torch.nn.Conv2d
- torch.nn.MaxPool2d, torch.nn.AvgPool2d
- torch.nn.BatchNorm2d
- torch.nn.ReLU
分割网络搭建
使用torch.nn里面的骨干网络直接出来，搭建分割模块，当前示例以resnet的经典模型先实践，后期将对不同模型都实践一遍。
-
模型设计，初步使用resnet18骨干
resnet18为骨干网络.png
resnet18为骨干网络-参数输入.png -
提取resnet18的layer3层的输出作为分割分支。
分割网络示意图.png
- 代码示例如下
######## py内置函数：help-文件架构, dir-代码架构 ########
import torch # 包含基本，加减乘除，张量操作，优化器'torch.optim', 数据索引 'torch.utils.data.DataLoader'
import torch.nn as nn # "类": 包含卷积，池化，激活，损失等 "nn.CrossEntropyLoss()"
import torch.nn.functional as F # "函数": 包含卷积，池化，激活，损失等 "F.cross_entropy()"
import torchvision # 包含图像算法的基本操作等 torchvision.models; torchvision.datasets;
import torchvision.transforms as T # "类": 包含图像增强方向等 "T.RandomCrop()"
import torchvision.transforms.functional as TF # "函数": 包含图像增强方向等 "TF.center_crop()"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#### model build ####
class FCNHead(nn.Sequential):
'''nn.Sequential默认一定调用forward，因此直接init即可'''
def __init__(self, in_channels, channels):
inter_channels = in_channels // 4
## 写法1: layers=nn.Sequential(...,...);
## 写法2: layers=[...,...], (list)格式 ;
layers = nn.Sequential(
nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
nn.BatchNorm2d(inter_channels),
nn.ReLU(),
nn.Dropout(0.1),
nn.Conv2d(inter_channels, channels, 1),
)
super().__init__(*layers)
pass
pass
class SimpleSegModel(nn.Module):
def __init__(self, backbone: nn.Module, head: nn.Module) -> None:
super().__init__()
self.backbone = backbone
self.head = head
pass
def forward(self, x):
result = {}
input_shape = x.shape[-2:]
features = self.backbone(x) # OrderedDict
C3, C4, C5 = features["C3"], features["C4"], features["C5"]
H3 = self.head(C3)
H3 = F.interpolate(H3, size=input_shape, mode="bilinear", align_corners=False)
result["H3"] = H3
return result
pass
backbone = torchvision.models.resnet.resnet18(False)
feature_op = torchvision.models._utils.IntermediateLayerGetter(backbone,
return_layers={"layer2": "C3",
"layer3": "C4",
"layer4": "C5"})
model = SimpleSegModel(feature_op, FCNHead(128, 2))
model.to(device)
inputs = torch.ones(8, 3, 256,256)
outputs = model(inputs.to(device))["H3"]
print(outputs.size())
torch.Size([8, 2, 256, 256])
总结
模型搭建完毕，使用骨干网络（轻量骨干搭建，复杂骨干搭建，多任务骨干搭建持续更新种）的节点引出自己的分割分支，运用辅助模型容器搭建分割头的类，完成分割网络的搭建过程。
参考链接
pytorch官方网站
github的torchvision
图像语义分割实践
