Pytorch – What module should I use to multiply the output of a layer using Sequential?


While defining a neural network using nn.Module, in the forward function I can multiply the output of the final layer using:

def forward(self, x):
    x = torch.mul(x, self.max_action)
    return x

I am trying to do the same but instead using nn.Sequential method to define the neural network

model = nn.Sequential()
model.add_module(name='activation_output', module=?)

What should I use there to have the previous layer multiply by the scalar self.max_action ?
Or should I build the sequential model in a different way ?

Asked By: T4l0n



You could define a custom nn.Module layer:

class Multiply(nn.Module):
    def __init__(self, alpha):
        self.alpha =  alpha
    def forward(self, x):
        x = torch.mul(x, self.alpha)
        return x

Then use it as:

>>> model = nn.Sequential()
>>> model.add_module(name='activation_output', module=Multiply(10))

>>> model(torch.ones(1))
Answered By: Ivan
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.