用pytorch里的children方法自定义网络

假设我们要自定义一个网络,这个网络由resnet-18的所有除了最后一层和倒数第二层,最后一层用我们的自定义层。这时我们可以用到children方法。

children()返回网络模型里的组成元素,且children()返回的是最外层的元素

举个例子:

m = nn.Sequential(nn.Linear(2,2), 
                  nn.ReLU(),
                 nn.Sequential(nn.Sigmoid(), nn.ReLU()))

m.children()返回的是(一个list):

[Linear(in_features=2, out_features=2), ReLU(), Sequential(
   (0): Sigmoid()
   (1): ReLU()
 )]

此时我们对children方法应该有了一些了解,那么回到我们刚才的问题:

pretrained_net = torchvision.models.resnet18(pretrained=True)
net = nn.Sequential(*list(pretrained_net.children())[:-2])

通过上面两行代码就可以获得resnet-18的除最后两层的所有层。

然后我们通过add_module方法添加自己想要的层:

num_classes = 21
net.add_module('final_conv', nn.Conv2d(512, num_classes, kernel_size=1))
net.add_module('transpose_conv', nn.ConvTranspose2d(num_classes, num_classes,
                                    kernel_size=64, padding=16, stride=32))

这样,我们就完成了定义自己的网络。

  • 4
    点赞
  • 9
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值