def main(args):
lr_train_patch_size = args.patch_size
layers_to_extract = [5, 9]
hr_train_patch_size = lr_train_patch_size * args.scale
if args.model == 'rdn':
model = RDN(arch_params={'C':4, 'D':3, 'G':64, 'G0':64, 'T':10, 'x':args.scale}, patch_size=lr_train_patch_size)
else:
model = RRDN(arch_params={'C':4, 'D':3, 'G':64, 'G0':64, 'T':10, 'x':args.scale}, patch_size=lr_train_patch_size)
f_ext = Cut_VGG19(patch_size=hr_train_patch_size, layers_to_extract=layers_to_extract)
discr = Discriminator(patch_size=hr_train_patch_size, kernel_size=3)
loss_weights = {
'generator': 0.0,
'feature_extractor': 0.0833,
'discriminator': 0.01,
}
trainer = Trainer(
generator=model,
discriminator=discr,
feature_extractor=f_ext,
lr_train_dir='low_res/training/images',
hr_train_dir='high_res/training/images',
lr_valid_dir='low_res/validation/images',
hr_valid_dir='high_res/validation/images',
loss_weights=loss_weights,
dataname=args.name,
logs_dir='./logs',
weights_dir='./weights',
weights_generator=None,
weights_discriminator=None,
n_validation=40,
lr_decay_frequency=30,
lr_decay_factor=0.5,
)
trainer.train(epochs=args.num_epochs,
steps_per_epoch=args.epoch_steps,
batch_size,args.batch_size)
lr_train_patch_size = 40
layers_to_extract = [5, 9]
scale = 8
hr_train_patch_size = lr_train_patch_size * scale
rrdn = RRDN(arch_params={
'C': 4,
'D': 3,
'G': 64,
'G0': 64,
'T': 10,
'x': scale
},
patch_size=lr_train_patch_size)
f_ext = Cut_VGG19(patch_size=hr_train_patch_size,
layers_to_extract=layers_to_extract)
discr = Discriminator(patch_size=hr_train_patch_size, kernel_size=3)
###################################################################################################################
#run = wandb.init(project='superres')
run = wandb.init(project='respick')
config = run.config
config.num_epochs = 50
config.batch_size = 32
config.input_height = 32
config.input_width = 32
config.output_height = 256
config.output_width = 256
from ISR.models import RRDN_ADV
from ISR.models import Discriminator
from ISR.models import Cut_VGG19
lr_train_patch_size = 64
layers_to_extract = [5, 9]
scale = 8 #[2,4,8,16.....]
hr_train_patch_size = lr_train_patch_size * scale
rrdn = RRDN_ADV(arch_params={'C':4, 'D':3, 'G':64, 'G0':64, 'T':10, 'x':scale}, patch_size=lr_train_patch_size)
f_ext = Cut_VGG19(patch_size=None, layers_to_extract=layers_to_extract)
discr = Discriminator(patch_size=hr_train_patch_size, kernel_size=3)
from ISR.train import Trainer_adv
loss_weights = {
'generator': 0.06,
'feature_extractor': 0.083,
'discriminator': 0.04
}
losses = {
'generator': 'mae',
'feature_extractor': 'mse',
'discriminator': 'binary_crossentropy'
}
log_dirs = {'logs': './logs', 'weights': './weights'}
learning_rate = {'initial_value': 0.0004, 'decay_factor': 0.5, 'decay_frequency': 30}
flatness = {'min': 0.0, 'max': 0.15, 'increase': 0.01, 'increase_frequency': 5}#5