Pytorch c++ extension failed

Hello everyone,
I am writing a simple python c++ extension for Linear. But I encountered a problem: :sob:
My C++ extension code is as follows:

#include<vector>
#include<torch/extension.h>

torch::Tensor d_sigmoid(const torch::Tensor &z)
{
    auto s = torch::sigmoid(z);
    return (1 - s) * s;
}

torch::Tensor dense_forward(const torch::Tensor &input,const torch::Tensor &weights,const torch::Tensor &bias)
{
    auto output = torch::mm(input,weights.t()) + bias; //mm(input,weights.t()) eg. input shape 1*2;  weights shape 4*2; output shape 1*4; bias 1*4
    output = torch::sigmoid(output);
    return output;
}

std::vector<torch::Tensor> dense_backward(const torch::Tensor &grad_output, const torch::Tensor &input, const torch::Tensor &output,const torch::Tensor &weights, const torch::Tensor &bias)
{
    auto output_d_sigmoid = d_sigmoid(output);
    auto grad = grad_output * output_d_sigmoid;
    auto grad_weights = torch::mm(grad.t(),input); // 本层的权重的梯度
    auto grad_bias = grad.sum(0,/*keepdim=*/false);                //偏置层梯度
    auto grad_input = torch::mm(grad,weights);     // 传给前一层的梯度
    return {grad_input, grad_weights, grad_bias};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("forward", &dense_forward, "dense forward");
    m.def("backward", &dense_backward,"dense backward");
}

My python code is as follows:

import torch
import torch.nn as nn
from torch.autograd import Variable
import math

import dnn

class densefunction(torch.autograd.Function):
    
    @staticmethod 
    def forward(ctx, x, weights, bias):
        outputs = dnn.forward(x, weights, bias)
        ctx.save_for_backward(x, weights, bias, outputs)
        return outputs

    @staticmethod 
    def backward(ctx,grad_output):
        x, weight, bias, output = ctx.saved_tensors
        grad_input, grad_weight, grad_bias = dnn.backward(grad_output, x, output, weight, bias)
        return grad_input, grad_weight, grad_bias

class dense(nn.Module):
    def __init__ (self,input_features, output_features):
        super(dense,self).__init__()
        self.input_features = input_features
        self.output_features = output_features
        self.weight = nn.Parameter(torch.empty(output_features,input_features))
        self.bias = nn.Parameter(torch.empty(output_features))

        self.weight.data.uniform_(-0.1,0.1)
        self.bias.data.uniform_(-0.1,0.1)
        
    def forward(self,x):
        return densefunction.apply(x,self.weight,self.bias)


if __name__ == '__main__':
    
    x = torch.randn((4,5))
    dnn = dense(5,3)
    print(dnn(x))

When running the code,I encountered a problem:

Traceback (most recent call last):
  File "main.py", line 41, in <module>
    print(dnn(x))
  File "/data/zh/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "main.py", line 34, in forward
    return densefunction.apply(x,self.weight,self.bias)
  File "main.py", line 12, in forward
    outputs = dnn.forward(x, weights, bias)
TypeError: forward() takes 2 positional arguments but 4 were given

Hope someone help me to solve this problem. Thanks!

Hi,

Are you sure that you recompiled the cpp module with the latest version? And that you don’t have a cached version with an older version of your code where the c++ module only took 2 arguments?

Thanks for your reply. I am sure that I don’t have a cached version with an older version of your code where the c++ module only took 2 argument. And the densefunction.apply() works alone. The process is as follows:

    x = torch.randn((4,5))
    y_label = torch.randn((4,3))

    weight = torch.randn((3,5),requires_grad=True)
    bias = torch.randn(3,requires_grad=True)
    output = densefunction.apply(x, weight, bias)

    loss = (output - y_label).pow(2).sum()
    loss.backward()
    print(weight.grad)
    print(bias.grad)

The result is as follows:

#print(weight.grad)
tensor([[ 0.4120, -0.3259, -2.5679, -0.8549, -1.7054],
        [ 0.2807, -0.0275, -0.4202,  1.0669,  0.7516],
        [ 1.0337,  1.1368, -1.2103,  1.2529, -0.1512]])
#print(bias.grad)
tensor([1.7098, 0.9224, 1.0619])

When I changed the c++ module to another name, the previous error disappeared. I guess there are other dnn.py files in my system, and this module was imported by mistake when I used import.

1 Like