Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,323 Bytes
1ea89dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import torch
import torch.nn as nn
class CvnxtBlock(nn.Module):
def __init__(
self,
dim,
kernel_size=7,
layer_scale=1.0,
expansion=4,
dilation=1,
padding_mode: str = "zeros",
):
super().__init__()
self.dwconv = nn.Conv2d(
dim,
dim,
kernel_size=kernel_size,
padding=dilation * (kernel_size - 1) // 2,
groups=dim,
dilation=dilation,
padding_mode=padding_mode,
) # depthwise conv
self.norm = nn.LayerNorm(dim)
self.pwconv1 = nn.Linear(dim, expansion * dim)
self.act = nn.GELU()
self.pwconv2 = nn.Linear(expansion * dim, dim)
self.gamma = (
nn.Parameter(layer_scale * torch.ones(1, dim, 1, 1))
if layer_scale > 0.0
else 1.0
)
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, x):
input = x
x = self.dwconv(x)
x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
return self.skip_add.add(self.gamma * x.permute(0, 3, 1, 2), input)
class SimpleCvnxtBlock(nn.Module):
def __init__(
self,
dim,
output_dim=None,
kernel_size=7,
expansion=4,
dilation=1,
padding_mode: str = "zeros",
):
super().__init__()
output_dim = output_dim if output_dim is not None else dim
self.dwconv = nn.Conv2d(
dim,
dim,
kernel_size=kernel_size,
padding=dilation * (kernel_size - 1) // 2,
groups=dim,
dilation=dilation,
padding_mode=padding_mode,
) # depthwise conv
self.norm = nn.LayerNorm(dim)
self.pwconv1 = nn.Linear(dim, expansion * dim)
self.act = nn.GELU()
self.pwconv2 = nn.Linear(expansion * dim, output_dim)
def forward(self, x):
x = self.dwconv(x)
x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
return x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|