Spaces:
Running
Running
File size: 1,467 Bytes
cec5823 |
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 |
import torch
from rstor.properties import LEAKY_RELU
from rstor.architecture.base import get_non_linearity
class BaseConvolutionBlock(torch.nn.Module):
def __init__(
self,
ch_in: int,
ch_out: int,
k_size: int,
activation=LEAKY_RELU,
bias: bool = True
) -> None:
super().__init__()
self.conv = torch.nn.Conv2d(ch_in, ch_out, k_size, padding=k_size//2, bias=bias)
self.non_linearity = get_non_linearity(activation)
self.conv_non_lin = torch.nn.Sequential(self.conv, self.non_linearity)
def forward(self, x_in: torch.Tensor) -> torch.Tensor:
return self.conv_non_lin(x_in)
class ResConvolutionBlock(torch.nn.Module):
def __init__(
self,
ch_in: int,
ch_out: int,
k_size: int,
activation=LEAKY_RELU,
bias: bool = True,
residual: bool = True
) -> None:
super().__init__()
self.conv1 = torch.nn.Conv2d(ch_in, ch_out, k_size, padding=k_size//2, bias=bias)
self.non_linearity = get_non_linearity(activation)
self.conv2 = torch.nn.Conv2d(ch_out, ch_out, k_size, padding=k_size//2, bias=bias)
self.residual = residual
def forward(self, x_in: torch.Tensor) -> torch.Tensor:
y = self.conv1(x_in)
y = self.non_linearity(y)
y = self.conv2(y)
if self.residual:
y = x_in + y
y = self.non_linearity(y)
return y
|