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
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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With