from torch import nn

__all__ = ['DepthwiseSeparableConv2d']


class DepthwiseSeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
        super(DepthwiseSeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels,
                                   kernel_size=kernel_size,
                                   dilation=dilation, padding=padding,
                                   stride=stride,
                                   bias=bias,
                                   groups=in_channels)
        self.pointwise = nn.Conv2d(in_channels, out_channels,
                                   kernel_size=1,
                                   groups=groups,
                                   bias=bias)

    def forward(self, x):
        out = self.depthwise(x)
        out = self.pointwise(out)
        return out
