change a lot

This commit is contained in:
Ray Wong 2020-09-26 17:48:26 +08:00
parent f67bcdf161
commit 776fe40199
3 changed files with 9 additions and 6 deletions

View File

@ -1,8 +1,10 @@
from model.registry import MODEL
from model.registry import MODEL, NORMALIZATION
import model.GAN.CycleGAN
import model.GAN.MUNIT
import model.GAN.TAFG
import model.GAN.UGATIT
import model.GAN.wrapper
import model.GAN.base
import model.GAN.TSIT
import model.GAN.MUNIT
import model.GAN.UGATIT
import model.GAN.base
import model.GAN.wrapper
import model.base.normalization

View File

@ -6,7 +6,7 @@ import torch.nn as nn
def select_norm_layer(norm_type):
if norm_type == "BN":
return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
return functools.partial(nn.BatchNorm2d)
elif norm_type == "IN":
return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
elif norm_type == "LN":

View File

@ -1,3 +1,4 @@
from util.registry import Registry
MODEL = Registry("model")
NORMALIZATION = Registry("normalization")