在设计深度神经网络的时候,往往需要迭代很多个结构相似、参数不同的模块,这时候如果把一个个的模块写出来当然可以,而且逻辑很清晰,但缺点是“工作量大”,形式上太笨。一种更为简洁的方法是类似定义函数或定义类一样把一个个模块给抽象出来,然后把模块在神经网络中参数化。

如下代码:

import torch.nn as nn
import torch
from torch import autograd

#这时候可以定义一个基本模块,该模块由“两层”卷积、BN和激活构成,以后使用的时候可以实例化该模块
class DoubleConv(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(DoubleConv,self).__init__()
        self.conv=nn.Sequential(
            nn.Conv2d(in_ch,out_ch,3,padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(out_ch,out_ch,3,padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    
    def forward(self,input):
        return self.conv(input)

#在实现深度神经网络时候,就可以使用不同的参数实例化以上模块
class Unet(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(Unet,self).__init__()
        
        self.conv1=Double(in_ch,64)
        ...

 

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