Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Can conditionality be added inside Pytorch nn.Sequential()

Tags:

pytorch

Is there a way to add conditional statements inside the nn.Sequential(). Something similar to the code below.

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
like image 776
aravind pothana Avatar asked Sep 13 '25 22:09

aravind pothana


1 Answers

No, but in your case it's easy to take if out of 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

You can always construct a list containing layers however you want and unpack it into torch.nn.Sequential afterwards.

like image 62
Szymon Maszke Avatar answered Sep 16 '25 00:09

Szymon Maszke