diff --git a/face_swapper/README.md b/face_swapper/README.md index d46b69c..098755e 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -70,7 +70,7 @@ kernel_size = 4 [training.model.masker] input_channels = 67 output_channels = 1 -base_channels = 16 +num_filters = 16 ``` ``` diff --git a/face_swapper/src/networks/masknet.py b/face_swapper/src/networks/masknet.py index 8e5dad5..03b9940 100644 --- a/face_swapper/src/networks/masknet.py +++ b/face_swapper/src/networks/masknet.py @@ -52,20 +52,23 @@ class BottleNeck(nn.Module): def __init__(self, num_filters : int): super().__init__() self.layers = self.create_layers(num_filters) + self.sequences = nn.Sequential(*self.layers) self.relu = nn.ReLU(inplace = True) - def create_layers(self, num_filters : int) -> nn.Sequential: - return nn.Sequential( + @staticmethod + def create_layers(num_filters : int) -> nn.ModuleList: + return nn.ModuleList( + [ nn.Conv2d(num_filters, num_filters, kernel_size = 3, padding = 1, bias = False), nn.BatchNorm2d(num_filters), nn.ReLU(inplace = True), nn.Conv2d(num_filters, num_filters, kernel_size = 3, padding = 1, bias = False), nn.BatchNorm2d(num_filters), nn.ReLU(inplace = True) - ) + ]) def forward(self, input_tensor : Tensor) -> Tensor: - output_tensor = self.layers(input_tensor) + input_tensor + output_tensor = self.sequences(input_tensor) + input_tensor output_tensor = self.relu(output_tensor) return output_tensor