devbox
不正经
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

U2Net网络结构搭建_do not use torch.sigmoid for amp safe

作者：不正经 | 2024-06-14 03:04:00

赞

踩

do not use torch.sigmoid for amp safe

论文题目：U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection

论文链接：https://arxiv.org/abs/2005.09007

一、RSU-7模块

  1. class RSU(nn.Module):
  2.     def __init__(self, height: int, in_ch: int, mid_ch: int, out_ch: int):
  3.         super().__init__()
  4.  
  5.         assert height >= 2 #断言
  6.         self.conv_in = ConvBNReLU(in_ch, out_ch) # 第一个ConvBNReLU
  7.  
  8.         # 构建列表encode_list
  9.         encode_list = [DownConvBNReLU(out_ch, mid_ch, flag=False)] #第一个encode的ConvBNReLU，没有下采样
  10.         # 构建列表decode_list
  11.         decode_list = [UpConvBNReLU(mid_ch * 2, mid_ch, flag=False)] #第一个decode的ConvBNReLU，没有上采样
  12.         for i in range(height - 2):
  13.             encode_list.append(DownConvBNReLU(mid_ch, mid_ch)) #encode中有下采样的ConvBNReLU
  14.             #decode中有上采样的ConvBNReLU，最后一个ConvBNReLU的输出通道数为out_ch，输入通道为mid_ch * 2是因为concat
  15.             decode_list.append(UpConvBNReLU(mid_ch * 2, mid_ch if i < height - 3 else out_ch))
  16.  
  17.         encode_list.append(ConvBNReLU(mid_ch, mid_ch, dilation=2)) #最下面一个含膨胀卷积 encode的ConvBNReLU，将其添加进列表encode_list
  18.         self.encode_modules = nn.ModuleList(encode_list)
  19.         self.decode_modules = nn.ModuleList(decode_list)
  20.  
  21.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  22.         x_in = self.conv_in(x)
  23.  
  24.         x = x_in
  25.         encode_outputs = []
  26.         for m in self.encode_modules:
  27.             x = m(x)
  28.             encode_outputs.append(x) # 依次通过encode_modules，并把每一层的输出添加进列表encode_outputs
  29.  
  30.         x = encode_outputs.pop() # 最后一个encode的ConvBNReLU输出，作为decode的输入
  31.         for m in self.decode_modules:
  32.             x2 = encode_outputs.pop() #依次弹出
  33.             x = m(x, x2) # 拼接后在传入decode_modules
  34.  
  35.         return x + x_in  #最后的输出与第一个ConvBNReLU的结果进行add操作

 

