【Pytorch学习】add_module()函数用法

本文通过一个具体的网络模型示例介绍了PyTorch中add_module函数的应用方式。该函数可以方便地为模型添加模块,使得代码更加清晰易读。通过实例展示了如何使用add_module来构建包含多个卷积层的神经网络。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

  最近在看大佬们写的代码时,看到使用add_module函数。所以就了解了一番,在这里做个介绍。add_module就好像是list的用法,看代码之后就懂了(代码来自文献【3】)

from torch import nn
from torchsummary import summary

class Net_test(nn.Module):
    def __init__(self):
        super(Net_test,self).__init__()
        self.conv_1 = nn.Conv2d(3,6,3)
        self.add_module('conv_2', nn.Conv2d(6,12,3))
        self.conv_3 = nn.Conv2d(12,24,3)
        
    def forward(self,x):
        x = self.conv_1(x)
        x = self.conv_2(x)
        x = self.conv_3(x)
        return x
    
model = Net_test()
print(model)
model.to('cuda')
summary(model,(3,128,128))

“”“
Net_test(
  (conv_1): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1))
  (conv_2): Conv2d(6, 12, kernel_size=(3, 3), stride=(1, 1))
  (conv_3): Conv2d(12, 24, kernel_size=(3, 3), stride=(1, 1))
)
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1          [-1, 6, 126, 126]             168
            Conv2d-2         [-1, 12, 124, 124]             660
            Conv2d-3         [-1, 24, 122, 122]           2,616
================================================================
Total params: 3,444
Trainable params: 3,444
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 4.86
Params size (MB): 0.01
Estimated Total Size (MB): 5.06
“””

参考文献

[1][Pytorch进阶技巧(一)] 使用add_module替换部分模型
[2]PyTorch-网络的创建,预训练模型的加载
[3]pytorch 使用 add_module 添加模块

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值