How to implement keras.layers.core.Lambda in pytorch

1

add a x -> x^2 layer

model.add(Lambda(lambda x: x ** 2))

1 Like
2

That is one of the cases where it’s better not to use nn.Sequential, and inherit from model yourself and perform the operations that you want.
For example

class MyModel(nn.Module):
    def forward(self, input):
        return input ** 2 + 1
model = MyModel()

But if you want an equivalent to a Lambda layer, you can write it very easily in pytorch

class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd
    def forward(self, x):
        return self.lambd(x)

And now you can use it as you would in keras

model.add(Lambda(lambda x: x ** 2))
15 Likes
3

Your method is valid，thank you very much

4 
from torchvision import transforms
import torch

LamdaLayer = torch.jit.script(torch.nn.Sequential(
      transforms.Lambda(lambda x: <expr> ))
)