【深度之眼】Pytorch框架班第五期-模型容器之nn.ModuleDict代码调试

模型容器(Containers)

nn.ModuleDict

nn.ModuleDict 是nn.module的容器，用于包装一组网络层，以索引方式调用网络层。
主要方法：

• clear(): 清空ModuleDict
• items(): 返回可迭代的键值对(key-value pairs)
• keys(): 返回字典的键（key）
• values(): 返回字典的值(value)
• pop(): 返回一对键值，并从字典中删除

class ModuleDict(nn.Module):
    def __init__(self):
        super(ModuleDict, self).__init__()
        self.choices = nn.ModuleDict({
            'conv': nn.Conv2d(10, 10, 3),
            'pool': nn.MaxPool2d(3)
        })

        self.activations = nn.ModuleDict({
            'relu': nn.ReLU(),
            'prelu': nn.PReLU()
        })

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

net = ModuleDict()

fake_img = torch.randn((4, 10, 32, 32))

output = net(fake_img, 'conv', 'relu')

print(output)

从上面的代码我们可以看到在forward处我们传入两个参数，第一个参数选择使用conv还是pool，第二个参数选择使用relu还是prelu。

1、使用relu时output没有负值。

2、使用prelu时，有负值。

容器总结