lama

1
import torch
2
import torch.nn as nn
3

4
class DepthWiseSeperableConv(nn.Module):
5
    def __init__(self, in_dim, out_dim, *args, **kwargs):
6
        super().__init__()
7
        if 'groups' in kwargs:
8
            # ignoring groups for Depthwise Sep Conv
9
            del kwargs['groups']
10
        
11
        self.depthwise = nn.Conv2d(in_dim, in_dim, *args, groups=in_dim, **kwargs)
12
        self.pointwise = nn.Conv2d(in_dim, out_dim, kernel_size=1)
13
        
14
    def forward(self, x):
15
        out = self.depthwise(x)
16
        out = self.pointwise(out)
17
        return out