How to define a new layer with autograd?
Hi all,
big_tree
Sure, this will be handled for you. For example: <SCODE>import torch.nn as nn
from torch.autograd import Variable

class Gaussian(nn.Module):
    def __init__(self):
        self.a = nn.Parameter(torch.zeros(1))
        self.b = nn.Parameter(torch.zeros(1))
        self.c = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        # unfortunately we don't have automatic broadcasting yet
        a = self.a.expand_as(x)
        b = self.b.expand_as(x)
        c = self.c.expand_as(x)
        return a * torch.exp((x - b)^2 / c)

module = Gaussian()
x = Variable(torch.randn(20))
out = module(x)
loss = loss_fn(out)
loss.backward()

# Now module.a.grad should be non-zero.
<ECODE>
apaszke
to update all the parameters?
big_tree
apaszke
Thank you so much for your help!
big_tree
Now we have defined the Gaussian layer, I want to know how to use it in a network of several layers. For example, <SCODE>class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(100, 20)
        self.gauss = Gaussian()
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.gauss(x)

net = MyNet()
<ECODE> Is the above code correct? Or do we have to set up other things in order to use the newly-defined layer in other network?
jdhao
the above is correct.
smth
Can I define my own layer with a forward/backward function or I will have to define the layer as a Function so that I can define the backward function too?
Soumava_Roy
Hao_Chu
herleeyandi
jdhao
herleeyandi
Most of the time we do not need to extend PyTorch using numpy. As long as you use the builtin method of Variable, you can only write forward method and backward gradient computation is handled by autograd. So using a composition of builtin Variable method to achieve what you want is more time-saving.
jdhao
colesbury
herleeyandi
Yes, you can do that. Writing new CUDA kernels usually requires a lot of effort. If you can express your layer in terms of existing Tensor operations, then that’s usually a better way to get started. If you can’t do that, then you might have to write new kernels.
colesbury
What if the forward function is not differentiable? How should I update the parameters and return the gradient w.r.t inputs?
MahdiNazemi
I tried the piece of code posted by apaszke, but I got the following error. <SCODE>RuntimeError: bitwise_xor(): functions with out=... arguments don't support automatic differentiation, but one of the arguments requires grad.
<ECODE>
jwang9
I figured that the problem arises from the power calculation in python. I changed ^2 to **2 and it works now.
jwang9