pytorch 使用 add_module 添加模块

最新推荐文章于 2024-01-23 09:44:52 发布

梅dá爷~

最新推荐文章于 2024-01-23 09:44:52 发布

阅读量4k

点赞数 4

文章标签：深度学习 python

本文链接：https://blog.csdn.net/weixin_44400443/article/details/111832267

版权

利用add_module ,添加conv_2.
add_module用法：add_module（name，module）

直接看代码

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,