PyTorch——add_module

torch.nn.Module.add_module

添加一个子模块到当前的模块。

模块可以用给定的名称以属性获取。

输入参数有：

• name：str类型，子模块名称；
• module：nn.Module类型，添加到模块的子模块；

import torch
from torch import nn

o_module = nn.Module()
print(o_module)  # Module()

o_module.add_module('my_conv2d', nn.Conv2d(3, 16, 3))
print(o_module)  # Module(((mv_conv2d): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1))

img = torch.randn(1, 3, 256, 256)
print(o_module.my_conv2d(img.shape))  # torch.Size([1, 16, 254, 254])