From d875ca83426f5472fa355e5b10a650660d115bb8 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 18 Dec 2020 09:24:31 -0700 Subject: [PATCH] More refactor changes --- codes/data/byol_attachment.py | 2 +- codes/data/stylegan2_dataset.py | 2 +- codes/models/ProgressiveSrg_arch.py | 4 +- codes/models/RRDBNet_arch.py | 2 +- .../models/SwitchedResidualGenerator_arch.py | 8 ++-- codes/models/byol/byol_structural.py | 2 +- codes/models/discriminator_vgg_arch.py | 6 +-- codes/models/spinenet_arch.py | 2 +- .../{srflow_orig => srflow}/FlowActNorms.py | 2 +- .../FlowAffineCouplingsAblation.py | 4 +- .../{srflow_orig => srflow}/FlowStep.py | 10 ++--- .../FlowUpsamplerNet.py | 16 ++++---- .../{srflow_orig => srflow}/Permutations.py | 2 +- .../{srflow_orig => srflow}/RRDBNet_arch.py | 4 +- .../{srflow_orig => srflow}/SRFlowNet_arch.py | 8 ++-- codes/models/{srflow_orig => srflow}/Split.py | 4 +- .../{srflow_orig => srflow}/__init__.py | 0 codes/models/{srflow_orig => srflow}/flow.py | 2 +- .../{srflow_orig => srflow}/glow_arch.py | 0 .../{srflow_orig => srflow}/module_util.py | 0 codes/models/{srflow_orig => srflow}/thops.py | 0 codes/models/srg2_classic.py | 6 +-- codes/models/stylegan/__init__.py | 2 +- codes/models/tecogan/teco_resgen.py | 2 +- codes/scripts/byol_extract_wrapped_model.py | 2 +- codes/scripts/byol_spinenet_playground.py | 2 +- .../stylegan2/convert_weights_rosinality.py | 2 +- .../stereoscopic.py | 4 +- .../tecogan_losses.py | 4 +- codes/trainer/eval/__init__.py | 2 +- codes/trainer/eval/flow_gaussian_nll.py | 2 +- codes/trainer/injectors.py | 4 +- codes/trainer/losses.py | 4 +- codes/trainer/networks.py | 40 +++++++++---------- codes/utils/numeric_stability.py | 4 +- 35 files changed, 80 insertions(+), 80 deletions(-) rename codes/models/{srflow_orig => srflow}/FlowActNorms.py (99%) rename codes/models/{srflow_orig => srflow}/FlowAffineCouplingsAblation.py (97%) rename codes/models/{srflow_orig => srflow}/FlowStep.py (91%) rename codes/models/{srflow_orig => srflow}/FlowUpsamplerNet.py (96%) rename codes/models/{srflow_orig => srflow}/Permutations.py (97%) rename codes/models/{srflow_orig => srflow}/RRDBNet_arch.py (98%) rename codes/models/{srflow_orig => srflow}/SRFlowNet_arch.py (97%) rename codes/models/{srflow_orig => srflow}/Split.py (95%) rename codes/models/{srflow_orig => srflow}/__init__.py (100%) rename codes/models/{srflow_orig => srflow}/flow.py (98%) rename codes/models/{srflow_orig => srflow}/glow_arch.py (100%) rename codes/models/{srflow_orig => srflow}/module_util.py (100%) rename codes/models/{srflow_orig => srflow}/thops.py (100%) diff --git a/codes/data/byol_attachment.py b/codes/data/byol_attachment.py index ace0b2c1..f5dda78b 100644 --- a/codes/data/byol_attachment.py +++ b/codes/data/byol_attachment.py @@ -15,7 +15,7 @@ import torch.nn.functional as F from tqdm import tqdm from data import create_dataset -from models.archs.arch_util import PixelUnshuffle +from models.arch_util import PixelUnshuffle from utils.util import opt_get diff --git a/codes/data/stylegan2_dataset.py b/codes/data/stylegan2_dataset.py index c52045d8..4281267f 100644 --- a/codes/data/stylegan2_dataset.py +++ b/codes/data/stylegan2_dataset.py @@ -9,7 +9,7 @@ from torchvision import transforms import torch.nn as nn from pathlib import Path -import models.archs.stylegan.stylegan2_lucidrains as sg2 +import models.stylegan.stylegan2_lucidrains as sg2 def convert_transparent_to_rgb(image): diff --git a/codes/models/ProgressiveSrg_arch.py b/codes/models/ProgressiveSrg_arch.py index f2ff9441..7e5c9dde 100644 --- a/codes/models/ProgressiveSrg_arch.py +++ b/codes/models/ProgressiveSrg_arch.py @@ -1,9 +1,9 @@ -import models.archs.SwitchedResidualGenerator_arch as srg +import models.SwitchedResidualGenerator_arch as srg import torch import torch.nn as nn from switched_conv.switched_conv_util import save_attention_to_image from switched_conv.switched_conv import compute_attention_specificity -from models.archs.arch_util import ConvGnLelu, ExpansionBlock, MultiConvBlock +from models.arch_util import ConvGnLelu, ExpansionBlock, MultiConvBlock import functools import torch.nn.functional as F diff --git a/codes/models/RRDBNet_arch.py b/codes/models/RRDBNet_arch.py index 229b5a76..6a43f158 100644 --- a/codes/models/RRDBNet_arch.py +++ b/codes/models/RRDBNet_arch.py @@ -5,7 +5,7 @@ import torch.nn as nn import torch.nn.functional as F import torchvision -from models.archs.arch_util import make_layer, default_init_weights, ConvGnSilu, ConvGnLelu +from models.arch_util import make_layer, default_init_weights, ConvGnSilu, ConvGnLelu from utils.util import checkpoint, sequential_checkpoint diff --git a/codes/models/SwitchedResidualGenerator_arch.py b/codes/models/SwitchedResidualGenerator_arch.py index 6fb49c13..feb20733 100644 --- a/codes/models/SwitchedResidualGenerator_arch.py +++ b/codes/models/SwitchedResidualGenerator_arch.py @@ -1,14 +1,14 @@ import torch from torch import nn -from switched_conv.switched_conv import BareConvSwitch, compute_attention_specificity, AttentionNorm +from models.switched_conv.switched_conv import BareConvSwitch, compute_attention_specificity, AttentionNorm import torch.nn.functional as F import functools from collections import OrderedDict -from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, ExpansionBlock2, ConvGnLelu, MultiConvBlock, \ +from models.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, ExpansionBlock2, ConvGnLelu, MultiConvBlock, \ SiLU, UpconvBlock, ReferenceJoinBlock -from switched_conv.switched_conv_util import save_attention_to_image_rgb +from models.switched_conv.switched_conv_util import save_attention_to_image_rgb import os -from models.archs.spinenet_arch import SpineNet +from models.spinenet_arch import SpineNet import torchvision from utils.util import checkpoint diff --git a/codes/models/byol/byol_structural.py b/codes/models/byol/byol_structural.py index 80a9af75..eb56a5d7 100644 --- a/codes/models/byol/byol_structural.py +++ b/codes/models/byol/byol_structural.py @@ -5,7 +5,7 @@ import torch.nn.functional as F from torch import nn from data.byol_attachment import reconstructed_shared_regions -from models.archs.byol.byol_model_wrapper import singleton, EMA, get_module_device, set_requires_grad, \ +from models.byol.byol_model_wrapper import singleton, EMA, get_module_device, set_requires_grad, \ update_moving_average from utils.util import checkpoint diff --git a/codes/models/discriminator_vgg_arch.py b/codes/models/discriminator_vgg_arch.py index 81af1f6b..3917ddc7 100644 --- a/codes/models/discriminator_vgg_arch.py +++ b/codes/models/discriminator_vgg_arch.py @@ -1,10 +1,10 @@ import torch import torch.nn as nn -from models.archs.RRDBNet_arch import RRDB, RRDBWithBypass -from models.archs.arch_util import ConvBnLelu, ConvGnLelu, ExpansionBlock, ConvGnSilu, ResidualBlockGN +from models.RRDBNet_arch import RRDB, RRDBWithBypass +from models.arch_util import ConvBnLelu, ConvGnLelu, ExpansionBlock, ConvGnSilu, ResidualBlockGN import torch.nn.functional as F -from models.archs.SwitchedResidualGenerator_arch import gather_2d +from models.SwitchedResidualGenerator_arch import gather_2d from utils.util import checkpoint diff --git a/codes/models/spinenet_arch.py b/codes/models/spinenet_arch.py index c16ef2a1..8bd723ef 100644 --- a/codes/models/spinenet_arch.py +++ b/codes/models/spinenet_arch.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from torch.nn.init import kaiming_normal from torchvision.models.resnet import BasicBlock, Bottleneck -from models.archs.arch_util import ConvGnSilu, ConvBnSilu, ConvBnRelu +from models.arch_util import ConvGnSilu, ConvBnSilu, ConvBnRelu def constant_init(module, val, bias=0): diff --git a/codes/models/srflow_orig/FlowActNorms.py b/codes/models/srflow/FlowActNorms.py similarity index 99% rename from codes/models/srflow_orig/FlowActNorms.py rename to codes/models/srflow/FlowActNorms.py index 6c1ef4e6..b37c3485 100644 --- a/codes/models/srflow_orig/FlowActNorms.py +++ b/codes/models/srflow/FlowActNorms.py @@ -1,7 +1,7 @@ import torch from torch import nn as nn -from models.srflow_orig import thops +from models.srflow import thops class _ActNorm(nn.Module): diff --git a/codes/models/srflow_orig/FlowAffineCouplingsAblation.py b/codes/models/srflow/FlowAffineCouplingsAblation.py similarity index 97% rename from codes/models/srflow_orig/FlowAffineCouplingsAblation.py rename to codes/models/srflow/FlowAffineCouplingsAblation.py index 0eddbc8d..62e8d41e 100644 --- a/codes/models/srflow_orig/FlowAffineCouplingsAblation.py +++ b/codes/models/srflow/FlowAffineCouplingsAblation.py @@ -1,8 +1,8 @@ import torch from torch import nn as nn -from models.srflow_orig import thops -from models.archs.srflow_orig.flow import Conv2d, Conv2dZeros +from models.srflow import thops +from models.srflow.flow import Conv2d, Conv2dZeros from utils.util import opt_get diff --git a/codes/models/srflow_orig/FlowStep.py b/codes/models/srflow/FlowStep.py similarity index 91% rename from codes/models/srflow_orig/FlowStep.py rename to codes/models/srflow/FlowStep.py index 174acdea..9f6e4f8d 100644 --- a/codes/models/srflow_orig/FlowStep.py +++ b/codes/models/srflow/FlowStep.py @@ -1,7 +1,7 @@ import torch from torch import nn as nn -import models.archs.srflow_orig.Permutations +import models.srflow.Permutations def getConditional(rrdbResults, position): @@ -44,17 +44,17 @@ class FlowStep(nn.Module): self.acOpt = acOpt # 1. actnorm - self.actnorm = models.archs.srflow_orig.FlowActNorms.ActNorm2d(in_channels, actnorm_scale) + self.actnorm = models.srflow.FlowActNorms.ActNorm2d(in_channels, actnorm_scale) # 2. permute if flow_permutation == "invconv": - self.invconv = models.archs.srflow_orig.Permutations.InvertibleConv1x1( + self.invconv = models.srflow.Permutations.InvertibleConv1x1( in_channels, LU_decomposed=LU_decomposed) # 3. coupling if flow_coupling == "CondAffineSeparatedAndCond": - self.affine = models.archs.srflow_orig.FlowAffineCouplingsAblation.CondAffineSeparatedAndCond(in_channels=in_channels, - opt=opt) + self.affine = models.srflow.FlowAffineCouplingsAblation.CondAffineSeparatedAndCond(in_channels=in_channels, + opt=opt) elif flow_coupling == "noCoupling": pass else: diff --git a/codes/models/srflow_orig/FlowUpsamplerNet.py b/codes/models/srflow/FlowUpsamplerNet.py similarity index 96% rename from codes/models/srflow_orig/FlowUpsamplerNet.py rename to codes/models/srflow/FlowUpsamplerNet.py index beb0122f..387ebbef 100644 --- a/codes/models/srflow_orig/FlowUpsamplerNet.py +++ b/codes/models/srflow/FlowUpsamplerNet.py @@ -2,12 +2,12 @@ import numpy as np import torch from torch import nn as nn -import models.archs.srflow_orig.Split -from models.archs.srflow_orig import flow -from models.srflow_orig import thops -from models.archs.srflow_orig.Split import Split2d -from models.archs.srflow_orig.glow_arch import f_conv2d_bias -from models.archs.srflow_orig.FlowStep import FlowStep +import models.srflow.Split +from models.srflow import flow +from models.srflow import thops +from models.srflow.Split import Split2d +from models.srflow.glow_arch import f_conv2d_bias +from models.srflow.FlowStep import FlowStep from utils.util import opt_get, checkpoint @@ -146,8 +146,8 @@ class FlowUpsamplerNet(nn.Module): t = opt_get(opt, ['networks', 'generator','flow', 'split', 'type'], 'Split2d') if t == 'Split2d': - split = models.archs.srflow_orig.Split.Split2d(num_channels=self.C, logs_eps=logs_eps, position=position, - cond_channels=cond_channels, consume_ratio=consume_ratio, opt=opt) + split = models.srflow.Split.Split2d(num_channels=self.C, logs_eps=logs_eps, position=position, + cond_channels=cond_channels, consume_ratio=consume_ratio, opt=opt) self.layers.append(split) self.output_shapes.append([-1, split.num_channels_pass, H, W]) self.C = split.num_channels_pass diff --git a/codes/models/srflow_orig/Permutations.py b/codes/models/srflow/Permutations.py similarity index 97% rename from codes/models/srflow_orig/Permutations.py rename to codes/models/srflow/Permutations.py index 7a0ea276..ba8d85f4 100644 --- a/codes/models/srflow_orig/Permutations.py +++ b/codes/models/srflow/Permutations.py @@ -3,7 +3,7 @@ import torch from torch import nn as nn from torch.nn import functional as F -from models.srflow_orig import thops +from models.srflow import thops class InvertibleConv1x1(nn.Module): diff --git a/codes/models/srflow_orig/RRDBNet_arch.py b/codes/models/srflow/RRDBNet_arch.py similarity index 98% rename from codes/models/srflow_orig/RRDBNet_arch.py rename to codes/models/srflow/RRDBNet_arch.py index 1fff45c1..36eee072 100644 --- a/codes/models/srflow_orig/RRDBNet_arch.py +++ b/codes/models/srflow/RRDBNet_arch.py @@ -2,8 +2,8 @@ import functools import torch import torch.nn as nn import torch.nn.functional as F -import models.archs.srflow_orig.module_util as mutil -from models.archs.arch_util import default_init_weights, ConvGnSilu, ConvGnLelu +import models.srflow.module_util as mutil +from models.arch_util import default_init_weights, ConvGnSilu, ConvGnLelu from utils.util import opt_get diff --git a/codes/models/srflow_orig/SRFlowNet_arch.py b/codes/models/srflow/SRFlowNet_arch.py similarity index 97% rename from codes/models/srflow_orig/SRFlowNet_arch.py rename to codes/models/srflow/SRFlowNet_arch.py index f91944bb..73314407 100644 --- a/codes/models/srflow_orig/SRFlowNet_arch.py +++ b/codes/models/srflow/SRFlowNet_arch.py @@ -4,10 +4,10 @@ import torch import torch.nn as nn import torch.nn.functional as F import numpy as np -from models.archs.srflow_orig.RRDBNet_arch import RRDBNet -from models.archs.srflow_orig.FlowUpsamplerNet import FlowUpsamplerNet -import models.srflow_orig.thops as thops -import models.archs.srflow_orig.flow as flow +from models.srflow.RRDBNet_arch import RRDBNet +from models.srflow.FlowUpsamplerNet import FlowUpsamplerNet +import models.srflow.thops as thops +import models.srflow.flow as flow from utils.util import opt_get diff --git a/codes/models/srflow_orig/Split.py b/codes/models/srflow/Split.py similarity index 95% rename from codes/models/srflow_orig/Split.py rename to codes/models/srflow/Split.py index 28bae29f..c3a1ffd5 100644 --- a/codes/models/srflow_orig/Split.py +++ b/codes/models/srflow/Split.py @@ -1,8 +1,8 @@ import torch from torch import nn as nn -from models.srflow_orig import thops -from models.archs.srflow_orig.flow import Conv2dZeros, GaussianDiag +from models.srflow import thops +from models.srflow.flow import Conv2dZeros, GaussianDiag from utils.util import opt_get diff --git a/codes/models/srflow_orig/__init__.py b/codes/models/srflow/__init__.py similarity index 100% rename from codes/models/srflow_orig/__init__.py rename to codes/models/srflow/__init__.py diff --git a/codes/models/srflow_orig/flow.py b/codes/models/srflow/flow.py similarity index 98% rename from codes/models/srflow_orig/flow.py rename to codes/models/srflow/flow.py index 71e03646..d5c06277 100644 --- a/codes/models/srflow_orig/flow.py +++ b/codes/models/srflow/flow.py @@ -3,7 +3,7 @@ import torch.nn as nn import torch.nn.functional as F import numpy as np -from models.archs.srflow_orig.FlowActNorms import ActNorm2d +from models.srflow.FlowActNorms import ActNorm2d from . import thops diff --git a/codes/models/srflow_orig/glow_arch.py b/codes/models/srflow/glow_arch.py similarity index 100% rename from codes/models/srflow_orig/glow_arch.py rename to codes/models/srflow/glow_arch.py diff --git a/codes/models/srflow_orig/module_util.py b/codes/models/srflow/module_util.py similarity index 100% rename from codes/models/srflow_orig/module_util.py rename to codes/models/srflow/module_util.py diff --git a/codes/models/srflow_orig/thops.py b/codes/models/srflow/thops.py similarity index 100% rename from codes/models/srflow_orig/thops.py rename to codes/models/srflow/thops.py diff --git a/codes/models/srg2_classic.py b/codes/models/srg2_classic.py index b33cccca..bece2781 100644 --- a/codes/models/srg2_classic.py +++ b/codes/models/srg2_classic.py @@ -8,9 +8,9 @@ import torch.nn.functional as F import functools from collections import OrderedDict -from models.archs.SwitchedResidualGenerator_arch import HalvingProcessingBlock, ConfigurableSwitchComputer -from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, MultiConvBlock -from switched_conv.switched_conv import BareConvSwitch, AttentionNorm +from models.SwitchedResidualGenerator_arch import HalvingProcessingBlock, ConfigurableSwitchComputer +from models.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, MultiConvBlock +from models.switched_conv.switched_conv import BareConvSwitch, AttentionNorm from utils.util import checkpoint diff --git a/codes/models/stylegan/__init__.py b/codes/models/stylegan/__init__.py index ab92d1d5..444808b8 100644 --- a/codes/models/stylegan/__init__.py +++ b/codes/models/stylegan/__init__.py @@ -1,4 +1,4 @@ -import models.archs.stylegan.stylegan2_lucidrains as stylegan2 +import models.stylegan.stylegan2_lucidrains as stylegan2 def create_stylegan2_loss(opt_loss, env): diff --git a/codes/models/tecogan/teco_resgen.py b/codes/models/tecogan/teco_resgen.py index 06d93fd4..9803df33 100644 --- a/codes/models/tecogan/teco_resgen.py +++ b/codes/models/tecogan/teco_resgen.py @@ -5,7 +5,7 @@ import torch.nn as nn import torchvision from utils.util import sequential_checkpoint -from models.archs.arch_util import ConvGnSilu, make_layer +from models.arch_util import ConvGnSilu, make_layer class TecoResblock(nn.Module): diff --git a/codes/scripts/byol_extract_wrapped_model.py b/codes/scripts/byol_extract_wrapped_model.py index f5e80c54..12818d31 100644 --- a/codes/scripts/byol_extract_wrapped_model.py +++ b/codes/scripts/byol_extract_wrapped_model.py @@ -1,6 +1,6 @@ import torch -from models.archs.spinenet_arch import SpineNet +from models.spinenet_arch import SpineNet if __name__ == '__main__': pretrained_path = '../../experiments/train_sbyol_512unsupervised_restart/models/48000_generator.pth' diff --git a/codes/scripts/byol_spinenet_playground.py b/codes/scripts/byol_spinenet_playground.py index c28b955e..3c8a8e61 100644 --- a/codes/scripts/byol_spinenet_playground.py +++ b/codes/scripts/byol_spinenet_playground.py @@ -13,7 +13,7 @@ import numpy as np import utils from data.image_folder_dataset import ImageFolderDataset -from models.archs.spinenet_arch import SpineNet +from models.spinenet_arch import SpineNet # Computes the structural euclidean distance between [x,y]. "Structural" here means the [h,w] dimensions are preserved diff --git a/codes/scripts/stylegan2/convert_weights_rosinality.py b/codes/scripts/stylegan2/convert_weights_rosinality.py index a2154669..2b9400fb 100644 --- a/codes/scripts/stylegan2/convert_weights_rosinality.py +++ b/codes/scripts/stylegan2/convert_weights_rosinality.py @@ -13,7 +13,7 @@ import torch import numpy as np from torchvision import utils -from models.archs.stylegan.stylegan2_rosinality import Generator, Discriminator +from models.stylegan.stylegan2_rosinality import Generator, Discriminator # Converts from the TF state_dict input provided into the vars originally expected from the rosinality converter. diff --git a/codes/trainer/custom_training_components/stereoscopic.py b/codes/trainer/custom_training_components/stereoscopic.py index 4dac8034..95b7945f 100644 --- a/codes/trainer/custom_training_components/stereoscopic.py +++ b/codes/trainer/custom_training_components/stereoscopic.py @@ -1,7 +1,7 @@ import torch from torch.cuda.amp import autocast -from models.archs.flownet2.networks import Resample2d -from models.archs.flownet2 import flow2img +from models.flownet2.networks import Resample2d +from models.flownet2 import flow2img from trainer.injectors import Injector diff --git a/codes/trainer/custom_training_components/tecogan_losses.py b/codes/trainer/custom_training_components/tecogan_losses.py index 687ac037..d093240f 100644 --- a/codes/trainer/custom_training_components/tecogan_losses.py +++ b/codes/trainer/custom_training_components/tecogan_losses.py @@ -1,8 +1,8 @@ from torch.cuda.amp import autocast -from models.archs.stylegan.stylegan2_lucidrains import gradient_penalty +from models.stylegan.stylegan2_lucidrains import gradient_penalty from trainer.losses import ConfigurableLoss, GANLoss, extract_params_from_state, get_basic_criterion_for_name -from models.archs.flownet2.networks import Resample2d +from models.flownet2.networks import Resample2d from trainer.injectors import Injector import torch import torch.nn.functional as F diff --git a/codes/trainer/eval/__init__.py b/codes/trainer/eval/__init__.py index 1b9ffb89..e8338096 100644 --- a/codes/trainer/eval/__init__.py +++ b/codes/trainer/eval/__init__.py @@ -1,6 +1,6 @@ from trainer.eval.flow_gaussian_nll import FlowGaussianNll from trainer.eval.sr_style import SrStyleTransferEvaluator -from trainer.eval import StyleTransferEvaluator +from trainer.eval.style import StyleTransferEvaluator def create_evaluator(model, opt_eval, env): diff --git a/codes/trainer/eval/flow_gaussian_nll.py b/codes/trainer/eval/flow_gaussian_nll.py index 7cef4b54..dacf51d1 100644 --- a/codes/trainer/eval/flow_gaussian_nll.py +++ b/codes/trainer/eval/flow_gaussian_nll.py @@ -6,7 +6,7 @@ import trainer.eval.evaluator as evaluator # Evaluate how close to true Gaussian a flow network predicts in a "normal" pass given a LQ/HQ image pair. from data.image_folder_dataset import ImageFolderDataset -from models.archs.srflow_orig.flow import GaussianDiag +from models.srflow.flow import GaussianDiag class FlowGaussianNll(evaluator.Evaluator): diff --git a/codes/trainer/injectors.py b/codes/trainer/injectors.py index d5b5236a..c8248867 100644 --- a/codes/trainer/injectors.py +++ b/codes/trainer/injectors.py @@ -19,7 +19,7 @@ def create_injector(opt_inject, env): from trainer.custom_training_components import create_stereoscopic_injector return create_stereoscopic_injector(opt_inject, env) elif 'igpt' in type: - from models.archs.transformers.igpt import gpt2 + from models.transformers.igpt import gpt2 return gpt2.create_injector(opt_inject, env) elif type == 'generator': return ImageGeneratorInjector(opt_inject, env) @@ -372,7 +372,7 @@ class MultiFrameCombiner(Injector): self.in_hq_key = opt['in_hq'] self.out_lq_key = opt['out'] self.out_hq_key = opt['out_hq'] - from models.archs.flownet2.networks import Resample2d + from models.flownet2.networks import Resample2d self.resampler = Resample2d() def combine(self, state): diff --git a/codes/trainer/losses.py b/codes/trainer/losses.py index 8f47a361..4ef9e437 100644 --- a/codes/trainer/losses.py +++ b/codes/trainer/losses.py @@ -14,7 +14,7 @@ def create_loss(opt_loss, env): from trainer.custom_training_components import create_teco_loss return create_teco_loss(opt_loss, env) elif 'stylegan2_' in type: - from models.archs.stylegan import create_stylegan2_loss + from models.stylegan import create_stylegan2_loss return create_stylegan2_loss(opt_loss, env) elif type == 'crossentropy': return CrossEntropy(opt_loss, env) @@ -311,7 +311,7 @@ class DiscriminatorGanLoss(ConfigurableLoss): if self.gradient_penalty: # Apply gradient penalty. TODO: migrate this elsewhere. - from models.archs.stylegan.stylegan2_lucidrains import gradient_penalty + from models.stylegan.stylegan2_lucidrains import gradient_penalty assert len(real) == 1 # Grad penalty doesn't currently support multi-input discriminators. gp = gradient_penalty(real[0], d_real) self.metrics.append(("gradient_penalty", gp.clone().detach())) diff --git a/codes/trainer/networks.py b/codes/trainer/networks.py index abbde8bc..43c7afe5 100644 --- a/codes/trainer/networks.py +++ b/codes/trainer/networks.py @@ -6,16 +6,16 @@ import munch import torch import torchvision from munch import munchify -import models.archs.stylegan.stylegan2_lucidrains as stylegan2 +import models.stylegan.stylegan2_lucidrains as stylegan2 -import models.archs.fixup_resnet.DiscriminatorResnet_arch as DiscriminatorResnet_arch -import models.archs.RRDBNet_arch as RRDBNet_arch -import models.archs.SwitchedResidualGenerator_arch as SwitchedGen_arch -import models.archs.discriminator_vgg_arch as SRGAN_arch -import models.archs.feature_arch as feature_arch -from models.archs import srg2_classic -from models.archs.stylegan.Discriminator_StyleGAN import StyleGanDiscriminator -from models.archs.tecogan.teco_resgen import TecoGen +import models.fixup_resnet.DiscriminatorResnet_arch as DiscriminatorResnet_arch +import models.RRDBNet_arch as RRDBNet_arch +import models.SwitchedResidualGenerator_arch as SwitchedGen_arch +import models.discriminator_vgg_arch as SRGAN_arch +import models.feature_arch as feature_arch +from models import srg2_classic +from models.stylegan.Discriminator_StyleGAN import StyleGanDiscriminator +from models.tecogan.teco_resgen import TecoGen from utils.util import opt_get logger = logging.getLogger('base') @@ -30,7 +30,7 @@ def define_G(opt, opt_net, scale=None): if which_model == 'RRDBNetBypass': block = RRDBNet_arch.RRDBWithBypass elif which_model == 'RRDBNetLambda': - from models.archs.lambda_rrdb import LambdaRRDB + from models.lambda_rrdb import LambdaRRDB block = LambdaRRDB else: block = RRDBNet_arch.RRDB @@ -62,7 +62,7 @@ def define_G(opt, opt_net, scale=None): heightened_temp_min=opt_net['heightened_temp_min'], heightened_final_step=opt_net['heightened_final_step'], upsample_factor=scale, add_scalable_noise_to_transforms=opt_net['add_noise']) elif which_model == "flownet2": - from models.archs.flownet2 import FlowNet2 + from models.flownet2 import FlowNet2 ld = 'load_path' in opt_net.keys() args = munch.Munch({'fp16': False, 'rgb_max': 1.0, 'checkpoint': not ld}) netG = FlowNet2(args) @@ -85,12 +85,12 @@ def define_G(opt, opt_net, scale=None): netG = stylegan2.StyleGan2GeneratorWithLatent(image_size=opt_net['image_size'], latent_dim=opt_net['latent_dim'], style_depth=opt_net['style_depth'], structure_input=is_structured, attn_layers=attn) - elif which_model == 'srflow_orig': - from models.archs.srflow_orig import SRFlowNet_arch + elif which_model == 'srflow': + from models.srflow import SRFlowNet_arch netG = SRFlowNet_arch.SRFlowNet(in_nc=3, out_nc=3, nf=opt_net['nf'], nb=opt_net['nb'], scale=opt_net['scale'], K=opt_net['K'], opt=opt) elif which_model == 'rrdb_latent_wrapper': - from models.archs.srflow_orig.RRDBNet_arch import RRDBLatentWrapper + from models.srflow.RRDBNet_arch import RRDBLatentWrapper netG = RRDBLatentWrapper(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], with_bypass=opt_net['with_bypass'], blocks=opt_net['blocks_for_latent'], scale=opt_net['scale'], pretrain_rrdb_path=opt_net['pretrain_path']) @@ -100,29 +100,29 @@ def define_G(opt, opt_net, scale=None): mid_channels=opt_net['nf'], num_blocks=opt_net['nb'], scale=opt_net['scale'], headless=True, output_mode=output_mode) elif which_model == 'rrdb_srflow': - from models.archs.srflow_orig.RRDBNet_arch import RRDBNet + from models.srflow.RRDBNet_arch import RRDBNet netG = RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], scale=opt_net['scale'], initial_conv_stride=opt_net['initial_stride']) elif which_model == 'igpt2': - from models.archs.transformers.igpt.gpt2 import iGPT2 + from models.transformers.igpt.gpt2 import iGPT2 netG = iGPT2(opt_net['embed_dim'], opt_net['num_heads'], opt_net['num_layers'], opt_net['num_pixels'] ** 2, opt_net['num_vocab'], centroids_file=opt_net['centroids_file']) elif which_model == 'byol': - from models.archs.byol.byol_model_wrapper import BYOL + from models.byol.byol_model_wrapper import BYOL subnet = define_G(opt, opt_net['subnet']) netG = BYOL(subnet, opt_net['image_size'], opt_net['hidden_layer'], structural_mlp=opt_get(opt_net, ['use_structural_mlp'], False)) elif which_model == 'structural_byol': - from models.archs.byol.byol_structural import StructuralBYOL + from models.byol.byol_structural import StructuralBYOL subnet = define_G(opt, opt_net['subnet']) netG = StructuralBYOL(subnet, opt_net['image_size'], opt_net['hidden_layer'], pretrained_state_dict=opt_get(opt_net, ["pretrained_path"]), freeze_until=opt_get(opt_net, ['freeze_until'], 0)) elif which_model == 'spinenet': - from models.archs.spinenet_arch import SpineNet + from models.spinenet_arch import SpineNet netG = SpineNet(str(opt_net['arch']), in_channels=3, use_input_norm=opt_net['use_input_norm']) elif which_model == 'spinenet_with_logits': - from models.archs.spinenet_arch import SpinenetWithLogits + from models.spinenet_arch import SpinenetWithLogits netG = SpinenetWithLogits(str(opt_net['arch']), opt_net['output_to_attach'], opt_net['num_labels'], in_channels=3, use_input_norm=opt_net['use_input_norm']) else: diff --git a/codes/utils/numeric_stability.py b/codes/utils/numeric_stability.py index 2fc0fb5f..63cc763f 100644 --- a/codes/utils/numeric_stability.py +++ b/codes/utils/numeric_stability.py @@ -1,7 +1,7 @@ import torch from torch import nn -import models.archs.SwitchedResidualGenerator_arch as srg -import models.archs.discriminator_vgg_arch as disc +import models.SwitchedResidualGenerator_arch as srg +import models.discriminator_vgg_arch as disc import functools blacklisted_modules = [nn.Conv2d, nn.ReLU, nn.LeakyReLU, nn.BatchNorm2d, nn.Softmax]