时间戳

首页 > 解决方案 > 可以在 Pytorch nn.Sequential() 中添加条件

pytorch - 可以在 Pytorch nn.Sequential() 中添加条件

问题描述

有没有办法在nn.Sequential(). 类似于下面的代码。

import torch


class Building_Blocks(torch.nn.Module):
  def conv_block (self, in_features, out_features, kernal_size, upsample=False):
    block = torch.nn.Sequential(
        torch.nn.Conv2d(in_features, out_features, kernal_size),
        torch.nn.ReLU(inplace = True),
        torch.nn.Conv2d(out_features, out_features, kernal_size),
        torch.nn.ReLU(inplace = True),
        if(upsample):
          torch.nn.ConvTranspose2d(out_features, out_features, kernal_size)
        )
    return block

  def __init__(self):
    super(Building_Blocks, self).__init__()
    self.contracting_layer1 = self.conv_block(3, 64, 3, upsample=True)

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

标签： pytorch

解决方案

不，但是在您的情况下，很容易if摆脱nn.Sequential：

class Building_Blocks(torch.nn.Module):
    def conv_block(self, in_features, out_features, kernal_size, upsample=False):
        layers = [
            torch.nn.Conv2d(in_features, out_features, kernal_size),
            torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(out_features, out_features, kernal_size),
            torch.nn.ReLU(inplace=True),
        ]
        if upsample:
            layers.append(
                torch.nn.ConvTranspose2d(out_features, out_features, kernal_size)
            )
        block = torch.nn.Sequential(*layers)
        return block

    def __init__(self):
        super(Building_Blocks, self).__init__()
        self.contracting_layer1 = self.conv_block(3, 64, 3, upsample=True)

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

您始终可以根据需要构建list包含层，然后将其解压缩torch.nn.Sequential。

推荐阅读