How do I implement my own activation function with the nn module?

1

I was looking at the tutorial (http://pytorch.org/tutorials/beginner/examples_nn/two_layer_net_nn.html#sphx-glr-beginner-examples-nn-two-layer-net-nn-py)

and there was a very nice line were the define the model:

model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out),
)

I wanted to change it to:

model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.My_Activation(),
    torch.nn.Linear(H, D_out),
)

say for the sake of an example that My_Activation is the element-wise quadratic (cuz why not something simple for an example?). Is it possible to write such code for custom nn modules?

1 Like
2

Have you tried simply writing it direclty as a mathematical formula in the forward method?