Use torch.square() inside a nn.Sequential layer in PyTorch

Question

I want to square the result of a maxpool layer. I tried the following:

class CNNClassifier(Classifier): # nn.Module def __init__(self, in_channels): super().__init__() self.save_hyperparameters('in_channels') self.cnn = nn.Sequential( # maxpool nn.MaxPool2d((1, 5), stride=(1, 5)), torch.square(), # layer1 nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=5, )

Which to the experienced PyTorch user for sure makes no sense.

Indeed, the error is quite clear:

TypeError: square() missing 1 required positional arguments: "input"

How can I feed in to square the tensor from the preceding layer?

Ivan · Accepted Answer · 2021-08-25 14:16:51Z

You can't put a PyTorch function in a nn.Sequential pipeline, it needs to be a nn.Module.

You could wrap it like this:

class Square(nn.Module): def forward(self, x): return torch.square(x)

Then use it inside your sequential layer like so:

class CNNClassifier(Classifier): # nn.Module def __init__(self, in_channels): super().__init__() self.save_hyperparameters('in_channels') self.cnn = nn.Sequential( nn.MaxPool2d((1, 5), stride=(1, 5)), Square(), nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=5))

Collectives™ on Stack Overflow

Use torch.square() inside a nn.Sequential layer in PyTorch

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related