Parameter setting for training the RDN ArtefactCancelling models?

[JunbinWang]

Dear cfrancesco,

I am very interested about how you train the ArtefactCancelling RDN model as you provided in the 'sample_weights/' dir.

I tried to try the model, but didn't get a very good result.

Below is my training scripts.

from ISR.models import RDN
from ISR.models import Discriminator
from ISR.models import Cut_VGG19


lr_train_patch_size = 40
layers_to_extract = [5, 9]
scale = 2
hr_train_patch_size = lr_train_patch_size * scale

rdn  = RDN(arch_params={'C':3, 'D':10, 'G':64, 'G0':64, '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)

from ISR.train import Trainer
loss_weights = {
  'generator': 0.0,
  'feature_extractor': 0.0833,
  'discriminator': 0.01
}
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}

trainer = Trainer(
    generator=rdn,
    discriminator=discr,
    feature_extractor=f_ext,
    lr_train_dir= '/data1/wjb/div2k/DIV2K_train_LR_bicubic/X2/',  #'div2k/DIV2K_train_LR_bicubic/X2/'
    hr_train_dir= '/data1/wjb/div2k/DIV2K_train_HR/',  #'div2k/DIV2K_train_HR/'
    lr_valid_dir= '/data1/wjb/div2k/DIV2K_valid_LR_bicubic/X2', #'div2k/DIV2K_train_LR_bicubic/X2/',
    hr_valid_dir='/data1/wjb/div2k/DIV2K_valid_HR/', #'div2k/DIV2K_train_HR/',
    loss_weights=loss_weights,
    learning_rate=learning_rate,
    flatness=flatness,
    dataname='div2k',
    log_dirs=log_dirs,
    weights_generator=None,
    weights_discriminator=None,
    n_validation=100,
)

trainer.train(
    epochs=300,
    steps_per_epoch=100,
    batch_size=16,
    #monitored_metrics={'val_generator_PSNR': 'max'}
    monitored_metrics={'val_generator_PSNR_Y': 'max'}
)

Best Regard, Even

[cfrancesco]

Hi,

the first few epochs (~20) you might want to train the model only with the MAE instead than adversarial loss (maybe without adding noise?). There is not automatic switch yet though, so you have to train and then train again loading the pre-trained weights you obtained in the first part

[JunbinWang]

Hi,

the first few epochs (~20) you might want to train the model only with the MAE instead than adversarial loss (maybe without adding noise?). There is not automatic switch yet though, so you have to train and then train again loading the pre-trained weights you obtained in the first part

HI @cfrancesco,

Thank you very much for your reply!

I will try it!

By the way, do you think it's possible to train the RDN ArtefactCancelling models with smaller models? like the one I used above (rdn = RDN(arch_params={'C':3, 'D':10, 'G':64, 'G0':64, 'x':scale})

And do you think my configurations to train the ArtefactCancelling model ( after loading the pre-trained weights ) is correct above?

[cfrancesco]

If you have an adequate dataset, I would use a 4x upscaling. For the rest, your configuration looks fine. In my experiment I observed no to very little qualitative difference between the larger and smaller models. I would however recommend to use the RRDN model, with a similar configuration to what we provide already: you could for instance retrain the provided RRDN GANS model, with proper preprocessing have it learn noise (et al.) cancelling and share here your results ;)

Also, experiment with the layers you extract from the vgg network.

[JunbinWang]

@cfrancesco

Hi cfrancesco,

My idea is to train a model that can resize images to any size. To achieve this,
during the training phase, I would like to train an SR model with 1x scale factor; during the inference phase, I would resize the image to the desired size first, and then use the SR model to enhance it. I am not sure if this idea is practical?

Second, the reason I would like to use a smaller model is that I found the original model is very slow during training and inference time even with the support of GPU resources.

It took me about 10 minutes to process 100 images, I don't know why. By the way, I found the model doesn't use too much GPU resources during the training process. I have a GPU with 32GB memory, but it looks like only about 300MB GPU memory are used during the training phase. Do you know how to make the model runs faster?

Best Regard, Even.

[cfrancesco]

Nice idea, but that kinda means to go back in time. Image restoration used to be the pre-2015(16?) approach, which according to the literature delivers worse results. Nevertheless, give it a go, maybe this suits your needs. However, I would probably upscale to a high scale and then shrink back, given that the computationally expensive part is the blocks part, not the actual upscaling. Yes, these models are slow due to the high number of operations needed (the model does not shrink through the layers). You could try to parallelize the inference of different images to speed up, if your gpu allows for it

[JunbinWang]

haha @cfrancesco , you are right !!! Upscale image to a larger size and then downscaling it would be a better idea!

But the speed of the model could still be a tricky problem~! I don't know how you solve it~ Maybe I should try to explore the light-weight model ！