0

我已将默认的 tensorboard 记录器 ( from pytorch_lightning.loggers import TensorBoardLogger) 添加到我的 pytorch 闪电训练器中log_graph=True。

当我训练我的模型时，我的图表的第一个视图显示了三个块：

输入 => MyNetworkClassName => 输出

到现在为止还挺好。

但是，当我扩展 MyNetworkClassName 时，它​​绝对给了我网络中发生的一切。到处都是箭。我想将此图组织成带有可扩展子图的更简单的块。所以在我的例子中，我的网络有一个典型的编码器 - 增强器 - 解码器结构，我想要更像这样的东西：

  • 第一张图：

    1. 输入 => MyNetworkClassName => 输出

  • 放大 MyNetworkClassName：

    1. 编码器 => 增强器 -> 解码器

  • 放大编码器：

    1. 编码器层1 => 编码器层2 => ...

  • 放大encoder_layer1：

    1. conv2d => 批处理规范

我在这里有什么选择？我应该把所有东西都放在单独的类中吗？是否有任何命令可以让我将某些操作组合在一起？

4

1 回答 1

0

将代码重构为类也会影响张量板图（其中重构为方法不会）。将显示为可扩展块的典型示例类：

class EncoderLayer(nn.Module):
"""Encoder layer class"""

    def __init__(self, activation_function, kernel_num, kernel_size, idx):
        super().__init__()
        self.layer = nn.Sequential(
            ComplexConv2d(
            kernel_num[idx],
            kernel_num[idx + 1],
            kernel_size=(kernel_size, 2),
        ),
        nn.BatchNorm2d(kernel_num[idx + 1])
        activation_function,
    )

def forward(self, x):
    return self.layer(x)
于 2020-11-18T08:01:33.680 回答