From 9e8e73c988d16c787e6fa3741004a6a1f5727d96 Mon Sep 17 00:00:00 2001 From: budui Date: Fri, 28 Aug 2020 08:16:07 +0800 Subject: [PATCH] move norm select to top --- model/GAN/residual_generator.py | 19 ++++--------------- model/normalization.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/model/GAN/residual_generator.py b/model/GAN/residual_generator.py index ea4abbc..8a288e0 100644 --- a/model/GAN/residual_generator.py +++ b/model/GAN/residual_generator.py @@ -1,18 +1,7 @@ import torch import torch.nn as nn -import functools from model.registry import MODEL - - -def _select_norm_layer(norm_type): - if norm_type == "BN": - return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) - elif norm_type == "IN": - return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) - elif norm_type == "NONE": - return lambda x: nn.Identity() - else: - raise NotImplemented(f'normalization layer {norm_type} is not found') +from model.normalization import select_norm_layer class GANImageBuffer(object): @@ -77,7 +66,7 @@ class ResidualBlock(nn.Module): if use_bias is None: # Only for IN, use bias since it does not have affine parameters. use_bias = norm_type == "IN" - norm_layer = _select_norm_layer(norm_type) + norm_layer = select_norm_layer(norm_type) models = [nn.Sequential( nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias), norm_layer(num_channels), @@ -101,7 +90,7 @@ class ResGenerator(nn.Module): norm_type="IN", use_dropout=False): super(ResGenerator, self).__init__() assert num_blocks >= 0, f'Number of residual blocks must be non-negative, but got {num_blocks}.' - norm_layer = _select_norm_layer(norm_type) + norm_layer = select_norm_layer(norm_type) use_bias = norm_type == "IN" self.start_conv = nn.Sequential( @@ -157,7 +146,7 @@ class PatchDiscriminator(nn.Module): def __init__(self, in_channels, base_channels=64, num_conv=3, norm_type="IN"): super(PatchDiscriminator, self).__init__() assert num_conv >= 0, f'Number of conv blocks must be non-negative, but got {num_conv}.' - norm_layer = _select_norm_layer(norm_type) + norm_layer = select_norm_layer(norm_type) use_bias = norm_type == "IN" kernel_size = 4 diff --git a/model/normalization.py b/model/normalization.py index e69de29..36413c0 100644 --- a/model/normalization.py +++ b/model/normalization.py @@ -0,0 +1,13 @@ +import torch.nn as nn +import functools + + +def select_norm_layer(norm_type): + if norm_type == "BN": + return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) + elif norm_type == "IN": + return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) + elif norm_type == "NONE": + return lambda x: nn.Identity() + else: + raise NotImplemented(f'normalization layer {norm_type} is not found')