I want to implement a custom activation function with learnable parameters. For example one that takes the input x and returns a polinomial of specified order, of x.
def poli_activation(x, order=2):
input_tens = []
# is this the way to make coeff a vector of parameters?
coeff = torch.nn.parameter(torch.randn(order+1))
# need a vector of powers of x , for example (x^2, x, 1)
for idx in range(order):
element = torch.pow(x, idx)
input_tens.append(element)
# make this vector a variable and implement polynomial
# as dot product.
input_tens = torch.Tensor(input_tens)
input_tens = torch.autograd.Variable(input_tens.cuda())
output = torch.dot(coeff, input_tens)
return output
Now this is 1) a very inefficient implementation of such a function and 2) does not work. Does anyone have an idea of how this can be efficiently implemented in pytorch? I want the network to be able to learn the polynomial coefficients.