是否可以为 nn.Module 对象添加子模块
时间: 2024-02-28 20:55:01 浏览: 55
是的,可以为 nn.Module 对象添加子模块。可以通过调用 nn.Module 的 add_module 方法,向 nn.Module 对象中添加子模块。add_module 方法需要传入两个参数:子模块的名称和子模块对象。例如,下面的代码创建了一个包含两个子模块的 nn.Module 对象:
```python
import torch.nn as nn
class MyModule(nn.Module):
def __init__(self):
super(MyModule, self).__init__()
self.conv = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.relu = nn.ReLU()
self.add_module('conv_layer', self.conv)
self.add_module('relu_layer', self.relu)
def forward(self, x):
x = self.conv_layer(x)
x = self.relu_layer(x)
return x
```
在这个例子中,我们创建了一个名为 MyModule 的 nn.Module 对象,它包含了两个子模块:一个卷积层 nn.Conv2d 和一个激活函数 nn.ReLU。我们首先定义了这两个子模块,然后通过 add_module 方法向 nn.Module 对象中添加子模块。在 forward 方法中,我们通过子模块的名称来逐个调用子模块,实现了一个简单的前向传播过程。
阅读全文