pytorch的python API略读--Sequential/ModuleList/ModuleDict

作者:机器视觉全栈er
网站:cvtutorials.com

torch.nn.Sequential:序列容器,顾名思义,就是将构造函数中的子模块会按照你定义的序列顺序被添加到模块中。这里有个注意的小点,要注意区分Sequential和torch.nn.ModuleList,后者就是一个简单的列表,里面的元素是模块,但是模块之间是孤立的,前者则是互相连通的模块。还是以上面的网络模块为例,使用Sequential可以实现快速搭建:

import torch.nn as nn
import torch

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(1, 3, 3),
            nn.ReLU(),
            nn.Conv2d(3, 3, 3),
            nn.ReLU()
            )

    def forward(self, x):
        x = self.layer(x)
        return x

my_module = MyModule()

cvtutorials = torch.randn(3, 1, 5, 5)
print("=====================================")
print(cvtutorials)
print("=====================================")
print(my_module(cvtutorials))

torch.nn.ModuleList():顾名思义,是多个Module组成的列表,还是以前面的网络模型举例:

import torch.nn as nn
import torch

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Conv2d(1, 3, 3),
            nn.ReLU(),
            nn.Conv2d(3, 3, 3),
            nn.ReLU()])
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
my_module = MyModule()
cvtutorials = torch.randn(3, 1, 5, 5)
print("=====================================")
print(cvtutorials)
print("=====================================")
my_module(cvtutorials)

torch.nn.ModuleDict():顾名思义,是多个Module组成的字典,还是以前面的网络模型举例(注意这里前向推理的时候,只选择了conv1和activation1):

import torch.nn as nn
import torch

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.convs = nn.ModuleDict({
            "conv1": nn.Conv2d(1, 3, 3),
            "conv2": nn.Conv2d(3, 3, 3)
        })
        self.activations = nn.ModuleDict({
            "activation1": nn.ReLU(),
            "activation2": nn.ReLU()
        })

    def forward(self, x, conv, activation):
        x = self.convs[conv](x)
        x = self.activations[activation](x)
        return x

my_module = MyModule()
cvtutorials = torch.randn(3, 1, 5, 5)
print("=====================================")
print(cvtutorials)
print("=====================================")
my_module(cvtutorials, "conv1", "activation1")

现在我们来总结下Sequential, ModuleList, ModuleDict三者常见的应用场景:

www.cvtutorials.com Sequential ModuleList ModuleDict
特点 顺序性 迭代性 索引性
应用场景 网络block搭建 大量重复网络结构 可选择网络模块

你可能感兴趣的