二、RSU-4F模块

  1. class RSU4F(nn.Module):
  2.     def __init__(self, in_ch: int, mid_ch: int, out_ch: int):
  3.         super().__init__()
  4.         self.conv_in = ConvBNReLU(in_ch, out_ch)
  5.         # encode模块，包含四个ConvBNReLU
  6.         self.encode_modules = nn.ModuleList([ConvBNReLU(out_ch, mid_ch),
  7.                                              ConvBNReLU(mid_ch, mid_ch, dilation=2),
  8.                                              ConvBNReLU(mid_ch, mid_ch, dilation=4),
  9.                                              ConvBNReLU(mid_ch, mid_ch, dilation=8)])
  10.  
  11.         # decode模块，包含三个ConvBNReLU
  12.         self.decode_modules = nn.ModuleList([ConvBNReLU(mid_ch * 2, mid_ch, dilation=4),
  13.                                              ConvBNReLU(mid_ch * 2, mid_ch, dilation=2),
  14.                                              ConvBNReLU(mid_ch * 2, out_ch)])
  15.  
  16.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  17.         x_in = self.conv_in(x)
  18.  
  19.         x = x_in
  20.         encode_outputs = []
  21.         for m in self.encode_modules:
  22.             x = m(x)
  23.             encode_outputs.append(x) # 依次通过encode_modules，并把每一层的输出添加进列表encode_outputs
  24.  
  25.         x = encode_outputs.pop() # 最后一个encode的ConvBNReLU输出，作为decode的输入
  26.         for m in self.decode_modules:
  27.             x2 = encode_outputs.pop()
  28.             # 没有用到UpConvBNReLU模块，故需要torch.cat([x, x2]
  29.             x = m(torch.cat([x, x2], dim=1)) 
  30.  
  31.         return x + x_in #最后的输出与第一个ConvBNReLU的结果进行add操作

 三、U2Net整体结构

  1. from typing import Union, List
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5.  
  6.  
  7. class ConvBNReLU(nn.Module):
  8.     def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, dilation: int = 1): #dilation=1代表普通卷积，大于1代表膨胀卷积
  9.         super().__init__()
  10.  
  11.         padding = kernel_size // 2 if dilation == 1 else dilation
  12.         self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding, dilation=dilation, bias=False)
  13.         self.bn = nn.BatchNorm2d(out_ch)
  14.         self.relu = nn.ReLU(inplace=True)
  15.  
  16.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  17.         return self.relu(self.bn(self.conv(x)))
  18.  
  19.  
  20. class DownConvBNReLU(ConvBNReLU): #继承自ConvBNReLU
  21.     def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, dilation: int = 1, flag: bool = True): #flag表示是否启用下采样
  22.         super().__init__(in_ch, out_ch, kernel_size, dilation) #调用父类的方法
  23.         self.down_flag = flag
  24.  
  25.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  26.         if self.down_flag: #如果flag为True，则启用下采样
  27.             x = F.max_pool2d(x, kernel_size=2, stride=2, ceil_mode=True)
  28.         return self.relu(self.bn(self.conv(x)))
  29.  
  30.  
  31. class UpConvBNReLU(ConvBNReLU): #继承自ConvBNReLU
  32.     def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, dilation: int = 1, flag: bool = True): #flag表示是否启用上采样
  33.         super().__init__(in_ch, out_ch, kernel_size, dilation)
  34.         self.up_flag = flag
  35.  
  36.     def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
  37.         if self.up_flag: #如果flag为True，则启用上采样
  38.             x1 = F.interpolate(x1, size=x2.shape[2:], mode='bilinear', align_corners=False) #x1通过双线性插值上采用到x2的宽高
  39.         return self.relu(self.bn(self.conv(torch.cat([x1, x2], dim=1)))) #x1, x2进行拼接，再ConvBNReLU
  40.  
  41.  
  42. class RSU(nn.Module):
  43.     def __init__(self, height: int, in_ch: int, mid_ch: int, out_ch: int):
  44.         super().__init__()
  45.  
  46.         assert height >= 2 #断言
  47.         self.conv_in = ConvBNReLU(in_ch, out_ch) # 第一个ConvBNReLU
  48.  
  49.         # 构建列表encode_list
  50.         encode_list = [DownConvBNReLU(out_ch, mid_ch, flag=False)] #第一个encode的ConvBNReLU，没有下采样
  51.         # 构建列表decode_list
  52.         decode_list = [UpConvBNReLU(mid_ch * 2, mid_ch, flag=False)] #第一个decode的ConvBNReLU，没有上采样
  53.         for i in range(height - 2):
  54.             encode_list.append(DownConvBNReLU(mid_ch, mid_ch)) #encode中有下采样的ConvBNReLU
  55.             #decode中有上采样的ConvBNReLU，最后一个ConvBNReLU的输出通道数为out_ch，输入通道为mid_ch * 2是因为concat
  56.             decode_list.append(UpConvBNReLU(mid_ch * 2, mid_ch if i < height - 3 else out_ch))
  57.  
  58.         encode_list.append(ConvBNReLU(mid_ch, mid_ch, dilation=2)) #最下面一个含膨胀卷积 encode的ConvBNReLU，将其添加进列表encode_list
  59.         self.encode_modules = nn.ModuleList(encode_list)
  60.         self.decode_modules = nn.ModuleList(decode_list)
  61.  
  62.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  63.         x_in = self.conv_in(x)
  64.  
  65.         x = x_in
  66.         encode_outputs = []
  67.         for m in self.encode_modules:
  68.             x = m(x)
  69.             encode_outputs.append(x) # 依次通过encode_modules，并把每一层的输出添加进列表encode_outputs
  70.  
  71.         x = encode_outputs.pop() # 最后一个encode的ConvBNReLU输出，作为decode的输入
  72.         for m in self.decode_modules:
  73.             x2 = encode_outputs.pop() #依次弹出
  74.             # UpConvBNReLU已经包含torch.cat([x1, x2]
  75.             x = m(x, x2) # 拼接后在传入decode_modules
  76.  
  77.         return x + x_in  #最后的输出与第一个ConvBNReLU的结果进行add操作
  78.  
  79.  
  80. class RSU4F(nn.Module):
  81.     def __init__(self, in_ch: int, mid_ch: int, out_ch: int):
  82.         super().__init__()
  83.         self.conv_in = ConvBNReLU(in_ch, out_ch)
  84.         # encode模块，包含四个ConvBNReLU
  85.         self.encode_modules = nn.ModuleList([ConvBNReLU(out_ch, mid_ch),
  86.                                              ConvBNReLU(mid_ch, mid_ch, dilation=2),
  87.                                              ConvBNReLU(mid_ch, mid_ch, dilation=4),
  88.                                              ConvBNReLU(mid_ch, mid_ch, dilation=8)])
  89.  
  90.         # decode模块，包含三个ConvBNReLU
  91.         self.decode_modules = nn.ModuleList([ConvBNReLU(mid_ch * 2, mid_ch, dilation=4),
  92.                                              ConvBNReLU(mid_ch * 2, mid_ch, dilation=2),
  93.                                              ConvBNReLU(mid_ch * 2, out_ch)])
  94.  
  95.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  96.         x_in = self.conv_in(x)
  97.  
  98.         x = x_in
  99.         encode_outputs = []
  100.         for m in self.encode_modules:
  101.             x = m(x)
  102.             encode_outputs.append(x) # 依次通过encode_modules，并把每一层的输出添加进列表encode_outputs
  103.  
  104.         x = encode_outputs.pop() # 最后一个encode的ConvBNReLU输出，作为decode的输入
  105.         for m in self.decode_modules:
  106.             x2 = encode_outputs.pop()
  107.             # 没有用到UpConvBNReLU模块，故需要torch.cat([x, x2]
  108.             x = m(torch.cat([x, x2], dim=1))
  109.  
  110.         return x + x_in #最后的输出与第一个ConvBNReLU的结果进行add操作
  111.  
  112.  
  113. class U2Net(nn.Module):
  114.     def __init__(self, cfg: dict, out_ch: int = 1): #此处的out_ch:是为1的
  115.         super().__init__()
  116.         assert "encode" in cfg
  117.         assert "decode" in cfg
  118.         self.encode_num = len(cfg["encode"]) #6
  119.  
  120.         encode_list = [] # 存储每个encode模块的输出
  121.         side_list = [] # 收集每个decode和最后一个encode的输出
  122.         for c in cfg["encode"]: #遍历encode
  123.             # c: [height, in_ch, mid_ch, out_ch, RSU4F, side]
  124.             assert len(c) == 6 # c中有六个参数
  125.             # 如果是RSU，则传入height, in_ch, mid_ch, out_ch 如果是RSU4F，则传入in_ch, mid_ch, out_ch
  126.             encode_list.append(RSU(*c[:4]) if c[4] is False else RSU4F(*c[1:4]))
  127.  
  128.             if c[5] is True: #有1个
  129.                 side_list.append(nn.Conv2d(c[3], out_ch, kernel_size=3, padding=1)) #每层的输出后接一个3*3的卷积，此处的out_ch:是为1的
  130.         self.encode_modules = nn.ModuleList(encode_list)
  131.  
  132.         decode_list = []
  133.         for c in cfg["decode"]: #遍历decode
  134.             # c: [height, in_ch, mid_ch, out_ch, RSU4F, side]
  135.             assert len(c) == 6
  136.             # 如果是RSU，则传入height, in_ch, mid_ch, out_ch 如果是RSU4F，则传入in_ch, mid_ch, out_ch
  137.             decode_list.append(RSU(*c[:4]) if c[4] is False else RSU4F(*c[1:4]))
  138.  
  139.             if c[5] is True: #有5个
  140.                 side_list.append(nn.Conv2d(c[3], out_ch, kernel_size=3, padding=1)) #每层的输出后接一个3*3的卷积，此处的out_ch:是为1的
  141.         self.decode_modules = nn.ModuleList(decode_list)
  142.         self.side_modules = nn.ModuleList(side_list) # encode和decode并在一起的
  143.         self.out_conv = nn.Conv2d(self.encode_num * out_ch, out_ch, kernel_size=1) #concat之后跟的1*1卷积，输入通道为6，输出通道为1
  144.  
  145.     def forward(self, x: torch.Tensor) -> Union[torch.Tensor, List[torch.Tensor]]:
  146.         _, _, h, w = x.shape #获取输入图片的高宽
  147.  
  148.         # collect encode outputs
  149.         encode_outputs = [] # 收集每个encode的输出
  150.         for i, m in enumerate(self.encode_modules):
  151.             x = m(x)
  152.             encode_outputs.append(x) #放入encode_outputs列表中的是下采样之前的
  153.             if i != self.encode_num - 1: #前五个encode层后都会进行下采样，最后一个encode后不会进行下采样
  154.                 x = F.max_pool2d(x, kernel_size=2, stride=2, ceil_mode=True) #下采样之后再输入下一层
  155.  
  156.         # collect decode outputs
  157.         x = encode_outputs.pop()
  158.         decode_outputs = [x] #先收集En_6的输出
  159.         for m in self.decode_modules:
  160.             x2 = encode_outputs.pop() #弹出En_5的输出
  161.             x = F.interpolate(x, size=x2.shape[2:], mode='bilinear', align_corners=False) #最后一个encode的输出x，通过双线性插值上采用到x2的宽高
  162.             x = m(torch.concat([x, x2], dim=1))
  163.             decode_outputs.insert(0, x) # 将x插入在列表decode_outputs的最前面
  164.             #最终列表decode_outputs = [De_1,De_2,De_3,De_4,De_5,En_6]
  165.  
  166.         # collect side outputs
  167.         side_outputs = [] #用于存储通过3*3卷积之后的输出
  168.         #side_modules构建的顺序是从En_6,De_5,De_4,De_3,De_2,De_1
  169.         for m in self.side_modules:
  170.             x = decode_outputs.pop()
  171.             x = F.interpolate(m(x), size=[h, w], mode='bilinear', align_corners=False) #经过3*卷积核，通过双线性插值将其还原回输入图片的h, w
  172.             side_outputs.insert(0, x) # 将x插入在列表side_outputs的最前面
  173.             # 最终列表side_outputs = [Sup1,Sup2,Sup3,Sup4,Sup5,Sup6]
  174.  
  175.         x = self.out_conv(torch.concat(side_outputs, dim=1)) #六个特征图经concat拼接之后，在经过1*1卷积
  176.  
  177.         if self.training:
  178.             # do not use torch.sigmoid for amp safe
  179.             return [x] + side_outputs #训练模式，返回x和[Sup1,Sup2,Sup3,Sup4,Sup5,Sup6]，在算损失的时候会用到
  180.         else:
  181.             return torch.sigmoid(x) #预测模式，直接sigmoid到 0~1
  182.  
  183.  
  184. def u2net_full(out_ch: int = 1):
  185.     cfg = {
  186.         # height为深度，side为是否收集其输出
  187.         # height, in_ch, mid_ch, out_ch, RSU4F, side
  188.         "encode": [[7, 3, 32, 64, False, False],      # En1
  189.                    [6, 64, 32, 128, False, False],    # En2
  190.                    [5, 128, 64, 256, False, False],   # En3
  191.                    [4, 256, 128, 512, False, False],  # En4
  192.                    [4, 512, 256, 512, True, False],   # En5
  193.                    [4, 512, 256, 512, True, True]],   # En6
  194.         # height, in_ch, mid_ch, out_ch, RSU4F, side
  195.         "decode": [[4, 1024, 256, 512, True, True],   # De5
  196.                    [4, 1024, 128, 256, False, True],  # De4
  197.                    [5, 512, 64, 128, False, True],    # De3
  198.                    [6, 256, 32, 64, False, True],     # De2
  199.                    [7, 128, 16, 64, False, True]]     # De1
  200.     }
  201.  
  202.     return U2Net(cfg, out_ch)
  203.  
  204.  
  205. def u2net_lite(out_ch: int = 1):
  206.     cfg = {
  207.         # height, in_ch, mid_ch, out_ch, RSU4F, side
  208.         "encode": [[7, 3, 16, 64, False, False],  # En1
  209.                    [6, 64, 16, 64, False, False],  # En2
  210.                    [5, 64, 16, 64, False, False],  # En3
  211.                    [4, 64, 16, 64, False, False],  # En4
  212.                    [4, 64, 16, 64, True, False],  # En5
  213.                    [4, 64, 16, 64, True, True]],  # En6
  214.         # height, in_ch, mid_ch, out_ch, RSU4F, side
  215.         "decode": [[4, 128, 16, 64, True, True],  # De5
  216.                    [4, 128, 16, 64, False, True],  # De4
  217.                    [5, 128, 16, 64, False, True],  # De3
  218.                    [6, 128, 16, 64, False, True],  # De2
  219.                    [7, 128, 16, 64, False, True]]  # De1
  220.     }
  221.  
  222.     return U2Net(cfg, out_ch)
  223.  
  224.  
  225. def convert_onnx(m, save_path):
  226.     m.eval()
  227.     x = torch.rand(1, 3, 288, 288, requires_grad=True)
  228.  
  229.     # export the model
  230.     torch.onnx.export(m,  # model being run
  231.                       x,  # model input (or a tuple for multiple inputs)
  232.                       save_path,  # where to save the model (can be a file or file-like object)
  233.                       export_params=True,
  234.                       opset_version=11)
  235.  
  236.  
  237. if __name__ == '__main__':
  238.     # n_m = RSU(height=7, in_ch=3, mid_ch=12, out_ch=3)
  239.     # convert_onnx(n_m, "RSU7.onnx")
  240.     #
  241.     # n_m = RSU4F(in_ch=3, mid_ch=12, out_ch=3)
  242.     # convert_onnx(n_m, "RSU4F.onnx")
  243.  
  244.     u2net = u2net_full()
  245.     convert_onnx(u2net, "u2net_full.onnx")

bulalalalalala~~~作者代码写的太妙了，看完源码后真的很佩服，点赞点赞！！！

bulalalalalala~~~废了好大劲才搞定！！！

reference

【U2Net源码解析(Pytorch)】 https://www.bilibili.com/video/BV1Kt4y137iS?p=2&share_source=copy_web&vd_source=95705b32f23f70b32dfa1721628d5874

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/不正经/article/detail/716006
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号