How to use DepthWiseConv2dImplicitGEMM, I only have input channels, kernal, bais

[CccJ23]

Excuse me, how to use DepthWiseConv2dImplicitGEMM? I only have input channels, kernal, bais, how to replace ordinary nn.conv(in_channels, out_channels, kernel_size, bais), the input parameters seem unreasonable

[TheGreatTreatsby]

DeepWiseConv2dImplicitGEMM default to use the DW，so u needn‘t to worry about other arguments

[DingXiaoH]

DepthWiseConv2dImplicitGEMM is designed for depthwise conv. So it assumes in_channels=out_channels=groups.

[xg416]

I wonder what if I want the in_channels not equal to the out_channels?

I made some modifications of the original depthwise_conv2d_implicit_gemm.py file:

import os
import torch.utils.cpp_extension as cpp_extension

import _depthwise_conv2d_implicit_gemm_C as _extension


__all__ = ["DepthWiseConv2dImplicitGEMM"]


class _DepthWiseConv2dImplicitGEMMFP32(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, x, w):
        ctx.save_for_backward(x, w)
        return _extension.forward_fp32(x, w)

    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad):
        x, w = ctx.saved_tensors
        dx = _extension.backward_data_fp32(grad, w)
        dw = _extension.backward_filter_fp32(grad, x, w)
        return dx, dw


class _DepthWiseConv2dImplicitGEMMFP16(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, x, w):
        ctx.save_for_backward(x, w)
        return _extension.forward_fp16(x, w)

    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad):
        x, w = ctx.saved_tensors
        dx = _extension.backward_data_fp16(grad, w)
        dw = _extension.backward_filter_fp16(grad, x, w)
        return dx, dw


class DepthWiseConv2dImplicitGEMM(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel, bias=False, padding_mode='reflect'):
        super().__init__(in_channels, out_channels, kernel, groups=in_channels, 
                         bias=bias, padding_mode=padding_mode)
        # _load_extension()

    def forward(self, x):
        if x.dtype == torch.float32:
            x = _DepthWiseConv2dImplicitGEMMFP32.apply(x, self.weight)
        elif x.dtype == torch.float16:
            x = _DepthWiseConv2dImplicitGEMMFP16.apply(x, self.weight)
        else:
            raise TypeError("Only support fp32 and fp16, get {}".format(x.dtype))
        if self.bias is not None:
            x = x + self.bias.to(x).view(1, -1, 1, 1)
        return x

So that I can set different in_channels and out_channels, but when I ran:

x = torch.randn(4, 1, 512, 512).cuda()
conv = DepthWiseConv2dImplicitGEMM(in_channels=1, out_channels=100, kernel=33, bias=False).cuda()
y = conv(x)
print(y.shape)

I got "torch.Size([4, 1, 512, 512])", so I assume the out_channels was intrinsically reset to =in_channels, but is there a way for me to use different input and output channels?

[DingXiaoH]

Hi, our low-level code works with the depth-wise case only. If you need to set out_channels != in_channels, you need to customize the low-level code as well. Only modifying the python code will not make a difference.