Spaces:
Runtime error
Runtime error
File size: 3,208 Bytes
47c46ea |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
from functools import partial
from torch import nn
def activation_func(activation: str):
return nn.ModuleDict([
['relu', nn.ReLU(inplace=True)],
['leaky_relu', nn.LeakyReLU(negative_slope=0.01, inplace=True)],
['selu', nn.SELU(inplace=True)],
['none', nn.Identity()]
])[activation]
def norm_module(norm: str):
return {
'batch': nn.BatchNorm2d,
'instance': nn.InstanceNorm2d,
}[norm]
class Conv2dAuto(nn.Conv2d):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# dynamic add padding based on the kernel_size
self.padding = (self.kernel_size[0] // 2, self.kernel_size[1] // 2)
conv3x3 = partial(Conv2dAuto, kernel_size=3)
class ResidualBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, activation: str = 'relu'):
super().__init__()
self.in_channels, self.out_channels = in_channels, out_channels
self.blocks = nn.Identity()
self.activate = activation_func(activation)
self.shortcut = nn.Identity()
def forward(self, x):
residual = x
if self.should_apply_shortcut:
residual = self.shortcut(x)
x = self.blocks(x)
x += residual
x = self.activate(x)
return x
@property
def should_apply_shortcut(self):
return self.in_channels != self.out_channels
class ResNetResidualBlock(ResidualBlock):
def __init__(
self, in_channels: int, out_channels: int,
expansion: int = 1, downsampling: int = 1,
conv=conv3x3, norm: str = 'batch', *args, **kwargs
):
super().__init__(in_channels, out_channels, *args, **kwargs)
self.expansion, self.downsampling = expansion, downsampling
self.conv, self.norm = conv, norm_module(norm)
self.shortcut = nn.Sequential(
nn.Conv2d(self.in_channels, self.expanded_channels, kernel_size=1,
stride=self.downsampling, bias=False),
self.norm(self.expanded_channels)) if self.should_apply_shortcut else None
@property
def expanded_channels(self):
return self.out_channels * self.expansion
@property
def should_apply_shortcut(self):
return self.in_channels != self.expanded_channels
def conv_norm(in_channels: int, out_channels: int, conv, norm, *args, **kwargs):
return nn.Sequential(conv(in_channels, out_channels, *args, **kwargs), norm(out_channels))
class ResNetBasicBlock(ResNetResidualBlock):
"""
Basic ResNet block composed by two layers of 3x3conv/batchnorm/activation
"""
expansion = 1
def __init__(
self, in_channels: int, out_channels: int, bias: bool = False, *args, **kwargs
):
super().__init__(in_channels, out_channels, *args, **kwargs)
self.blocks = nn.Sequential(
conv_norm(
self.in_channels, self.out_channels, conv=self.conv, norm=self.norm,
bias=bias, stride=self.downsampling
),
self.activate,
conv_norm(self.out_channels, self.expanded_channels, conv=self.conv, norm=self.norm, bias=bias),
)
|