Hi, I’m trying to write a custom conv backward function. This is what I’ve archived so far (I’ve delete extra code, just leave the basics)
class ConvFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, bias, stride, padding, dilation, groups, transposed=False, output_padding=[0]):
ctx.save_for_backward(input, weight, bias)
ctx.infos = {"stride": stride,
"padding": padding,
"dilation": dilation,
"transposed": transposed,
"output_padding": output_padding,
"groups": groups
}
return cudnn_convolution.convolution(input, weight, bias,
stride, padding, dilation, transposed,
output_padding, groups)
@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_variables
infos = ctx.infos
input_grad = weight_grad = bias_grad = stride_grad = padding_grad = dilation_grad = groups_grad = None
output_mask = [ctx.needs_input_grad[0],
ctx.needs_input_grad[1],
bias is not None and ctx.needs_input_grad[2],
]
input_grad, weight_grad, bias_grad = cudnn_convolution.convolution_backward(grad_output,
input,
weight,
None, # <--- Problem here, list and tuple not working either
infos["stride"],
infos["padding"],
infos["dilation"],
infos["transposed"],
infos["output_padding"],
infos["groups"],
output_mask)
return input_grad, weight_grad, bias_grad, stride_grad, padding_grad, dilation_grad, groups_grad
Following is the definition in my cpp file cudnn_convolution.cpp
std::tuple<at::Tensor, at::Tensor, at::Tensor> convolution_backward(
const at::Tensor& grad_output_,
const at::Tensor& input_,
const at::Tensor& weight_,
const at::OptionalIntArrayRef bias_sizes_opt,
at::IntArrayRef stride,
at::IntArrayRef padding,
at::IntArrayRef dilation,
bool transposed,
at::IntArrayRef output_padding,
int64_t groups,
std::array<bool, 3> output_mask) {
return at::convolution_backward(grad_output_,
input_,
weight_,
bias_sizes_opt
stride,
padding,
dilation,
transposed,
output_padding,
groups,
output_mask);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
//skip convolution
m.def("convolution_backward", &convolution_backward, "convolution backward");
}
My error is:
TypeError: convolution_backward(): incompatible function arguments. The following argument types are supported:
1. (arg0: at::Tensor, arg1: at::Tensor, arg2: at::Tensor, arg3: c10::OptionalArrayRef, arg4: at::IntArrayRef, arg5: at::IntArrayRef, arg6: at::IntArrayRef, arg7: bool, arg8: at::IntArrayRef, arg9: int, arg10: List[bool[3]]) → Tuple[at::Tensor, at::Tensor, at::Tensor]
I think the problem is the at::OptionalIntArrayRef type of bias_sizes_opt, what kind of variable in python should I pass to this type?