deviandice/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Issues Pull Requests Packages Projects Releases Wiki Activity
78276afcaa
DL-Art-School/codes/models/networks.py

145 lines
11 KiB
Python
Raw Normal View History

mmsr
2019-08-23 13:42:47 +00:00
import torch
import models.archs.SRResNet_arch as SRResNet_arch
import models.archs.discriminator_vgg_arch as SRGAN_arch
Discriminator part 1 New discriminator. Includes spectral norming.
2020-04-29 05:00:29 +00:00
import models.archs.DiscriminatorResnet_arch as DiscriminatorResnet_arch
Add support for passthrough disc/gen Add RRDBNetXL, which performs processing at multiple image sizes. Add DiscResnet_passthrough, which allows passthrough of image at different sizes for discrimination. Adjust the rest of the repo to allow generators that return more than just a single image.
2020-05-04 20:01:43 +00:00
import models.archs.DiscriminatorResnet_arch_passthrough as DiscriminatorResnet_arch_passthrough
Full resnet corrupt, no BN And it works! Thanks fixup..
2020-05-01 01:17:30 +00:00
import models.archs.FlatProcessorNetNew_arch as FlatProcessorNetNew_arch
mmsr
2019-08-23 13:42:47 +00:00
import models.archs.RRDBNet_arch as RRDBNet_arch
Implement downsample GAN This bad boy is for a workflow where you train a model on disjoint image sets to downsample a "good" set of images like a "bad" set of images looks. You then use that downsampler to generate a training set of paired images for supersampling.
2020-04-24 06:00:46 +00:00
import models.archs.HighToLowResNet as HighToLowResNet
Add NestedSwitchGenerator An evolution of SwitchedResidualGenerator, this variant nests attention modules upon themselves to extend the representative capacity of the model significantly.
2020-06-29 03:21:57 +00:00
import models.archs.NestedSwitchGenerator as ng
Separate feature extractors out, add resnet feature extractor
2020-05-29 02:26:30 +00:00
import models.archs.feature_arch as feature_arch
New arch: SwitchedResidualGenerator_arch The concept here is to use switching to split the generator into two functions: interpretation and transformation. Transformation is done at the pixel level by relatively simple conv layers, while interpretation is computed at various levels by far more complicated conv stacks. The two are merged using the switching mechanism. This architecture is far less computationally intensive that RRDB.
2020-06-16 17:23:50 +00:00
import models.archs.SwitchedResidualGenerator_arch as SwitchedGen_arch
Fix attention conv bugs
2020-06-07 00:29:25 +00:00
import functools
mmsr
2019-08-23 13:42:47 +00:00
# Generator
Allow "corruptor" network to be specified This network is just a fixed (pre-trained) generator that performs a corruption transformation that the generator-in-training is expected to undo alongside SR.
2020-05-13 21:26:55 +00:00
def define_G(opt, net_key='network_G'):
opt_net = opt[net_key]
mmsr
2019-08-23 13:42:47 +00:00
which_model = opt_net['which_model_G']
Change GT_size to target_size
2020-04-22 06:37:41 +00:00
scale = opt['scale']
mmsr
2019-08-23 13:42:47 +00:00
# image restoration
if which_model == 'MSRResNet':
netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale'])
elif which_model == 'RRDBNet':
Change GT_size to target_size
2020-04-22 06:37:41 +00:00
# RRDB does scaling in two steps, so take the sqrt of the scale we actually want to achieve and feed it to RRDB.
Add RRDB Initial Stride Allows downsampling immediately before processing, which reduces network complexity on higher resolution images but keeps a higher filter count.
2020-06-02 16:47:15 +00:00
initial_stride = 1 if 'initial_stride' not in opt_net else opt_net['initial_stride']
assert initial_stride == 1 or initial_stride == 2
# Need to adjust the scale the generator sees by the stride since the stride causes a down-sample.
gen_scale = scale * initial_stride
mmsr
2019-08-23 13:42:47 +00:00
netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
Fix rrdb bug
2020-06-02 17:15:55 +00:00
nf=opt_net['nf'], nb=opt_net['nb'], scale=gen_scale, initial_stride=initial_stride)
Add AssistedRRDB and remove RRDBNetXL
2020-05-24 03:09:21 +00:00
elif which_model == 'AssistedRRDBNet':
netG = RRDBNet_arch.AssistedRRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
nf=opt_net['nf'], nb=opt_net['nb'], scale=scale)
Fix attention conv bugs
2020-06-07 00:29:25 +00:00
elif which_model == 'AttentiveRRDBNet':
netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
nf=opt_net['nf'], nb=opt_net['nb'], scale=scale,
Add alternative first block for PixShuffleRRDB
2020-06-11 03:45:24 +00:00
rrdb_block_f=functools.partial(RRDBNet_arch.SwitchedRRDB, nf=opt_net['nf'], gc=opt_net['gc'],
Allow switched RRDBNet to record metrics and decay temperature
2020-06-08 17:10:38 +00:00
init_temperature=opt_net['temperature'],
final_temperature_step=opt_net['temperature_final_step']))
Multiple modifications for experimental RRDB architectures - Add LowDimRRDB; essentially a "normal RRDB" but the RDB blocks process at a low dimension using PixelShuffle - Add switching wrappers around it - Add support for switching on top of multi-headed inputs and outputs - Moves PixelUnshuffle to arch_util
2020-06-13 17:37:27 +00:00
elif which_model == 'LowDimRRDBNet':
Fix initial_stride on lowdim models
2020-06-14 17:02:16 +00:00
gen_scale = scale * opt_net['initial_stride']
Multiple modifications for experimental RRDB architectures - Add LowDimRRDB; essentially a "normal RRDB" but the RDB blocks process at a low dimension using PixelShuffle - Add switching wrappers around it - Add support for switching on top of multi-headed inputs and outputs - Moves PixelUnshuffle to arch_util
2020-06-13 17:37:27 +00:00
rrdb = functools.partial(RRDBNet_arch.LowDimRRDB, nf=opt_net['nf'], gc=opt_net['gc'], dimensional_adjustment=opt_net['dim'])
netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
Fix initial_stride on lowdim models
2020-06-14 17:02:16 +00:00
nf=opt_net['nf'], nb=opt_net['nb'], scale=gen_scale, rrdb_block_f=rrdb, initial_stride=opt_net['initial_stride'])
Multiple modifications for experimental RRDB architectures - Add LowDimRRDB; essentially a "normal RRDB" but the RDB blocks process at a low dimension using PixelShuffle - Add switching wrappers around it - Add support for switching on top of multi-headed inputs and outputs - Moves PixelUnshuffle to arch_util
2020-06-13 17:37:27 +00:00
elif which_model == "LowDimRRDBWithMultiHeadSwitching":
Fix initial_stride on lowdim models
2020-06-14 17:02:16 +00:00
gen_scale = scale * opt_net['initial_stride']
Multiple modifications for experimental RRDB architectures - Add LowDimRRDB; essentially a "normal RRDB" but the RDB blocks process at a low dimension using PixelShuffle - Add switching wrappers around it - Add support for switching on top of multi-headed inputs and outputs - Moves PixelUnshuffle to arch_util
2020-06-13 17:37:27 +00:00
switcher = functools.partial(RRDBNet_arch.SwitchedMultiHeadRRDB, num_convs=opt_net['num_convs'], num_heads=opt_net['num_heads'],
init_temperature=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'])
rrdb = functools.partial(RRDBNet_arch.LowDimRRDBWrapper, nf=opt_net['nf'], gc=opt_net['gc'], dimensional_adjustment=opt_net['dim'],
partial_rrdb=switcher)
netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
Fix initial_stride on lowdim models
2020-06-14 17:02:16 +00:00
nf=opt_net['nf'], nb=opt_net['nb'], scale=gen_scale, rrdb_block_f=rrdb, initial_stride=opt_net['initial_stride'])
Add serveral new RRDB architectures
2020-06-09 19:28:55 +00:00
elif which_model == 'PixRRDBNet':
block_f = None
if opt_net['attention']:
Add alternative first block for PixShuffleRRDB
2020-06-11 03:45:24 +00:00
block_f = functools.partial(RRDBNet_arch.SwitchedRRDB, nf=opt_net['nf'], gc=opt_net['gc'],
init_temperature=opt_net['temperature'],
Add skip heads to switcher These pass through the input so that it can be selected by the attention mechanism.
2020-06-14 18:46:54 +00:00
final_temperature_step=opt_net['temperature_final_step'])
if opt_net['mhattention']:
block_f = functools.partial(RRDBNet_arch.SwitchedMultiHeadRRDB, num_convs=8, num_heads=2, nf=opt_net['nf'], gc=opt_net['gc'],
init_temperature=opt_net['temperature'],
Add alternative first block for PixShuffleRRDB
2020-06-11 03:45:24 +00:00
final_temperature_step=opt_net['temperature_final_step'])
Add serveral new RRDB architectures
2020-06-09 19:28:55 +00:00
netG = RRDBNet_arch.PixShuffleRRDB(nf=opt_net['nf'], nb=opt_net['nb'], gc=opt_net['gc'], scale=scale, rrdb_block_f=block_f)
Create a configurable SwichedResidualGenerator Also move attention image generator out of repo
2020-06-16 19:24:07 +00:00
elif which_model == "ConfigurableSwitchedResidualGenerator":
Add growth channel to switch_growths for flat networks
2020-06-22 16:40:16 +00:00
netG = SwitchedGen_arch.ConfigurableSwitchedResidualGenerator(switch_filters=opt_net['switch_filters'], switch_growths=opt_net['switch_growths'],
switch_reductions=opt_net['switch_reductions'],
Mods to SwitchedResidualGenerator_arch - Increased processing for high-resolution switches - Do stride=2 first in HalvingProcessingBlock
2020-06-16 20:19:12 +00:00
switch_processing_layers=opt_net['switch_processing_layers'], trans_counts=opt_net['trans_counts'],
Create a configurable SwichedResidualGenerator Also move attention image generator out of repo
2020-06-16 19:24:07 +00:00
trans_kernel_sizes=opt_net['trans_kernel_sizes'], trans_layers=opt_net['trans_layers'],
ResidualGen mods - Add filters_mid spec which allows a expansion->squeeze for the transformation layers. - Add scale and bias AFTER the switch - Remove identity transform (models were converging on this) - Move attention image generation and temperature setting into new function which gets called every step with a save path
2020-06-17 23:18:28 +00:00
trans_filters_mid=opt_net['trans_filters_mid'],
Add a double-step to attention temperature
2020-06-18 17:29:31 +00:00
initial_temp=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'],
Add ConfigurableSwitchComputer
2020-06-25 01:49:37 +00:00
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 == "ConfigurableSwitchedResidualGenerator2":
netG = SwitchedGen_arch.ConfigurableSwitchedResidualGenerator2(switch_filters=opt_net['switch_filters'], switch_growths=opt_net['switch_growths'],
switch_reductions=opt_net['switch_reductions'],
switch_processing_layers=opt_net['switch_processing_layers'], trans_counts=opt_net['trans_counts'],
trans_kernel_sizes=opt_net['trans_kernel_sizes'], trans_layers=opt_net['trans_layers'],
Missed networks.py in last commit
2020-06-26 00:36:06 +00:00
transformation_filters=opt_net['transformation_filters'],
Add ConfigurableSwitchComputer
2020-06-25 01:49:37 +00:00
initial_temp=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'],
Add NestedSwitchGenerator An evolution of SwitchedResidualGenerator, this variant nests attention modules upon themselves to extend the representative capacity of the model significantly.
2020-06-29 03:21:57 +00:00
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 == "NestedSwitchGenerator":
netG = ng.NestedSwitchedGenerator(switch_filters=opt_net['switch_filters'],
switch_reductions=opt_net['switch_reductions'],
switch_processing_layers=opt_net['switch_processing_layers'], trans_counts=opt_net['trans_counts'],
trans_kernel_sizes=opt_net['trans_kernel_sizes'], trans_layers=opt_net['trans_layers'],
transformation_filters=opt_net['transformation_filters'],
initial_temp=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'],
Fix inverse temperature curve logic and add upsample factor
2020-06-19 15:18:30 +00:00
heightened_temp_min=opt_net['heightened_temp_min'], heightened_final_step=opt_net['heightened_final_step'],
Add parameterized noise injection into resgen
2020-06-23 16:16:02 +00:00
upsample_factor=scale, add_scalable_noise_to_transforms=opt_net['add_noise'])
Implement ResGen arch This is a simpler resnet-based generator which performs mutations on an input interspersed with interpolate-upsampling. It is a two part generator: 1) A component that "fixes" LQ images with a long string of resnet blocks. This component is intended to remove compression artifacts and other noise from a LQ image. 2) A component that can double the image size. The idea is that this component be trained so that it can work at most reasonable resolutions, such that it can be repeatedly applied to itself to perform multiple upsamples. The motivation here is to simplify what is being done inside of RRDB. I don't believe the complexity inside of that network is justified.
2020-05-05 17:59:46 +00:00
Implement downsample GAN This bad boy is for a workflow where you train a model on disjoint image sets to downsample a "good" set of images like a "bad" set of images looks. You then use that downsampler to generate a training set of paired images for supersampling.
2020-04-24 06:00:46 +00:00
# image corruption
elif which_model == 'HighToLowResNet':
netG = HighToLowResNet.HighToLowResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
nf=opt_net['nf'], nb=opt_net['nb'], downscale=opt_net['scale'])
Add FlatProcessorNet After doing some thinking and reading on the subject, it occurred to me that I was treating the generator like a discriminator by focusing the network complexity at the feature levels. It makes far more sense to process each conv level equally for the generator, hence the FlatProcessorNet in this commit. This network borrows some of the residual pass-through logic from RRDB which makes the gradient path exceptionally short for pretty much all model parameters and can be trained in O1 optimization mode without overflows again.
2020-04-28 17:48:05 +00:00
elif which_model == 'FlatProcessorNet':
Full resnet corrupt, no BN And it works! Thanks fixup..
2020-05-01 01:17:30 +00:00
'''netG = FlatProcessorNet_arch.FlatProcessorNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
Add FlatProcessorNet After doing some thinking and reading on the subject, it occurred to me that I was treating the generator like a discriminator by focusing the network complexity at the feature levels. It makes far more sense to process each conv level equally for the generator, hence the FlatProcessorNet in this commit. This network borrows some of the residual pass-through logic from RRDB which makes the gradient path exceptionally short for pretty much all model parameters and can be trained in O1 optimization mode without overflows again.
2020-04-28 17:48:05 +00:00
nf=opt_net['nf'], downscale=opt_net['scale'], reduce_anneal_blocks=opt_net['ra_blocks'],
Full resnet corrupt, no BN And it works! Thanks fixup..
2020-05-01 01:17:30 +00:00
assembler_blocks=opt_net['assembler_blocks'])'''
Fix initialization in mhead switched rrdb
2020-06-16 03:32:03 +00:00
netG = FlatProcessorNetNew_arch.fixup_resnet34(num_filters=opt_net['nf'])\
Change GT_size to target_size
2020-04-22 06:37:41 +00:00
mmsr
2019-08-23 13:42:47 +00:00
else:
raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))
return netG
# Discriminator
def define_D(opt):
Change GT_size to target_size
2020-04-22 06:37:41 +00:00
img_sz = opt['datasets']['train']['target_size']
mmsr
2019-08-23 13:42:47 +00:00
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
if which_model == 'discriminator_vgg_128':
Add capability to place additional conv into discriminator This should allow us to support larger images sizes. May need to add another one of these.
2020-06-23 15:40:33 +00:00
netD = SRGAN_arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], nf=opt_net['nf'], input_img_factor=img_sz // 128, extra_conv=opt_net['extra_conv'])
Discriminator part 1 New discriminator. Includes spectral norming.
2020-04-29 05:00:29 +00:00
elif which_model == 'discriminator_resnet':
Fixup upconv for the next attempt!
2020-05-02 01:56:14 +00:00
netD = DiscriminatorResnet_arch.fixup_resnet34(num_filters=opt_net['nf'], num_classes=1, input_img_size=img_sz)
Add support for passthrough disc/gen Add RRDBNetXL, which performs processing at multiple image sizes. Add DiscResnet_passthrough, which allows passthrough of image at different sizes for discrimination. Adjust the rest of the repo to allow generators that return more than just a single image.
2020-05-04 20:01:43 +00:00
elif which_model == 'discriminator_resnet_passthrough':
Fix skips & images samples - Makes skip connections between the generator and discriminator more extensible by adding additional configuration options for them and supporting 1 and 0 skips. - Places the temp/ directory with sample images from the training process appear in the training directory instead of the codes/ directory.
2020-05-15 19:50:49 +00:00
netD = DiscriminatorResnet_arch_passthrough.fixup_resnet34(num_filters=opt_net['nf'], num_classes=1, input_img_size=img_sz,
Allow passthrough discriminator to have passthrough disabled from config
2020-05-19 15:41:16 +00:00
number_skips=opt_net['number_skips'], use_bn=True,
disable_passthrough=opt_net['disable_passthrough'])
mmsr
2019-08-23 13:42:47 +00:00
else:
raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
return netD
# Define network used for perceptual loss
def define_F(opt, use_bn=False):
gpu_ids = opt['gpu_ids']
device = torch.device('cuda' if gpu_ids else 'cpu')
Separate feature extractors out, add resnet feature extractor
2020-05-29 02:26:30 +00:00
if 'which_model_F' not in opt['train'].keys() or opt['train']['which_model_F'] == 'vgg':
# PyTorch pretrained VGG19-54, before ReLU.
if use_bn:
feature_layer = 49
else:
feature_layer = 34
netF = feature_arch.VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn,
use_input_norm=True, device=device)
elif opt['train']['which_model_F'] == 'wide_resnet':
netF = feature_arch.WideResnetFeatureExtractor(use_input_norm=True, device=device)
mmsr
2019-08-23 13:42:47 +00:00
netF.eval() # No need to train
return netF
Reference in New Issue Copy Permalink