add SRFlow with srflow.py

models/SRFlow/35000_G.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:040fcffde66ec3ef658a843d58832b8aa153734a2f04342d841e3018b498a511
+size 158819348

models/SRFlow/code/Measure.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import time
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import cv2
+import argparse
+
+from natsort import natsort
+from skimage.metrics import structural_similarity as ssim
+from skimage.metrics import peak_signal_noise_ratio as psnr
+import lpips
+
+
+class Measure():
+    def __init__(self, net='alex', use_gpu=False):
+        self.device = 'cuda' if use_gpu else 'cpu'
+        self.model = lpips.LPIPS(net=net)
+        self.model.to(self.device)
+
+    def measure(self, imgA, imgB):
+        return [float(f(imgA, imgB)) for f in [self.psnr, self.ssim, self.lpips]]
+
+    def lpips(self, imgA, imgB, model=None):
+        tA = t(imgA).to(self.device)
+        tB = t(imgB).to(self.device)
+        dist01 = self.model.forward(tA, tB).item()
+        return dist01
+
+    def ssim(self, imgA, imgB):
+        # multichannel: If True, treat the last dimension of the array as channels. Similarity calculations are done independently for each channel then averaged.
+        score, diff = ssim(imgA, imgB, full=True, multichannel=True, channel_axis=-1)
+        return score
+
+    def psnr(self, imgA, imgB):
+        psnr_val = psnr(imgA, imgB)
+        return psnr_val
+
+
+def t(img):
+    def to_4d(img):
+        assert len(img.shape) == 3
+        assert img.dtype == np.uint8
+        img_new = np.expand_dims(img, axis=0)
+        assert len(img_new.shape) == 4
+        return img_new
+
+    def to_CHW(img):
+        return np.transpose(img, [2, 0, 1])
+
+    def to_tensor(img):
+        return torch.Tensor(img)
+
+    return to_tensor(to_4d(to_CHW(img))) / 127.5 - 1
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+def imread(path):
+    return cv2.imread(path)[:, :, [2, 1, 0]]
+
+
+def format_result(psnr, ssim, lpips):
+    return f'{psnr:0.2f}, {ssim:0.3f}, {lpips:0.3f}'
+
+def measure_dirs(dirA, dirB, use_gpu, verbose=False):
+    if verbose:
+        vprint = lambda x: print(x)
+    else:
+        vprint = lambda x: None
+
+
+    t_init = time.time()
+
+    paths_A = fiFindByWildcard(os.path.join(dirA, f'*.{type}'))
+    paths_B = fiFindByWildcard(os.path.join(dirB, f'*.{type}'))
+
+    vprint("Comparing: ")
+    vprint(dirA)
+    vprint(dirB)
+
+    measure = Measure(use_gpu=use_gpu)
+
+    results = []
+    for pathA, pathB in zip(paths_A, paths_B):
+        result = OrderedDict()
+
+        t = time.time()
+        result['psnr'], result['ssim'], result['lpips'] = measure.measure(imread(pathA), imread(pathB))
+        d = time.time() - t
+        vprint(f"{pathA.split('/')[-1]}, {pathB.split('/')[-1]}, {format_result(**result)}, {d:0.1f}")
+
+        results.append(result)
+
+    psnr = np.mean([result['psnr'] for result in results])
+    ssim = np.mean([result['ssim'] for result in results])
+    lpips = np.mean([result['lpips'] for result in results])
+
+    vprint(f"Final Result: {format_result(psnr, ssim, lpips)}, {time.time() - t_init:0.1f}s")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-dirA', default='', type=str)
+    parser.add_argument('-dirB', default='', type=str)
+    parser.add_argument('-type', default='png')
+    parser.add_argument('--use_gpu', action='store_true', default=False)
+    args = parser.parse_args()
+
+    dirA = args.dirA
+    dirB = args.dirB
+    type = args.type
+    use_gpu = args.use_gpu
+
+    if len(dirA) > 0 and len(dirB) > 0:
+        measure_dirs(dirA, dirB, use_gpu=use_gpu, verbose=True)

models/SRFlow/code/a.py

@@ -0,0 +1,27 @@
+import pickle
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+
+def load_pkls(path):
+    assert os.path.isfile(path), path
+    images = []
+    with open(path, "rb") as f:
+        images += pickle.load(f)
+    assert len(images) > 0, path
+    images = [np.transpose(image, [2, 0, 1]) for image in images]
+    return images
+
+path = 'datasets/DIV2K-va.pklv4'
+loaded_images = load_pkls(path)
+print(len(loaded_images))
+# Display the first image
+if loaded_images:
+    first_image = loaded_images[11]
+    plt.imshow(np.transpose(first_image, [1, 2, 0]))  # Transpose image to original shape [height, width, channels]
+    plt.title('First Image')
+    plt.axis('off')  # Hide axis
+    plt.show()
+else:
+    print("No images loaded from the pickle file.")
+print(loaded_images[11])

models/SRFlow/code/confs/RRDB_CelebA_8X.yml

@@ -0,0 +1,83 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 8
+#gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8_1pct.pklv4
+
+    use_shuffle: true
+    n_workers: 0  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    use_rot: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/celebA-valid-x8_1pct.pklv4
+    n_max: 10
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 200000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/RRDB_DF2K_4X.yml

@@ -0,0 +1,85 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 4
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-train-x4_1pct.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-valid-x4_1pct.pklv4
+    quant: 32
+    n_max: 20
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  use_orig: True
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 1000000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/RRDB_DF2K_8X.yml

@@ -0,0 +1,85 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-train-x8_1pct.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-valid-x8_1pct.pklv4
+    quant: 32
+    n_max: 20
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 200000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/SRFlow_CelebA_8X.yml

@@ -0,0 +1,107 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot_GT: ../datasets/celebA-validation-gt
+dataroot_LR: ../datasets/celebA-validation-x8
+model_path: ../pretrained_models/SRFlow_CelebA_8X.pth
+heat: 0.9 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 8
+  upscale: 8
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 4
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 3, 5, 7 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: ../pretrained_models/RRDB_CelebA_8X.pth
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 200000
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/SRFlow_DF2K_4X.yml

@@ -0,0 +1,106 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 4
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: DF2K_256_tr
+    mode: LRHR_PKL
+    dataroot_GT: /kaggle/input/srflow0103/SRFlow/datasets/DF2K-tr.pklv4
+    dataroot_LQ: /kaggle/input/srflow0103/SRFlow/datasets/DF2K-tr_X4.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 12
+    GT_size: 256
+    use_flip: true
+    color: RGB
+  val:
+    name: DF2K_256_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DIV2K-va.pklv4
+    dataroot_LQ: ../datasets/DIV2K-va_X4.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot: /kaggle/input/test-set/test set
+model_path: /models/SRFlow/35000_G
+heat: 0.6 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+  upscale: 4
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 3
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 8, 15, 22 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: 
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 2.5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 64185
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3

models/SRFlow/code/confs/SRFlow_DF2K_8X.yml

@@ -0,0 +1,112 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-tr.pklv4
+    dataroot_LQ: ../datasets/DF2K-tr_X8.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DIV2K-va.pklv4
+    dataroot_LQ: ../datasets/DIV2K-va_X8.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot_GT: ../datasets/div2k-validation-modcrop8-gt
+dataroot_LR: ../datasets/div2k-validation-modcrop8-x8
+model_path: ../pretrained_models/SRFlow_DF2K_8X.pth
+heat: 0.9 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+  upscale: 8
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 4
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 3, 5, 7 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: ../pretrained_models/RRDB_DF2K_8X.pth
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 200000
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+test:
+  heats: [ 0.0, 0.7, 0.8, 0.9 ]
+
+#### logger
+logger:
+  # Debug print_freq: 100
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/data/LRHR_PKL_dataset.py

@@ -0,0 +1,179 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+import subprocess
+import torch.utils.data as data
+import numpy as np
+import time
+import torch
+
+import pickle
+
+
+class LRHR_PKLDataset(data.Dataset):
+    def __init__(self, opt):
+        super(LRHR_PKLDataset, self).__init__()
+        self.opt = opt
+        self.crop_size = opt.get("GT_size", None)
+        self.scale = None
+        self.random_scale_list = [1]
+
+        hr_file_path = opt["dataroot_GT"]
+        lr_file_path = opt["dataroot_LQ"]
+        y_labels_file_path = opt['dataroot_y_labels']
+
+        gpu = True
+        augment = True
+
+        self.use_flip = opt["use_flip"] if "use_flip" in opt.keys() else False
+        self.use_rot = opt["use_rot"] if "use_rot" in opt.keys() else False
+        self.use_crop = opt["use_crop"] if "use_crop" in opt.keys() else False
+        self.center_crop_hr_size = opt.get("center_crop_hr_size", None)
+
+        n_max = opt["n_max"] if "n_max" in opt.keys() else int(1e8)
+
+        t = time.time()
+        self.lr_images = self.load_pkls(lr_file_path, n_max)
+        self.hr_images = self.load_pkls(hr_file_path, n_max)
+
+        min_val_hr = np.min([i.min() for i in self.hr_images[:20]])
+        max_val_hr = np.max([i.max() for i in self.hr_images[:20]])
+
+        min_val_lr = np.min([i.min() for i in self.lr_images[:20]])
+        max_val_lr = np.max([i.max() for i in self.lr_images[:20]])
+
+        t = time.time() - t
+        print("Loaded {} HR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+              format(len(self.hr_images), min_val_hr, max_val_hr, t, hr_file_path))
+        print("Loaded {} LR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+              format(len(self.lr_images), min_val_lr, max_val_lr, t, lr_file_path))
+
+        self.gpu = gpu
+        self.augment = augment
+
+        self.measures = None
+
+    def load_pkls(self, path, n_max):
+        assert os.path.isfile(path), path
+        images = []
+        with open(path, "rb") as f:
+            images += pickle.load(f)
+        assert len(images) > 0, path
+        images = images[:n_max]
+        images = [np.transpose(image, [2, 0, 1]) for image in images]
+        return images
+
+    def __len__(self):
+        return len(self.hr_images)
+
+    def __getitem__(self, item):
+        hr = self.hr_images[item]
+        lr = self.lr_images[item]
+
+        if self.scale == None:
+            self.scale = hr.shape[1] // lr.shape[1]
+            assert hr.shape[1] == self.scale * lr.shape[1], ('non-fractional ratio', lr.shape, hr.shape)
+
+        if self.use_crop:
+            hr, lr = random_crop(hr, lr, self.crop_size, self.scale, self.use_crop)
+
+        if self.center_crop_hr_size:
+            hr, lr = center_crop(hr, self.center_crop_hr_size), center_crop(lr, self.center_crop_hr_size // self.scale)
+
+        if self.use_flip:
+            hr, lr = random_flip(hr, lr)
+
+        if self.use_rot:
+            hr, lr = random_rotation(hr, lr)
+
+        hr = hr / 255.0
+        lr = lr / 255.0
+
+        if self.measures is None or np.random.random() < 0.05:
+            if self.measures is None:
+                self.measures = {}
+            self.measures['hr_means'] = np.mean(hr)
+            self.measures['hr_stds'] = np.std(hr)
+            self.measures['lr_means'] = np.mean(lr)
+            self.measures['lr_stds'] = np.std(lr)
+
+        hr = torch.Tensor(hr)
+        lr = torch.Tensor(lr)
+
+        # if self.gpu:
+        #    hr = hr.cuda()
+        #    lr = lr.cuda()
+
+        return {'LQ': lr, 'GT': hr, 'LQ_path': str(item), 'GT_path': str(item)}
+
+    def print_and_reset(self, tag):
+        m = self.measures
+        kvs = []
+        for k in sorted(m.keys()):
+            kvs.append("{}={:.2f}".format(k, m[k]))
+        print("[KPI] " + tag + ": " + ", ".join(kvs))
+        self.measures = None
+
+
+def random_flip(img, seg):
+    random_choice = np.random.choice([True, False])
+    img = img if random_choice else np.flip(img, 2).copy()
+    seg = seg if random_choice else np.flip(seg, 2).copy()
+    return img, seg
+
+
+def random_rotation(img, seg):
+    random_choice = np.random.choice([0, 1, 3])
+    img = np.rot90(img, random_choice, axes=(1, 2)).copy()
+    seg = np.rot90(seg, random_choice, axes=(1, 2)).copy()
+    return img, seg
+
+
+def random_crop(hr, lr, size_hr, scale, random):
+    size_lr = size_hr // scale
+
+    size_lr_x = lr.shape[1]
+    size_lr_y = lr.shape[2]
+
+    start_x_lr = np.random.randint(low=0, high=(size_lr_x - size_lr) + 1) if size_lr_x > size_lr else 0
+    start_y_lr = np.random.randint(low=0, high=(size_lr_y - size_lr) + 1) if size_lr_y > size_lr else 0
+
+    # LR Patch
+    lr_patch = lr[:, start_x_lr:start_x_lr + size_lr, start_y_lr:start_y_lr + size_lr]
+
+    # HR Patch
+    start_x_hr = start_x_lr * scale
+    start_y_hr = start_y_lr * scale
+    hr_patch = hr[:, start_x_hr:start_x_hr + size_hr, start_y_hr:start_y_hr + size_hr]
+
+    return hr_patch, lr_patch
+
+
+def center_crop(img, size):
+    assert img.shape[1] == img.shape[2], img.shape
+    border_double = img.shape[1] - size
+    assert border_double % 2 == 0, (img.shape, size)
+    border = border_double // 2
+    return img[:, border:-border, border:-border]
+
+
+def center_crop_tensor(img, size):
+    assert img.shape[2] == img.shape[3], img.shape
+    border_double = img.shape[2] - size
+    assert border_double % 2 == 0, (img.shape, size)
+    border = border_double // 2
+    return img[:, :, border:-border, border:-border]

models/SRFlow/code/data/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+'''create dataset and dataloader'''
+import logging
+import torch
+import torch.utils.data
+
+
+def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
+    phase = dataset_opt.get('phase', 'test')
+    if phase == 'train':
+        gpu_ids = opt.get('gpu_ids', None)
+        gpu_ids = gpu_ids if gpu_ids else []
+        num_workers = dataset_opt['n_workers'] * len(gpu_ids)
+        batch_size = dataset_opt['batch_size']
+        shuffle = True
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
+                                           num_workers=num_workers, sampler=sampler, drop_last=True,
+                                           pin_memory=False)
+    else:
+        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1,
+                                           pin_memory=True)
+
+
+def create_dataset(dataset_opt):
+    print(dataset_opt)
+    mode = dataset_opt['mode']
+    if mode == 'LRHR_PKL':
+        from data.LRHR_PKL_dataset import LRHR_PKLDataset as D
+    else:
+        raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
+    dataset = D(dataset_opt)
+
+    logger = logging.getLogger('base')
+    logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
+                                                           dataset_opt['name']))
+    return dataset

models/SRFlow/code/imresize.py

@@ -0,0 +1,180 @@
+        # https://github.com/fatheral/matlab_imresize
+#
+# MIT License
+#
+# Copyright (c) 2020 Alex
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+from __future__ import print_function
+import numpy as np
+from math import ceil, floor
+
+
+def deriveSizeFromScale(img_shape, scale):
+    output_shape = []
+    for k in range(2):
+        output_shape.append(int(ceil(scale[k] * img_shape[k])))
+    return output_shape
+
+
+def deriveScaleFromSize(img_shape_in, img_shape_out):
+    scale = []
+    for k in range(2):
+        scale.append(1.0 * img_shape_out[k] / img_shape_in[k])
+    return scale
+
+
+def triangle(x):
+    x = np.array(x).astype(np.float64)
+    lessthanzero = np.logical_and((x >= -1), x < 0)
+    greaterthanzero = np.logical_and((x <= 1), x >= 0)
+    f = np.multiply((x + 1), lessthanzero) + np.multiply((1 - x), greaterthanzero)
+    return f
+
+
+def cubic(x):
+    x = np.array(x).astype(np.float64)
+    absx = np.absolute(x)
+    absx2 = np.multiply(absx, absx)
+    absx3 = np.multiply(absx2, absx)
+    f = np.multiply(1.5 * absx3 - 2.5 * absx2 + 1, absx <= 1) + np.multiply(-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2,
+                                                                            (1 < absx) & (absx <= 2))
+    return f
+
+
+def contributions(in_length, out_length, scale, kernel, k_width):
+    if scale < 1:
+        h = lambda x: scale * kernel(scale * x)
+        kernel_width = 1.0 * k_width / scale
+    else:
+        h = kernel
+        kernel_width = k_width
+    x = np.arange(1, out_length + 1).astype(np.float64)
+    u = x / scale + 0.5 * (1 - 1 / scale)
+    left = np.floor(u - kernel_width / 2)
+    P = int(ceil(kernel_width)) + 2
+    ind = np.expand_dims(left, axis=1) + np.arange(P) - 1  # -1 because indexing from 0
+    indices = ind.astype(np.int32)
+    weights = h(np.expand_dims(u, axis=1) - indices - 1)  # -1 because indexing from 0
+    weights = np.divide(weights, np.expand_dims(np.sum(weights, axis=1), axis=1))
+    aux = np.concatenate((np.arange(in_length), np.arange(in_length - 1, -1, step=-1))).astype(np.int32)
+    indices = aux[np.mod(indices, aux.size)]
+    ind2store = np.nonzero(np.any(weights, axis=0))
+    weights = weights[:, ind2store]
+    indices = indices[:, ind2store]
+    return weights, indices
+
+
+def imresizemex(inimg, weights, indices, dim):
+    in_shape = inimg.shape
+    w_shape = weights.shape
+    out_shape = list(in_shape)
+    out_shape[dim] = w_shape[0]
+    outimg = np.zeros(out_shape)
+    if dim == 0:
+        for i_img in range(in_shape[1]):
+            for i_w in range(w_shape[0]):
+                w = weights[i_w, :]
+                ind = indices[i_w, :]
+                im_slice = inimg[ind, i_img].astype(np.float64)
+                outimg[i_w, i_img] = np.sum(np.multiply(np.squeeze(im_slice, axis=0), w.T), axis=0)
+    elif dim == 1:
+        for i_img in range(in_shape[0]):
+            for i_w in range(w_shape[0]):
+                w = weights[i_w, :]
+                ind = indices[i_w, :]
+                im_slice = inimg[i_img, ind].astype(np.float64)
+                outimg[i_img, i_w] = np.sum(np.multiply(np.squeeze(im_slice, axis=0), w.T), axis=0)
+    if inimg.dtype == np.uint8:
+        outimg = np.clip(outimg, 0, 255)
+        return np.around(outimg).astype(np.uint8)
+    else:
+        return outimg
+
+
+def imresizevec(inimg, weights, indices, dim):
+    wshape = weights.shape
+    if dim == 0:
+        weights = weights.reshape((wshape[0], wshape[2], 1, 1))
+        outimg = np.sum(weights * ((inimg[indices].squeeze(axis=1)).astype(np.float64)), axis=1)
+    elif dim == 1:
+        weights = weights.reshape((1, wshape[0], wshape[2], 1))
+        outimg = np.sum(weights * ((inimg[:, indices].squeeze(axis=2)).astype(np.float64)), axis=2)
+    if inimg.dtype == np.uint8:
+        outimg = np.clip(outimg, 0, 255)
+        return np.around(outimg).astype(np.uint8)
+    else:
+        return outimg
+
+
+def resizeAlongDim(A, dim, weights, indices, mode="vec"):
+    if mode == "org":
+        out = imresizemex(A, weights, indices, dim)
+    else:
+        out = imresizevec(A, weights, indices, dim)
+    return out
+
+
+def imresize(I, scalar_scale=None, method='bicubic', output_shape=None, mode="vec"):
+    if method is 'bicubic':
+        kernel = cubic
+    elif method is 'bilinear':
+        kernel = triangle
+    else:
+        print('Error: Unidentified method supplied')
+
+    kernel_width = 4.0
+    # Fill scale and output_size
+    if scalar_scale is not None:
+        scalar_scale = float(scalar_scale)
+        scale = [scalar_scale, scalar_scale]
+        output_size = deriveSizeFromScale(I.shape, scale)
+    elif output_shape is not None:
+        scale = deriveScaleFromSize(I.shape, output_shape)
+        output_size = list(output_shape)
+    else:
+        print('Error: scalar_scale OR output_shape should be defined!')
+        return
+    scale_np = np.array(scale)
+    order = np.argsort(scale_np)
+    weights = []
+    indices = []
+    for k in range(2):
+        w, ind = contributions(I.shape[k], output_size[k], scale[k], kernel, kernel_width)
+        weights.append(w)
+        indices.append(ind)
+    B = np.copy(I)
+    flag2D = False
+    if B.ndim == 2:
+        B = np.expand_dims(B, axis=2)
+        flag2D = True
+    for k in range(2):
+        dim = order[k]
+        B = resizeAlongDim(B, dim, weights[dim], indices[dim], mode)
+    if flag2D:
+        B = np.squeeze(B, axis=2)
+    return B
+
+
+def convertDouble2Byte(I):
+    B = np.clip(I, 0.0, 1.0)
+    B = 255 * B
+    return np.around(B).astype(np.uint8)

models/SRFlow/code/models/SRFlow_model.py

@@ -0,0 +1,278 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import logging
+from collections import OrderedDict
+from utils.util import get_resume_paths, opt_get
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from .base_model import BaseModel
+
+logger = logging.getLogger('base')
+
+
+class SRFlowModel(BaseModel):
+    def __init__(self, opt, step):
+        super(SRFlowModel, self).__init__(opt)
+        self.opt = opt
+
+        self.heats = opt['val']['heats']
+        self.n_sample = opt['val']['n_sample']
+        self.hr_size = opt_get(opt, ['datasets', 'train', 'center_crop_hr_size'])
+        self.hr_size = 256 if self.hr_size is None else self.hr_size
+        self.lr_size = self.hr_size // opt['scale']
+
+        if opt['dist']:
+            self.rank = torch.distributed.get_rank()
+        else:
+            self.rank = -1  # non dist training
+        train_opt = opt['train']
+
+        # define network and load pretrained models
+        self.netG = networks.define_Flow(opt, step).to(self.device)
+        if opt['dist']:
+            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+        else:
+            self.netG = DataParallel(self.netG)
+        # print network
+        self.print_network()
+
+        if opt_get(opt, ['path', 'resume_state'], 1) is not None:
+            self.load()
+        else:
+            print("WARNING: skipping initial loading, due to resume_state None")
+
+        if self.is_train:
+            self.netG.train()
+
+            self.init_optimizer_and_scheduler(train_opt)
+            self.log_dict = OrderedDict()
+
+    def to(self, device):
+        self.device = device
+        self.netG.to(device)
+
+    def init_optimizer_and_scheduler(self, train_opt):
+        # optimizers
+        self.optimizers = []
+        wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+        optim_params_RRDB = []
+        optim_params_other = []
+        for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+            print(k, v.requires_grad)
+            if v.requires_grad:
+                if '.RRDB.' in k:
+                    optim_params_RRDB.append(v)
+                    print('opt', k)
+                else:
+                    optim_params_other.append(v)
+                if self.rank <= 0:
+                    logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+        print('rrdb params', len(optim_params_RRDB))
+
+        self.optimizer_G = torch.optim.Adam(
+            [
+                {"params": optim_params_other, "lr": train_opt['lr_G'], 'beta1': train_opt['beta1'],
+                 'beta2': train_opt['beta2'], 'weight_decay': wd_G},
+                {"params": optim_params_RRDB, "lr": train_opt.get('lr_RRDB', train_opt['lr_G']),
+                 'beta1': train_opt['beta1'],
+                 'beta2': train_opt['beta2'], 'weight_decay': wd_G}
+            ],
+        )
+
+        self.optimizers.append(self.optimizer_G)
+        # schedulers
+        if train_opt['lr_scheme'] == 'MultiStepLR':
+            for optimizer in self.optimizers:
+                self.schedulers.append(
+                    lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+                                                     restarts=train_opt['restarts'],
+                                                     weights=train_opt['restart_weights'],
+                                                     gamma=train_opt['lr_gamma'],
+                                                     clear_state=train_opt['clear_state'],
+                                                     lr_steps_invese=train_opt.get('lr_steps_inverse', [])))
+        elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+            for optimizer in self.optimizers:
+                self.schedulers.append(
+                    lr_scheduler.CosineAnnealingLR_Restart(
+                        optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+                        restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+        else:
+            raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+    def add_optimizer_and_scheduler_RRDB(self, train_opt):
+        # optimizers
+        assert len(self.optimizers) == 1, self.optimizers
+        assert len(self.optimizer_G.param_groups[1]['params']) == 0, self.optimizer_G.param_groups[1]
+        for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+            if v.requires_grad:
+                if '.RRDB.' in k:
+                    self.optimizer_G.param_groups[1]['params'].append(v)
+        assert len(self.optimizer_G.param_groups[1]['params']) > 0
+
+    def feed_data(self, data, need_GT=True):
+        self.var_L = data['LQ'].to(self.device)  # LQ
+        if need_GT:
+            self.real_H = data['GT'].to(self.device)  # GT
+
+    def optimize_parameters(self, step):
+
+        train_RRDB_delay = opt_get(self.opt, ['network_G', 'train_RRDB_delay'])
+        if train_RRDB_delay is not None and step > int(train_RRDB_delay * self.opt['train']['niter']) \
+                and not self.netG.module.RRDB_training:
+            if self.netG.module.set_rrdb_training(True):
+                self.add_optimizer_and_scheduler_RRDB(self.opt['train'])
+
+        # self.print_rrdb_state()
+
+        self.netG.train()
+        self.log_dict = OrderedDict()
+        self.optimizer_G.zero_grad()
+
+        losses = {}
+        weight_fl = opt_get(self.opt, ['train', 'weight_fl'])
+        weight_fl = 1 if weight_fl is None else weight_fl
+        if weight_fl > 0:
+            z, nll, y_logits = self.netG(gt=self.real_H, lr=self.var_L, reverse=False)
+            nll_loss = torch.mean(nll)
+            losses['nll_loss'] = nll_loss * weight_fl
+
+        weight_l1 = opt_get(self.opt, ['train', 'weight_l1']) or 0
+        if weight_l1 > 0:
+            z = self.get_z(heat=0, seed=None, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+            sr, logdet = self.netG(lr=self.var_L, z=z, eps_std=0, reverse=True, reverse_with_grad=True)
+            l1_loss = (sr - self.real_H).abs().mean()
+            losses['l1_loss'] = l1_loss * weight_l1
+
+        total_loss = sum(losses.values())
+        total_loss.backward()
+        self.optimizer_G.step()
+
+        mean = total_loss.item()
+        return mean
+
+    def print_rrdb_state(self):
+        for name, param in self.netG.module.named_parameters():
+            if "RRDB.conv_first.weight" in name:
+                print(name, param.requires_grad, param.data.abs().sum())
+        print('params', [len(p['params']) for p in self.optimizer_G.param_groups])
+
+    def test(self):
+        self.netG.eval()
+        self.fake_H = {}
+        for heat in self.heats:
+            for i in range(self.n_sample):
+                z = self.get_z(heat, seed=None, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+                with torch.no_grad():
+                    self.fake_H[(heat, i)], logdet = self.netG(lr=self.var_L, z=z, eps_std=heat, reverse=True)
+        with torch.no_grad():
+            _, nll, _ = self.netG(gt=self.real_H, lr=self.var_L, reverse=False)
+        self.netG.train()
+        return nll.mean().item()
+
+    def get_encode_nll(self, lq, gt):
+        self.netG.eval()
+        with torch.no_grad():
+            _, nll, _ = self.netG(gt=gt, lr=lq, reverse=False)
+        self.netG.train()
+        return nll.mean().item()
+
+    def get_sr(self, lq, heat=None, seed=None, z=None, epses=None):
+        return self.get_sr_with_z(lq, heat, seed, z, epses)[0]
+
+    def get_encode_z(self, lq, gt, epses=None, add_gt_noise=True):
+        self.netG.eval()
+        with torch.no_grad():
+            z, _, _ = self.netG(gt=gt, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise)
+        self.netG.train()
+        return z
+
+    def get_encode_z_and_nll(self, lq, gt, epses=None, add_gt_noise=True):
+        self.netG.eval()
+        with torch.no_grad():
+            z, nll, _ = self.netG(gt=gt, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise)
+        self.netG.train()
+        return z, nll
+
+    def get_sr_with_z(self, lq, heat=None, seed=None, z=None, epses=None):
+        self.netG.eval()
+
+        z = self.get_z(heat, seed, batch_size=lq.shape[0], lr_shape=lq.shape) if z is None and epses is None else z
+
+        with torch.no_grad():
+            sr, logdet = self.netG(lr=lq, z=z, eps_std=heat, reverse=True, epses=epses)
+        self.netG.train()
+        return sr, z
+
+    def get_z(self, heat, seed=None, batch_size=1, lr_shape=None):
+        if seed: torch.manual_seed(seed)
+        if opt_get(self.opt, ['network_G', 'flow', 'split', 'enable']):
+            C = self.netG.module.flowUpsamplerNet.C
+            H = int(self.opt['scale'] * lr_shape[2] // self.netG.module.flowUpsamplerNet.scaleH)
+            W = int(self.opt['scale'] * lr_shape[3] // self.netG.module.flowUpsamplerNet.scaleW)
+            z = torch.normal(mean=0, std=heat, size=(batch_size, C, H, W)) if heat > 0 else torch.zeros(
+                (batch_size, C, H, W))
+        else:
+            L = opt_get(self.opt, ['network_G', 'flow', 'L']) or 3
+            fac = 2 ** (L - 3)
+            z_size = int(self.lr_size // (2 ** (L - 3)))
+            z = torch.normal(mean=0, std=heat, size=(batch_size, 3 * 8 * 8 * fac * fac, z_size, z_size))
+        return z
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_GT=True):
+        out_dict = OrderedDict()
+        out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+        for heat in self.heats:
+            for i in range(self.n_sample):
+                out_dict[('SR', heat, i)] = self.fake_H[(heat, i)].detach()[0].float().cpu()
+        if need_GT:
+            out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        if self.rank <= 0:
+            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+    def load(self):
+        _, get_resume_model_path = get_resume_paths(self.opt)
+        if get_resume_model_path is not None:
+            self.load_network(get_resume_model_path, self.netG, strict=True, submodule=None)
+            return
+
+        load_path_G = self.opt['path']['pretrain_model_G']
+        load_submodule = self.opt['path']['load_submodule'] if 'load_submodule' in self.opt['path'].keys() else 'RRDB'
+        if load_path_G is not None:
+            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG, self.opt['path'].get('strict_load', True),
+                              submodule=load_submodule)
+
+    def save(self, iter_label):
+        self.save_network(self.netG, 'G', iter_label)

models/SRFlow/code/models/SR_model.py

@@ -0,0 +1,217 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from utils.util import opt_get
+from .base_model import BaseModel
+from models.modules.loss import CharbonnierLoss
+
+logger = logging.getLogger('base')
+
+
+class SRModel(BaseModel):
+    def __init__(self, opt, step):
+        super(SRModel, self).__init__(opt)
+
+        self.step = step
+
+        if opt['dist']:
+            self.rank = torch.distributed.get_rank()
+        else:
+            self.rank = -1  # non dist training
+        train_opt = opt['train']
+
+        # define network and load pretrained_models models
+        self.netG = networks.define_G(opt).to(self.device)
+        if opt['dist']:
+            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+        else:
+            self.netG = DataParallel(self.netG)
+        # print network
+        self.print_network()
+        self.load()
+
+        if self.is_train:
+            self.netG.train()
+
+            # loss
+            loss_type = train_opt['pixel_criterion']
+            if loss_type == 'l1':
+                self.cri_pix = nn.L1Loss().to(self.device)
+            elif loss_type == 'l2':
+                self.cri_pix = nn.MSELoss().to(self.device)
+            elif loss_type == 'cb':
+                self.cri_pix = CharbonnierLoss().to(self.device)
+            else:
+                raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+            self.l_pix_w = train_opt['pixel_weight']
+
+            # optimizers
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if v.requires_grad:
+                    optim_params.append(v)
+                else:
+                    if self.rank <= 0:
+                        logger.warning('Params [{:s}] will not optimize.'.format(k))
+            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
+                                                weight_decay=wd_G,
+                                                betas=(train_opt['beta1'], train_opt['beta2']))
+            self.optimizers.append(self.optimizer_G)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+                                                         restarts=train_opt['restarts'],
+                                                         weights=train_opt['restart_weights'],
+                                                         gamma=train_opt['lr_gamma'],
+                                                         clear_state=train_opt['clear_state']))
+            elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.CosineAnnealingLR_Restart(
+                            optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+                            restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+
+    def feed_data(self, data, need_GT=True):
+        self.var_L = data['LQ'].to(self.device)  # LQ
+        if need_GT:
+            self.real_H = data['GT'].to(self.device)  # GT
+
+    def to(self, device):
+        self.device = device
+        self.netG.to(device)
+
+    def optimize_parameters(self, step):
+        def getEnv(name): import os; return True if name in os.environ.keys() else False
+
+        if getEnv("DEBUG_FEED_IMAGES"):
+            import imageio
+            import random
+            i = random.randint(0, 10000)
+            label = self.var_L.cpu().numpy()[0].transpose([1, 2, 0])
+            print("var_L", label.min(), label.max(), label.shape)
+            imageio.imwrite("/tmp/{}_l.png".format(i), label)
+            image = self.real_H.cpu().numpy()[0].transpose([1, 2, 0])
+            print("self.real_H", image.min(), image.max(), image.shape)
+            imageio.imwrite("/tmp/{}_gt.png".format(i), image)
+        self.optimizer_G.zero_grad()
+        self.fake_H = self.netG(self.var_L)
+        l_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.real_H.to(self.fake_H.device))
+        l_pix.backward()
+        self.optimizer_G.step()
+
+        # set log
+        self.log_dict['l_pix'] = l_pix.item()
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG(self.var_L)
+        self.netG.train()
+
+    def get_encode_nll(self, lq, gt):
+        return torch.ones(1) * 1e14
+
+    def get_sr(self, lq, heat=None, seed=None):
+        self.netG.eval()
+        sr = self.netG(lq)
+        self.netG.train()
+        return sr
+
+    def test_x8(self):
+        # from https://github.com/thstkdgus35/EDSR-PyTorch
+        self.netG.eval()
+
+        def _transform(v, op):
+            # if self.precision != 'single': v = v.float()
+            v2np = v.data.cpu().numpy()
+            if op == 'v':
+                tfnp = v2np[:, :, :, ::-1].copy()
+            elif op == 'h':
+                tfnp = v2np[:, :, ::-1, :].copy()
+            elif op == 't':
+                tfnp = v2np.transpose((0, 1, 3, 2)).copy()
+
+            ret = torch.Tensor(tfnp).to(self.device)
+            # if self.precision == 'half': ret = ret.half()
+
+            return ret
+
+        lr_list = [self.var_L]
+        for tf in 'v', 'h', 't':
+            lr_list.extend([_transform(t, tf) for t in lr_list])
+        with torch.no_grad():
+            sr_list = [self.netG(aug) for aug in lr_list]
+        for i in range(len(sr_list)):
+            if i > 3:
+                sr_list[i] = _transform(sr_list[i], 't')
+            if i % 4 > 1:
+                sr_list[i] = _transform(sr_list[i], 'h')
+            if (i % 4) % 2 == 1:
+                sr_list[i] = _transform(sr_list[i], 'v')
+
+        output_cat = torch.cat(sr_list, dim=0)
+        self.fake_H = output_cat.mean(dim=0, keepdim=True)
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_GT=True):
+        out_dict = OrderedDict()
+        out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_GT:
+            out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        if self.rank <= 0:
+            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
+
+    def save(self, iter_label):
+        self.save_network(self.netG, 'G', iter_label)
+
+    def get_encode_z_and_nll(self, *args, **kwargs):
+        return [], torch.zeros(1)

models/SRFlow/code/models/__init__.py

@@ -0,0 +1,52 @@
+import importlib
+import logging
+import os
+
+try:
+    import local_config
+except:
+    local_config = None
+
+
+logger = logging.getLogger('base')
+
+
+def find_model_using_name(model_name):
+    # Given the option --model [modelname],
+    # the file "models/modelname_model.py"
+    # will be imported.
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+
+    # In the file, the class called ModelNameModel() will
+    # be instantiated. It has to be a subclass of torch.nn.Module,
+    # and it is case-insensitive.
+    model = None
+    target_model_name = model_name.replace('_', '') + 'Model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower():
+            model = cls
+
+    if model is None:
+        print(
+            "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s." % (
+            model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def create_model(opt, step=0, **opt_kwargs):
+    if local_config is not None:
+        opt['path']['pretrain_model_G'] = os.path.join(local_config.checkpoint_path, os.path.basename(opt['path']['results_root'] + '.pth'))
+
+    for k, v in opt_kwargs.items():
+        opt[k] = v
+
+    model = opt['model']
+
+    M = find_model_using_name(model)
+
+    m = M(opt, step)
+    logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
+    return m

models/SRFlow/code/models/base_model.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+import natsort
+import glob
+
+
+class BaseModel():
+    def __init__(self, opt):
+        self.opt = opt
+        self.device = torch.device('cuda' if opt.get('gpu_ids', None) is not None else 'cpu')
+        self.is_train = opt['is_train']
+        self.schedulers = []
+        self.optimizers = []
+
+    def feed_data(self, data):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        pass
+
+    def get_current_losses(self):
+        pass
+
+    def print_network(self):
+        pass
+
+    def save(self, label):
+        pass
+
+    def load(self):
+        pass
+
+    def _set_lr(self, lr_groups_l):
+        ''' set learning rate for warmup,
+        lr_groups_l: list for lr_groups. each for a optimizer'''
+        for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
+            for param_group, lr in zip(optimizer.param_groups, lr_groups):
+                param_group['lr'] = lr
+
+    def _get_init_lr(self):
+        # get the initial lr, which is set by the scheduler
+        init_lr_groups_l = []
+        for optimizer in self.optimizers:
+            init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
+        return init_lr_groups_l
+
+    def update_learning_rate(self, cur_iter, warmup_iter=-1):
+        for scheduler in self.schedulers:
+            scheduler.step()
+        #### set up warm up learning rate
+        if cur_iter < warmup_iter:
+            # get initial lr for each group
+            init_lr_g_l = self._get_init_lr()
+            # modify warming-up learning rates
+            warm_up_lr_l = []
+            for init_lr_g in init_lr_g_l:
+                warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g])
+            # set learning rate
+            self._set_lr(warm_up_lr_l)
+
+    def get_current_learning_rate(self):
+        # return self.schedulers[0].get_lr()[0]
+        return self.optimizers[0].param_groups[0]['lr']
+
+    def get_network_description(self, network):
+        '''Get the string and total parameters of the network'''
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        s = str(network)
+        n = sum(map(lambda x: x.numel(), network.parameters()))
+        return s, n
+
+    def save_network(self, network, network_label, iter_label):
+        paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['models'], "*_{}.pth".format(network_label))),
+                                  reverse=True)
+        paths = [p for p in paths if
+                 "latest_" not in p and not any([str(i * 10000) in p.split("/")[-1].split("_") for i in range(101)])]
+        if len(paths) > 2:
+            for path in paths[2:]:
+                os.remove(path)
+        save_filename = '{}_{}.pth'.format(iter_label, network_label)
+        save_path = os.path.join(self.opt['path']['models'], save_filename)
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        state_dict = network.state_dict()
+        for key, param in state_dict.items():
+            state_dict[key] = param.cpu()
+        torch.save(state_dict, save_path)
+
+    def load_network(self, load_path, network, strict=True, submodule=None):
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        if not (submodule is None or submodule.lower() == 'none'.lower()):
+            network = network.__getattr__(submodule)
+        load_net = torch.load(load_path)
+        load_net_clean = OrderedDict()  # remove unnecessary 'module.'
+        for k, v in load_net.items():
+            if k.startswith('module.'):
+                load_net_clean[k[7:]] = v
+            else:
+                load_net_clean[k] = v
+        network.load_state_dict(load_net_clean, strict=strict)
+
+    def save_training_state(self, epoch, iter_step):
+        '''Saves training state during training, which will be used for resuming'''
+        state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []}
+        for s in self.schedulers:
+            state['schedulers'].append(s.state_dict())
+        for o in self.optimizers:
+            state['optimizers'].append(o.state_dict())
+        save_filename = '{}.state'.format(iter_step)
+        save_path = os.path.join(self.opt['path']['training_state'], save_filename)
+
+        paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['training_state'], "*.state")),
+                                  reverse=True)
+        paths = [p for p in paths if "latest_" not in p]
+        if len(paths) > 2:
+            for path in paths[2:]:
+                os.remove(path)
+
+        torch.save(state, save_path)
+
+    def resume_training(self, resume_state):
+        '''Resume the optimizers and schedulers for training'''
+        resume_optimizers = resume_state['optimizers']
+        resume_schedulers = resume_state['schedulers']
+        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+        for i, o in enumerate(resume_optimizers):
+            self.optimizers[i].load_state_dict(o)
+        for i, s in enumerate(resume_schedulers):
+            self.schedulers[i].load_state_dict(s)

models/SRFlow/code/models/lr_scheduler.py

@@ -0,0 +1,163 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import math
+from collections import Counter
+from collections import defaultdict
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class MultiStepLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, milestones, restarts=None, weights=None, gamma=0.1,
+                 clear_state=False, last_epoch=-1, lr_steps_invese=None):
+        assert lr_steps_invese is not None, "Use empty list"
+        self.milestones = Counter(milestones)
+        self.lr_steps_inverse = Counter(lr_steps_invese)
+        self.gamma = gamma
+        self.clear_state = clear_state
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(MultiStepLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch in self.restarts:
+            if self.clear_state:
+                self.optimizer.state = defaultdict(dict)
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        if self.last_epoch not in self.milestones and self.last_epoch not in self.lr_steps_inverse:
+            return [group['lr'] for group in self.optimizer.param_groups]
+        return [
+            group['lr'] * (self.gamma ** self.milestones[self.last_epoch]) *
+            (self.gamma ** (-self.lr_steps_inverse[self.last_epoch]))
+            for group in self.optimizer.param_groups
+        ]
+
+
+class CosineAnnealingLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, T_period, restarts=None, weights=None, eta_min=0, last_epoch=-1):
+        self.T_period = T_period
+        self.T_max = self.T_period[0]  # current T period
+        self.eta_min = eta_min
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        self.last_restart = 0
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(CosineAnnealingLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch == 0:
+            return self.base_lrs
+        elif self.last_epoch in self.restarts:
+            self.last_restart = self.last_epoch
+            self.T_max = self.T_period[self.restarts.index(self.last_epoch) + 1]
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        elif (self.last_epoch - self.last_restart - 1 - self.T_max) % (2 * self.T_max) == 0:
+            return [
+                group['lr'] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2
+                for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            ]
+        return [(1 + math.cos(math.pi * (self.last_epoch - self.last_restart) / self.T_max)) /
+                (1 + math.cos(math.pi * ((self.last_epoch - self.last_restart) - 1) / self.T_max)) *
+                (group['lr'] - self.eta_min) + self.eta_min
+                for group in self.optimizer.param_groups]
+
+
+if __name__ == "__main__":
+    optimizer = torch.optim.Adam([torch.zeros(3, 64, 3, 3)], lr=2e-4, weight_decay=0,
+                                 betas=(0.9, 0.99))
+    ##############################
+    # MultiStepLR_Restart
+    ##############################
+    ## Original
+    lr_steps = [200000, 400000, 600000, 800000]
+    restarts = None
+    restart_weights = None
+
+    ## two
+    lr_steps = [100000, 200000, 300000, 400000, 490000, 600000, 700000, 800000, 900000, 990000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    lr_steps = [
+        50000, 100000, 150000, 200000, 240000, 300000, 350000, 400000, 450000, 490000, 550000,
+        600000, 650000, 700000, 740000, 800000, 850000, 900000, 950000, 990000
+    ]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = MultiStepLR_Restart(optimizer, lr_steps, restarts, restart_weights, gamma=0.5,
+                                    clear_state=False)
+
+    ##############################
+    # Cosine Annealing Restart
+    ##############################
+    ## two
+    T_period = [500000, 500000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    T_period = [250000, 250000, 250000, 250000]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = CosineAnnealingLR_Restart(optimizer, T_period, eta_min=1e-7, restarts=restarts,
+                                          weights=restart_weights)
+
+    ##############################
+    # Draw figure
+    ##############################
+    N_iter = 1000000
+    lr_l = list(range(N_iter))
+    for i in range(N_iter):
+        scheduler.step()
+        current_lr = optimizer.param_groups[0]['lr']
+        lr_l[i] = current_lr
+
+    import matplotlib as mpl
+    from matplotlib import pyplot as plt
+    import matplotlib.ticker as mtick
+
+    mpl.style.use('default')
+    import seaborn
+
+    seaborn.set(style='whitegrid')
+    seaborn.set_context('paper')
+
+    plt.figure(1)
+    plt.subplot(111)
+    plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0))
+    plt.title('Title', fontsize=16, color='k')
+    plt.plot(list(range(N_iter)), lr_l, linewidth=1.5, label='learning rate scheme')
+    legend = plt.legend(loc='upper right', shadow=False)
+    ax = plt.gca()
+    labels = ax.get_xticks().tolist()
+    for k, v in enumerate(labels):
+        labels[k] = str(int(v / 1000)) + 'K'
+    ax.set_xticklabels(labels)
+    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
+
+    ax.set_ylabel('Learning rate')
+    ax.set_xlabel('Iteration')
+    fig = plt.gcf()
+    plt.show()

models/SRFlow/code/models/modules/FlowActNorms.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+
+from models.modules import thops
+
+
+class _ActNorm(nn.Module):
+    """
+    Activation Normalization
+    Initialize the bias and scale with a given minibatch,
+    so that the output per-channel have zero mean and unit variance for that.
+
+    After initialization, `bias` and `logs` will be trained as parameters.
+    """
+
+    def __init__(self, num_features, scale=1.):
+        super().__init__()
+        # register mean and scale
+        size = [1, num_features, 1, 1]
+        self.register_parameter("bias", nn.Parameter(torch.zeros(*size)))
+        self.register_parameter("logs", nn.Parameter(torch.zeros(*size)))
+        self.num_features = num_features
+        self.scale = float(scale)
+        self.inited = False
+
+    def _check_input_dim(self, input):
+        return NotImplemented
+
+    def initialize_parameters(self, input):
+        self._check_input_dim(input)
+        if not self.training:
+            return
+        if (self.bias != 0).any():
+            self.inited = True
+            return
+        assert input.device == self.bias.device, (input.device, self.bias.device)
+        with torch.no_grad():
+            bias = thops.mean(input.clone(), dim=[0, 2, 3], keepdim=True) * -1.0
+            vars = thops.mean((input.clone() + bias) ** 2, dim=[0, 2, 3], keepdim=True)
+            logs = torch.log(self.scale / (torch.sqrt(vars) + 1e-6))
+            self.bias.data.copy_(bias.data)
+            self.logs.data.copy_(logs.data)
+            self.inited = True
+
+    def _center(self, input, reverse=False, offset=None):
+        bias = self.bias
+
+        if offset is not None:
+            bias = bias + offset
+
+        if not reverse:
+            return input + bias
+        else:
+            return input - bias
+
+    def _scale(self, input, logdet=None, reverse=False, offset=None):
+        logs = self.logs
+
+        if offset is not None:
+            logs = logs + offset
+
+        if not reverse:
+            input = input * torch.exp(logs) # should have shape batchsize, n_channels, 1, 1
+            # input = input * torch.exp(logs+logs_offset)
+        else:
+            input = input * torch.exp(-logs)
+        if logdet is not None:
+            """
+            logs is log_std of `mean of channels`
+            so we need to multiply pixels
+            """
+            dlogdet = thops.sum(logs) * thops.pixels(input)
+            if reverse:
+                dlogdet *= -1
+            logdet = logdet + dlogdet
+        return input, logdet
+
+    def forward(self, input, logdet=None, reverse=False, offset_mask=None, logs_offset=None, bias_offset=None):
+        if not self.inited:
+            self.initialize_parameters(input)
+        self._check_input_dim(input)
+
+        if offset_mask is not None:
+            logs_offset *= offset_mask
+            bias_offset *= offset_mask
+        # no need to permute dims as old version
+        if not reverse:
+            # center and scale
+
+            # self.input = input
+            input = self._center(input, reverse, bias_offset)
+            input, logdet = self._scale(input, logdet, reverse, logs_offset)
+        else:
+            # scale and center
+            input, logdet = self._scale(input, logdet, reverse, logs_offset)
+            input = self._center(input, reverse, bias_offset)
+        return input, logdet
+
+
+class ActNorm2d(_ActNorm):
+    def __init__(self, num_features, scale=1.):
+        super().__init__(num_features, scale)
+
+    def _check_input_dim(self, input):
+        assert len(input.size()) == 4
+        assert input.size(1) == self.num_features, (
+            "[ActNorm]: input should be in shape as `BCHW`,"
+            " channels should be {} rather than {}".format(
+                self.num_features, input.size()))
+
+
+class MaskedActNorm2d(ActNorm2d):
+    def __init__(self, num_features, scale=1.):
+        super().__init__(num_features, scale)
+
+    def forward(self, input, mask, logdet=None, reverse=False):
+
+        assert mask.dtype == torch.bool
+        output, logdet_out = super().forward(input, logdet, reverse)
+
+        input[mask] = output[mask]
+        logdet[mask] = logdet_out[mask]
+
+        return input, logdet
+

models/SRFlow/code/models/modules/FlowAffineCouplingsAblation.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+
+from models.modules import thops
+from models.modules.flow import Conv2d, Conv2dZeros
+from utils.util import opt_get
+
+
+class CondAffineSeparatedAndCond(nn.Module):
+    def __init__(self, in_channels, opt):
+        super().__init__()
+        self.need_features = True
+        self.in_channels = in_channels
+        self.in_channels_rrdb = 320
+        self.kernel_hidden = 1
+        self.affine_eps = 0.0001
+        self.n_hidden_layers = 1
+        hidden_channels = opt_get(opt, ['network_G', 'flow', 'CondAffineSeparatedAndCond', 'hidden_channels'])
+        self.hidden_channels = 64 if hidden_channels is None else hidden_channels
+
+        self.affine_eps = opt_get(opt, ['network_G', 'flow', 'CondAffineSeparatedAndCond', 'eps'],  0.0001)
+
+        self.channels_for_nn = self.in_channels // 2
+        self.channels_for_co = self.in_channels - self.channels_for_nn
+
+        if self.channels_for_nn is None:
+            self.channels_for_nn = self.in_channels // 2
+
+        self.fAffine = self.F(in_channels=self.channels_for_nn + self.in_channels_rrdb,
+                              out_channels=self.channels_for_co * 2,
+                              hidden_channels=self.hidden_channels,
+                              kernel_hidden=self.kernel_hidden,
+                              n_hidden_layers=self.n_hidden_layers)
+
+        self.fFeatures = self.F(in_channels=self.in_channels_rrdb,
+                                out_channels=self.in_channels * 2,
+                                hidden_channels=self.hidden_channels,
+                                kernel_hidden=self.kernel_hidden,
+                                n_hidden_layers=self.n_hidden_layers)
+
+    def forward(self, input: torch.Tensor, logdet=None, reverse=False, ft=None):
+        if not reverse:
+            z = input
+            assert z.shape[1] == self.in_channels, (z.shape[1], self.in_channels)
+
+            # Feature Conditional
+            scaleFt, shiftFt = self.feature_extract(ft, self.fFeatures)
+            z = z + shiftFt
+            z = z * scaleFt
+            logdet = logdet + self.get_logdet(scaleFt)
+
+            # Self Conditional
+            z1, z2 = self.split(z)
+            scale, shift = self.feature_extract_aff(z1, ft, self.fAffine)
+            self.asserts(scale, shift, z1, z2)
+            z2 = z2 + shift
+            z2 = z2 * scale
+
+            logdet = logdet + self.get_logdet(scale)
+            z = thops.cat_feature(z1, z2)
+            output = z
+        else:
+            z = input
+
+            # Self Conditional
+            z1, z2 = self.split(z)
+            scale, shift = self.feature_extract_aff(z1, ft, self.fAffine)
+            self.asserts(scale, shift, z1, z2)
+            z2 = z2 / scale
+            z2 = z2 - shift
+            z = thops.cat_feature(z1, z2)
+            logdet = logdet - self.get_logdet(scale)
+
+            # Feature Conditional
+            scaleFt, shiftFt = self.feature_extract(ft, self.fFeatures)
+            z = z / scaleFt
+            z = z - shiftFt
+            logdet = logdet - self.get_logdet(scaleFt)
+
+            output = z
+        return output, logdet
+
+    def asserts(self, scale, shift, z1, z2):
+        assert z1.shape[1] == self.channels_for_nn, (z1.shape[1], self.channels_for_nn)
+        assert z2.shape[1] == self.channels_for_co, (z2.shape[1], self.channels_for_co)
+        assert scale.shape[1] == shift.shape[1], (scale.shape[1], shift.shape[1])
+        assert scale.shape[1] == z2.shape[1], (scale.shape[1], z1.shape[1], z2.shape[1])
+
+    def get_logdet(self, scale):
+        return thops.sum(torch.log(scale), dim=[1, 2, 3])
+
+    def feature_extract(self, z, f):
+        h = f(z)
+        shift, scale = thops.split_feature(h, "cross")
+        scale = (torch.sigmoid(scale + 2.) + self.affine_eps)
+        return scale, shift
+
+    def feature_extract_aff(self, z1, ft, f):
+        z = torch.cat([z1, ft], dim=1)
+        h = f(z)
+        shift, scale = thops.split_feature(h, "cross")
+        scale = (torch.sigmoid(scale + 2.) + self.affine_eps)
+        return scale, shift
+
+    def split(self, z):
+        z1 = z[:, :self.channels_for_nn]
+        z2 = z[:, self.channels_for_nn:]
+        assert z1.shape[1] + z2.shape[1] == z.shape[1], (z1.shape[1], z2.shape[1], z.shape[1])
+        return z1, z2
+
+    def F(self, in_channels, out_channels, hidden_channels, kernel_hidden=1, n_hidden_layers=1):
+        layers = [Conv2d(in_channels, hidden_channels), nn.ReLU(inplace=False)]
+
+        for _ in range(n_hidden_layers):
+            layers.append(Conv2d(hidden_channels, hidden_channels, kernel_size=[kernel_hidden, kernel_hidden]))
+            layers.append(nn.ReLU(inplace=False))
+        layers.append(Conv2dZeros(hidden_channels, out_channels))
+
+        return nn.Sequential(*layers)

models/SRFlow/code/models/modules/FlowStep.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+
+import models.modules
+import models.modules.Permutations
+from models.modules import flow, thops, FlowAffineCouplingsAblation
+from utils.util import opt_get
+
+
+def getConditional(rrdbResults, position):
+    img_ft = rrdbResults if isinstance(rrdbResults, torch.Tensor) else rrdbResults[position]
+    return img_ft
+
+
+class FlowStep(nn.Module):
+    FlowPermutation = {
+        "reverse": lambda obj, z, logdet, rev: (obj.reverse(z, rev), logdet),
+        "shuffle": lambda obj, z, logdet, rev: (obj.shuffle(z, rev), logdet),
+        "invconv": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "squeeze_invconv": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "resqueeze_invconv_alternating_2_3": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "resqueeze_invconv_3": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlign": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1SubblocksShuf": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlignIndepBorder": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlignIndepBorder4": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+    }
+
+    def __init__(self, in_channels, hidden_channels,
+                 actnorm_scale=1.0, flow_permutation="invconv", flow_coupling="additive",
+                 LU_decomposed=False, opt=None, image_injector=None, idx=None, acOpt=None, normOpt=None, in_shape=None,
+                 position=None):
+        # check configures
+        assert flow_permutation in FlowStep.FlowPermutation, \
+            "float_permutation should be in `{}`".format(
+                FlowStep.FlowPermutation.keys())
+        super().__init__()
+        self.flow_permutation = flow_permutation
+        self.flow_coupling = flow_coupling
+        self.image_injector = image_injector
+
+        self.norm_type = normOpt['type'] if normOpt else 'ActNorm2d'
+        self.position = normOpt['position'] if normOpt else None
+
+        self.in_shape = in_shape
+        self.position = position
+        self.acOpt = acOpt
+
+        # 1. actnorm
+        self.actnorm = models.modules.FlowActNorms.ActNorm2d(in_channels, actnorm_scale)
+
+        # 2. permute
+        if flow_permutation == "invconv":
+            self.invconv = models.modules.Permutations.InvertibleConv1x1(
+                in_channels, LU_decomposed=LU_decomposed)
+
+        # 3. coupling
+        if flow_coupling == "CondAffineSeparatedAndCond":
+            self.affine = models.modules.FlowAffineCouplingsAblation.CondAffineSeparatedAndCond(in_channels=in_channels,
+                                                                                                opt=opt)
+        elif flow_coupling == "noCoupling":
+            pass
+        else:
+            raise RuntimeError("coupling not Found:", flow_coupling)
+
+    def forward(self, input, logdet=None, reverse=False, rrdbResults=None):
+        if not reverse:
+            return self.normal_flow(input, logdet, rrdbResults)
+        else:
+            return self.reverse_flow(input, logdet, rrdbResults)
+
+    def normal_flow(self, z, logdet, rrdbResults=None):
+        if self.flow_coupling == "bentIdentityPreAct":
+            z, logdet = self.bentIdentPar(z, logdet, reverse=False)
+
+        # 1. actnorm
+        if self.norm_type == "ConditionalActNormImageInjector":
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.actnorm(z, img_ft=img_ft, logdet=logdet, reverse=False)
+        elif self.norm_type == "noNorm":
+            pass
+        else:
+            z, logdet = self.actnorm(z, logdet=logdet, reverse=False)
+
+        # 2. permute
+        z, logdet = FlowStep.FlowPermutation[self.flow_permutation](
+            self, z, logdet, False)
+
+        need_features = self.affine_need_features()
+
+        # 3. coupling
+        if need_features or self.flow_coupling in ["condAffine", "condFtAffine", "condNormAffine"]:
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.affine(input=z, logdet=logdet, reverse=False, ft=img_ft)
+        return z, logdet
+
+    def reverse_flow(self, z, logdet, rrdbResults=None):
+
+        need_features = self.affine_need_features()
+
+        # 1.coupling
+        if need_features or self.flow_coupling in ["condAffine", "condFtAffine", "condNormAffine"]:
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.affine(input=z, logdet=logdet, reverse=True, ft=img_ft)
+
+        # 2. permute
+        z, logdet = FlowStep.FlowPermutation[self.flow_permutation](
+            self, z, logdet, True)
+
+        # 3. actnorm
+        z, logdet = self.actnorm(z, logdet=logdet, reverse=True)
+
+        return z, logdet
+
+    def affine_need_features(self):
+        need_features = False
+        try:
+            need_features = self.affine.need_features
+        except:
+            pass
+        return need_features

models/SRFlow/code/models/modules/FlowUpsamplerNet.py

@@ -0,0 +1,309 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import numpy as np
+import torch
+from torch import nn as nn
+
+import models.modules.Split
+from models.modules import flow, thops
+from models.modules.Split import Split2d
+from models.modules.glow_arch import f_conv2d_bias
+from models.modules.FlowStep import FlowStep
+from utils.util import opt_get
+
+
+class FlowUpsamplerNet(nn.Module):
+    def __init__(self, image_shape, hidden_channels, K, L=None,
+                 actnorm_scale=1.0,
+                 flow_permutation=None,
+                 flow_coupling="affine",
+                 LU_decomposed=False, opt=None):
+
+        super().__init__()
+
+        self.layers = nn.ModuleList()
+        self.output_shapes = []
+        self.L = opt_get(opt, ['network_G', 'flow', 'L'])
+        self.K = opt_get(opt, ['network_G', 'flow', 'K'])
+        if isinstance(self.K, int):
+            self.K = [K for K in [K, ] * (self.L + 1)]
+
+        self.opt = opt
+        H, W, self.C = image_shape
+        self.check_image_shape()
+
+        if opt['scale'] == 16:
+            self.levelToName = {
+                0: 'fea_up16',
+                1: 'fea_up8',
+                2: 'fea_up4',
+                3: 'fea_up2',
+                4: 'fea_up1',
+            }
+
+        if opt['scale'] == 8:
+            self.levelToName = {
+                0: 'fea_up8',
+                1: 'fea_up4',
+                2: 'fea_up2',
+                3: 'fea_up1',
+                4: 'fea_up0'
+            }
+
+        elif opt['scale'] == 4:
+            self.levelToName = {
+                0: 'fea_up4',
+                1: 'fea_up2',
+                2: 'fea_up1',
+                3: 'fea_up0',
+                4: 'fea_up-1'
+            }
+
+        affineInCh = self.get_affineInCh(opt_get)
+        flow_permutation = self.get_flow_permutation(flow_permutation, opt)
+
+        normOpt = opt_get(opt, ['network_G', 'flow', 'norm'])
+
+        conditional_channels = {}
+        n_rrdb = self.get_n_rrdb_channels(opt, opt_get)
+        n_bypass_channels = opt_get(opt, ['network_G', 'flow', 'levelConditional', 'n_channels'])
+        conditional_channels[0] = n_rrdb
+        for level in range(1, self.L + 1):
+            # Level 1 gets conditionals from 2, 3, 4 => L - level
+            # Level 2 gets conditionals from 3, 4
+            # Level 3 gets conditionals from 4
+            # Level 4 gets conditionals from None
+            n_bypass = 0 if n_bypass_channels is None else (self.L - level) * n_bypass_channels
+            conditional_channels[level] = n_rrdb + n_bypass
+
+        # Upsampler
+        for level in range(1, self.L + 1):
+            # 1. Squeeze
+            H, W = self.arch_squeeze(H, W)
+
+            # 2. K FlowStep
+            self.arch_additionalFlowAffine(H, LU_decomposed, W, actnorm_scale, hidden_channels, opt)
+            self.arch_FlowStep(H, self.K[level], LU_decomposed, W, actnorm_scale, affineInCh, flow_coupling,
+                               flow_permutation,
+                               hidden_channels, normOpt, opt, opt_get,
+                               n_conditinal_channels=conditional_channels[level])
+            # Split
+            self.arch_split(H, W, level, self.L, opt, opt_get)
+
+        if opt_get(opt, ['network_G', 'flow', 'split', 'enable']):
+            self.f = f_conv2d_bias(affineInCh, 2 * 3 * 64 // 2 // 2)
+        else:
+            self.f = f_conv2d_bias(affineInCh, 2 * 3 * 64)
+
+        self.H = H
+        self.W = W
+        self.scaleH = 160 / H
+        self.scaleW = 160 / W
+
+    def get_n_rrdb_channels(self, opt, opt_get):
+        blocks = opt_get(opt, ['network_G', 'flow', 'stackRRDB', 'blocks'])
+        n_rrdb = 64 if blocks is None else (len(blocks) + 1) * 64
+        return n_rrdb
+
+    def arch_FlowStep(self, H, K, LU_decomposed, W, actnorm_scale, affineInCh, flow_coupling, flow_permutation,
+                      hidden_channels, normOpt, opt, opt_get, n_conditinal_channels=None):
+        condAff = self.get_condAffSetting(opt, opt_get)
+        if condAff is not None:
+            condAff['in_channels_rrdb'] = n_conditinal_channels
+
+        for k in range(K):
+            position_name = get_position_name(H, self.opt['scale'])
+            if normOpt: normOpt['position'] = position_name
+
+            self.layers.append(
+                FlowStep(in_channels=self.C,
+                         hidden_channels=hidden_channels,
+                         actnorm_scale=actnorm_scale,
+                         flow_permutation=flow_permutation,
+                         flow_coupling=flow_coupling,
+                         acOpt=condAff,
+                         position=position_name,
+                         LU_decomposed=LU_decomposed, opt=opt, idx=k, normOpt=normOpt))
+            self.output_shapes.append(
+                [-1, self.C, H, W])
+
+    def get_condAffSetting(self, opt, opt_get):
+        condAff = opt_get(opt, ['network_G', 'flow', 'condAff']) or None
+        condAff = opt_get(opt, ['network_G', 'flow', 'condFtAffine']) or condAff
+        return condAff
+
+    def arch_split(self, H, W, L, levels, opt, opt_get):
+        correct_splits = opt_get(opt, ['network_G', 'flow', 'split', 'correct_splits'], False)
+        correction = 0 if correct_splits else 1
+        if opt_get(opt, ['network_G', 'flow', 'split', 'enable']) and L < levels - correction:
+            logs_eps = opt_get(opt, ['network_G', 'flow', 'split', 'logs_eps']) or 0
+            consume_ratio = opt_get(opt, ['network_G', 'flow', 'split', 'consume_ratio']) or 0.5
+            position_name = get_position_name(H, self.opt['scale'])
+            position = position_name if opt_get(opt, ['network_G', 'flow', 'split', 'conditional']) else None
+            cond_channels = opt_get(opt, ['network_G', 'flow', 'split', 'cond_channels'])
+            cond_channels = 0 if cond_channels is None else cond_channels
+
+            t = opt_get(opt, ['network_G', 'flow', 'split', 'type'], 'Split2d')
+
+            if t == 'Split2d':
+                split = models.modules.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
+
+    def arch_additionalFlowAffine(self, H, LU_decomposed, W, actnorm_scale, hidden_channels, opt):
+        if 'additionalFlowNoAffine' in opt['network_G']['flow']:
+            n_additionalFlowNoAffine = int(opt['network_G']['flow']['additionalFlowNoAffine'])
+            for _ in range(n_additionalFlowNoAffine):
+                self.layers.append(
+                    FlowStep(in_channels=self.C,
+                             hidden_channels=hidden_channels,
+                             actnorm_scale=actnorm_scale,
+                             flow_permutation='invconv',
+                             flow_coupling='noCoupling',
+                             LU_decomposed=LU_decomposed, opt=opt))
+                self.output_shapes.append(
+                    [-1, self.C, H, W])
+
+    def arch_squeeze(self, H, W):
+        self.C, H, W = self.C * 4, H // 2, W // 2
+        self.layers.append(flow.SqueezeLayer(factor=2))
+        self.output_shapes.append([-1, self.C, H, W])
+        return H, W
+
+    def get_flow_permutation(self, flow_permutation, opt):
+        flow_permutation = opt['network_G']['flow'].get('flow_permutation', 'invconv')
+        return flow_permutation
+
+    def get_affineInCh(self, opt_get):
+        affineInCh = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        affineInCh = (len(affineInCh) + 1) * 64
+        return affineInCh
+
+    def check_image_shape(self):
+        assert self.C == 1 or self.C == 3, ("image_shape should be HWC, like (64, 64, 3)"
+                                            "self.C == 1 or self.C == 3")
+
+    def forward(self, gt=None, rrdbResults=None, z=None, epses=None, logdet=0., reverse=False, eps_std=None,
+                y_onehot=None):
+
+        if reverse:
+            epses_copy = [eps for eps in epses] if isinstance(epses, list) else epses
+
+            sr, logdet = self.decode(rrdbResults, z, eps_std, epses=epses_copy, logdet=logdet, y_onehot=y_onehot)
+            return sr, logdet
+        else:
+            assert gt is not None
+            assert rrdbResults is not None
+            z, logdet = self.encode(gt, rrdbResults, logdet=logdet, epses=epses, y_onehot=y_onehot)
+
+            return z, logdet
+
+    def encode(self, gt, rrdbResults, logdet=0.0, epses=None, y_onehot=None):
+        fl_fea = gt
+        reverse = False
+        level_conditionals = {}
+        bypasses = {}
+
+        L = opt_get(self.opt, ['network_G', 'flow', 'L'])
+
+        for level in range(1, L + 1):
+            bypasses[level] = torch.nn.functional.interpolate(gt, scale_factor=2 ** -level, mode='bilinear', align_corners=False)
+
+        for layer, shape in zip(self.layers, self.output_shapes):
+            size = shape[2]
+            level = int(np.log(160 / size) / np.log(2))
+
+            if level > 0 and level not in level_conditionals.keys():
+                level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+            level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+            if isinstance(layer, FlowStep):
+                fl_fea, logdet = layer(fl_fea, logdet, reverse=reverse, rrdbResults=level_conditionals[level])
+            elif isinstance(layer, Split2d):
+                fl_fea, logdet = self.forward_split2d(epses, fl_fea, layer, logdet, reverse, level_conditionals[level],
+                                                      y_onehot=y_onehot)
+            else:
+                fl_fea, logdet = layer(fl_fea, logdet, reverse=reverse)
+
+        z = fl_fea
+
+        if not isinstance(epses, list):
+            return z, logdet
+
+        epses.append(z)
+        return epses, logdet
+
+    def forward_preFlow(self, fl_fea, logdet, reverse):
+        if hasattr(self, 'preFlow'):
+            for l in self.preFlow:
+                fl_fea, logdet = l(fl_fea, logdet, reverse=reverse)
+        return fl_fea, logdet
+
+    def forward_split2d(self, epses, fl_fea, layer, logdet, reverse, rrdbResults, y_onehot=None):
+        ft = None if layer.position is None else rrdbResults[layer.position]
+        fl_fea, logdet, eps = layer(fl_fea, logdet, reverse=reverse, eps=epses, ft=ft, y_onehot=y_onehot)
+
+        if isinstance(epses, list):
+            epses.append(eps)
+        return fl_fea, logdet
+
+    def decode(self, rrdbResults, z, eps_std=None, epses=None, logdet=0.0, y_onehot=None):
+        z = epses.pop() if isinstance(epses, list) else z
+
+        fl_fea = z
+        # debug.imwrite("fl_fea", fl_fea)
+        bypasses = {}
+        level_conditionals = {}
+        if not opt_get(self.opt, ['network_G', 'flow', 'levelConditional', 'conditional']) == True:
+            for level in range(self.L + 1):
+                level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+        for layer, shape in zip(reversed(self.layers), reversed(self.output_shapes)):
+            size = shape[2]
+            level = int(np.log(160 / size) / np.log(2))
+            # size = fl_fea.shape[2]
+            # level = int(np.log(160 / size) / np.log(2))
+
+            if isinstance(layer, Split2d):
+                fl_fea, logdet = self.forward_split2d_reverse(eps_std, epses, fl_fea, layer,
+                                                              rrdbResults[self.levelToName[level]], logdet=logdet,
+                                                              y_onehot=y_onehot)
+            elif isinstance(layer, FlowStep):
+                fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True, rrdbResults=level_conditionals[level])
+            else:
+                fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True)
+
+        sr = fl_fea
+
+        assert sr.shape[1] == 3
+        return sr, logdet
+
+    def forward_split2d_reverse(self, eps_std, epses, fl_fea, layer, rrdbResults, logdet, y_onehot=None):
+        ft = None if layer.position is None else rrdbResults[layer.position]
+        fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True,
+                               eps=epses.pop() if isinstance(epses, list) else None,
+                               eps_std=eps_std, ft=ft, y_onehot=y_onehot)
+        return fl_fea, logdet
+
+
+def get_position_name(H, scale):
+    downscale_factor = 160 // H
+    position_name = 'fea_up{}'.format(scale / downscale_factor)
+    return position_name

models/SRFlow/code/models/modules/Permutations.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import numpy as np
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from models.modules import thops
+
+
+class InvertibleConv1x1(nn.Module):
+    def __init__(self, num_channels, LU_decomposed=False):
+        super().__init__()
+        w_shape = [num_channels, num_channels]
+        w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32)
+        self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))
+        self.w_shape = w_shape
+        self.LU = LU_decomposed
+
+    def get_weight(self, input, reverse):
+        w_shape = self.w_shape
+        pixels = thops.pixels(input)
+        dlogdet = torch.slogdet(self.weight)[1] * pixels
+        if not reverse:
+            weight = self.weight.view(w_shape[0], w_shape[1], 1, 1)
+        else:
+            weight = torch.inverse(self.weight.double()).float() \
+                .view(w_shape[0], w_shape[1], 1, 1)
+        return weight, dlogdet
+    def forward(self, input, logdet=None, reverse=False):
+        """
+        log-det = log|abs(|W|)| * pixels
+        """
+        weight, dlogdet = self.get_weight(input, reverse)
+        if not reverse:
+            z = F.conv2d(input, weight)
+            if logdet is not None:
+                logdet = logdet + dlogdet
+            return z, logdet
+        else:
+            z = F.conv2d(input, weight)
+            if logdet is not None:
+                logdet = logdet - dlogdet
+            return z, logdet

models/SRFlow/code/models/modules/RRDBNet_arch.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import models.modules.module_util as mutil
+from utils.util import opt_get
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    def __init__(self, nf=64, gc=32, bias=True):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    '''Residual in Residual Dense Block'''
+
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB3 = ResidualDenseBlock_5C(nf, gc)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, opt=None):
+        self.opt = opt
+        super(RRDBNet, self).__init__()
+        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
+        self.scale = scale
+
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+        self.RRDB_trunk = mutil.make_layer(RRDB_block_f, nb)
+        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        #### upsampling
+        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 8:
+            self.upconv3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 16:
+            self.upconv4 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 32:
+            self.upconv5 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+
+        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x, get_steps=False):
+        fea = self.conv_first(x)
+
+        block_idxs = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        block_results = {}
+
+        for idx, m in enumerate(self.RRDB_trunk.children()):
+            fea = m(fea)
+            for b in block_idxs:
+                if b == idx:
+                    block_results["block_{}".format(idx)] = fea
+
+        trunk = self.trunk_conv(fea)
+
+        last_lr_fea = fea + trunk
+
+        fea_up2 = self.upconv1(F.interpolate(last_lr_fea, scale_factor=2, mode='nearest'))
+        fea = self.lrelu(fea_up2)
+
+        fea_up4 = self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest'))
+        fea = self.lrelu(fea_up4)
+
+        fea_up8 = None
+        fea_up16 = None
+        fea_up32 = None
+
+        if self.scale >= 8:
+            fea_up8 = self.upconv3(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up8)
+        if self.scale >= 16:
+            fea_up16 = self.upconv4(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up16)
+        if self.scale >= 32:
+            fea_up32 = self.upconv5(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up32)
+
+        out = self.conv_last(self.lrelu(self.HRconv(fea)))
+
+        results = {'last_lr_fea': last_lr_fea,
+                   'fea_up1': last_lr_fea,
+                   'fea_up2': fea_up2,
+                   'fea_up4': fea_up4,
+                   'fea_up8': fea_up8,
+                   'fea_up16': fea_up16,
+                   'fea_up32': fea_up32,
+                   'out': out}
+
+        fea_up0_en = opt_get(self.opt, ['network_G', 'flow', 'fea_up0']) or False
+        if fea_up0_en:
+            results['fea_up0'] = F.interpolate(last_lr_fea, scale_factor=1/2, mode='bilinear', align_corners=False, recompute_scale_factor=True)
+        fea_upn1_en = opt_get(self.opt, ['network_G', 'flow', 'fea_up-1']) or False
+        if fea_upn1_en:
+            results['fea_up-1'] = F.interpolate(last_lr_fea, scale_factor=1/4, mode='bilinear', align_corners=False, recompute_scale_factor=True)
+
+        if get_steps:
+            for k, v in block_results.items():
+                results[k] = v
+            return results
+        else:
+            return out

models/SRFlow/code/models/modules/SRFlowNet_arch.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from models.modules.RRDBNet_arch import RRDBNet
+from models.modules.FlowUpsamplerNet import FlowUpsamplerNet
+import models.modules.thops as thops
+import models.modules.flow as flow
+from utils.util import opt_get
+
+
+class SRFlowNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, K=None, opt=None, step=None):
+        super(SRFlowNet, self).__init__()
+
+        self.opt = opt
+        self.quant = 255 if opt_get(opt, ['datasets', 'train', 'quant']) is \
+                            None else opt_get(opt, ['datasets', 'train', 'quant'])
+        self.RRDB = RRDBNet(in_nc, out_nc, nf, nb, gc, scale, opt)
+        hidden_channels = opt_get(opt, ['network_G', 'flow', 'hidden_channels'])
+        hidden_channels = hidden_channels or 64
+        self.RRDB_training = True  # Default is true
+
+        train_RRDB_delay = opt_get(self.opt, ['network_G', 'train_RRDB_delay'])
+        set_RRDB_to_train = False
+        if set_RRDB_to_train:
+            self.set_rrdb_training(True)
+
+        self.flowUpsamplerNet = \
+            FlowUpsamplerNet((160, 160, 3), hidden_channels, K,
+                             flow_coupling=opt['network_G']['flow']['coupling'], opt=opt)
+        self.i = 0
+
+    def set_rrdb_training(self, trainable):
+        if self.RRDB_training != trainable:
+            for p in self.RRDB.parameters():
+                p.requires_grad = trainable
+            self.RRDB_training = trainable
+            return True
+        return False
+
+    def forward(self, gt=None, lr=None, z=None, eps_std=None, reverse=False, epses=None, reverse_with_grad=False,
+                lr_enc=None,
+                add_gt_noise=False, step=None, y_label=None):
+        if not reverse:
+            return self.normal_flow(gt, lr, epses=epses, lr_enc=lr_enc, add_gt_noise=add_gt_noise, step=step,
+                                    y_onehot=y_label)
+        else:
+            # assert lr.shape[0] == 1
+            assert lr.shape[1] == 3
+            # assert lr.shape[2] == 20
+            # assert lr.shape[3] == 20
+            # assert z.shape[0] == 1
+            # assert z.shape[1] == 3 * 8 * 8
+            # assert z.shape[2] == 20
+            # assert z.shape[3] == 20
+            if reverse_with_grad:
+                return self.reverse_flow(lr, z, y_onehot=y_label, eps_std=eps_std, epses=epses, lr_enc=lr_enc,
+                                         add_gt_noise=add_gt_noise)
+            else:
+                with torch.no_grad():
+                    return self.reverse_flow(lr, z, y_onehot=y_label, eps_std=eps_std, epses=epses, lr_enc=lr_enc,
+                                             add_gt_noise=add_gt_noise)
+
+    def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noise=True, step=None):
+        if lr_enc is None:
+            lr_enc = self.rrdbPreprocessing(lr)
+
+        logdet = torch.zeros_like(gt[:, 0, 0, 0])
+        pixels = thops.pixels(gt)
+
+        z = gt
+
+        if add_gt_noise:
+            # Setup
+            noiseQuant = opt_get(self.opt, ['network_G', 'flow', 'augmentation', 'noiseQuant'], True)
+            if noiseQuant:
+                z = z + ((torch.rand(z.shape, device=z.device) - 0.5) / self.quant)
+            logdet = logdet + float(-np.log(self.quant) * pixels)
+
+        # Encode
+        epses, logdet = self.flowUpsamplerNet(rrdbResults=lr_enc, gt=z, logdet=logdet, reverse=False, epses=epses,
+                                              y_onehot=y_onehot)
+
+        objective = logdet.clone()
+
+        if isinstance(epses, (list, tuple)):
+            z = epses[-1]
+        else:
+            z = epses
+
+        objective = objective + flow.GaussianDiag.logp(None, None, z)
+
+        nll = (-objective) / float(np.log(2.) * pixels)
+
+        if isinstance(epses, list):
+            return epses, nll, logdet
+        return z, nll, logdet
+
+    def rrdbPreprocessing(self, lr):
+        rrdbResults = self.RRDB(lr, get_steps=True)
+        block_idxs = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        if len(block_idxs) > 0:
+            concat = torch.cat([rrdbResults["block_{}".format(idx)] for idx in block_idxs], dim=1)
+
+            if opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'concat']) or False:
+                keys = ['last_lr_fea', 'fea_up1', 'fea_up2', 'fea_up4']
+                if 'fea_up0' in rrdbResults.keys():
+                    keys.append('fea_up0')
+                if 'fea_up-1' in rrdbResults.keys():
+                    keys.append('fea_up-1')
+                if self.opt['scale'] >= 8:
+                    keys.append('fea_up8')
+                if self.opt['scale'] == 16:
+                    keys.append('fea_up16')
+                for k in keys:
+                    h = rrdbResults[k].shape[2]
+                    w = rrdbResults[k].shape[3]
+                    rrdbResults[k] = torch.cat([rrdbResults[k], F.interpolate(concat, (h, w))], dim=1)
+        return rrdbResults
+
+    def get_score(self, disc_loss_sigma, z):
+        score_real = 0.5 * (1 - 1 / (disc_loss_sigma ** 2)) * thops.sum(z ** 2, dim=[1, 2, 3]) - \
+                     z.shape[1] * z.shape[2] * z.shape[3] * math.log(disc_loss_sigma)
+        return -score_real
+
+    def reverse_flow(self, lr, z, y_onehot, eps_std, epses=None, lr_enc=None, add_gt_noise=True):
+        logdet = torch.zeros_like(lr[:, 0, 0, 0])
+        pixels = thops.pixels(lr) * self.opt['scale'] ** 2
+
+        if add_gt_noise:
+            logdet = logdet - float(-np.log(self.quant) * pixels)
+
+        if lr_enc is None:
+            lr_enc = self.rrdbPreprocessing(lr)
+
+        x, logdet = self.flowUpsamplerNet(rrdbResults=lr_enc, z=z, eps_std=eps_std, reverse=True, epses=epses,
+                                          logdet=logdet)
+
+        return x, logdet

models/SRFlow/code/models/modules/Split.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+
+from models.modules import thops
+from models.modules.FlowStep import FlowStep
+from models.modules.flow import Conv2dZeros, GaussianDiag
+from utils.util import opt_get
+
+
+class Split2d(nn.Module):
+    def __init__(self, num_channels, logs_eps=0, cond_channels=0, position=None, consume_ratio=0.5, opt=None):
+        super().__init__()
+
+        self.num_channels_consume = int(round(num_channels * consume_ratio))
+        self.num_channels_pass = num_channels - self.num_channels_consume
+
+        self.conv = Conv2dZeros(in_channels=self.num_channels_pass + cond_channels,
+                                out_channels=self.num_channels_consume * 2)
+        self.logs_eps = logs_eps
+        self.position = position
+        self.opt = opt
+
+    def split2d_prior(self, z, ft):
+        if ft is not None:
+            z = torch.cat([z, ft], dim=1)
+        h = self.conv(z)
+        return thops.split_feature(h, "cross")
+
+    def exp_eps(self, logs):
+        return torch.exp(logs) + self.logs_eps
+
+    def forward(self, input, logdet=0., reverse=False, eps_std=None, eps=None, ft=None, y_onehot=None):
+        if not reverse:
+            # self.input = input
+            z1, z2 = self.split_ratio(input)
+            mean, logs = self.split2d_prior(z1, ft)
+            
+            eps = (z2 - mean) / self.exp_eps(logs)
+
+            logdet = logdet + self.get_logdet(logs, mean, z2)
+
+            # print(logs.shape, mean.shape, z2.shape)
+            # self.eps = eps
+            # print('split, enc eps:', eps)
+            return z1, logdet, eps
+        else:
+            z1 = input
+            mean, logs = self.split2d_prior(z1, ft)
+
+            if eps is None:
+                #print("WARNING: eps is None, generating eps untested functionality!")
+                eps = GaussianDiag.sample_eps(mean.shape, eps_std)
+
+            eps = eps.to(mean.device)
+            z2 = mean + self.exp_eps(logs) * eps
+
+            z = thops.cat_feature(z1, z2)
+            logdet = logdet - self.get_logdet(logs, mean, z2)
+
+            return z, logdet
+            # return z, logdet, eps
+
+    def get_logdet(self, logs, mean, z2):
+        logdet_diff = GaussianDiag.logp(mean, logs, z2)
+        # print("Split2D: logdet diff", logdet_diff.item())
+        return logdet_diff
+
+    def split_ratio(self, input):
+        z1, z2 = input[:, :self.num_channels_pass, ...], input[:, self.num_channels_pass:, ...]
+        return z1, z2

models/SRFlow/code/models/modules/flow.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from models.modules.FlowActNorms import ActNorm2d
+from . import thops
+
+
+class Conv2d(nn.Conv2d):
+    pad_dict = {
+        "same": lambda kernel, stride: [((k - 1) * s + 1) // 2 for k, s in zip(kernel, stride)],
+        "valid": lambda kernel, stride: [0 for _ in kernel]
+    }
+
+    @staticmethod
+    def get_padding(padding, kernel_size, stride):
+        # make paddding
+        if isinstance(padding, str):
+            if isinstance(kernel_size, int):
+                kernel_size = [kernel_size, kernel_size]
+            if isinstance(stride, int):
+                stride = [stride, stride]
+            padding = padding.lower()
+            try:
+                padding = Conv2d.pad_dict[padding](kernel_size, stride)
+            except KeyError:
+                raise ValueError("{} is not supported".format(padding))
+        return padding
+
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=[3, 3], stride=[1, 1],
+                 padding="same", do_actnorm=True, weight_std=0.05):
+        padding = Conv2d.get_padding(padding, kernel_size, stride)
+        super().__init__(in_channels, out_channels, kernel_size, stride,
+                         padding, bias=(not do_actnorm))
+        # init weight with std
+        self.weight.data.normal_(mean=0.0, std=weight_std)
+        if not do_actnorm:
+            self.bias.data.zero_()
+        else:
+            self.actnorm = ActNorm2d(out_channels)
+        self.do_actnorm = do_actnorm
+
+    def forward(self, input):
+        x = super().forward(input)
+        if self.do_actnorm:
+            x, _ = self.actnorm(x)
+        return x
+
+
+class Conv2dZeros(nn.Conv2d):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=[3, 3], stride=[1, 1],
+                 padding="same", logscale_factor=3):
+        padding = Conv2d.get_padding(padding, kernel_size, stride)
+        super().__init__(in_channels, out_channels, kernel_size, stride, padding)
+        # logscale_factor
+        self.logscale_factor = logscale_factor
+        self.register_parameter("logs", nn.Parameter(torch.zeros(out_channels, 1, 1)))
+        # init
+        self.weight.data.zero_()
+        self.bias.data.zero_()
+
+    def forward(self, input):
+        output = super().forward(input)
+        return output * torch.exp(self.logs * self.logscale_factor)
+
+
+class GaussianDiag:
+    Log2PI = float(np.log(2 * np.pi))
+
+    @staticmethod
+    def likelihood(mean, logs, x):
+        """
+        lnL = -1/2 * { ln|Var| + ((X - Mu)^T)(Var^-1)(X - Mu) + kln(2*PI) }
+              k = 1 (Independent)
+              Var = logs ** 2
+        """
+        if mean is None and logs is None:
+            return -0.5 * (x ** 2 + GaussianDiag.Log2PI)
+        else:
+            return -0.5 * (logs * 2. + ((x - mean) ** 2) / torch.exp(logs * 2.) + GaussianDiag.Log2PI)
+
+    @staticmethod
+    def logp(mean, logs, x):
+        likelihood = GaussianDiag.likelihood(mean, logs, x)
+        return thops.sum(likelihood, dim=[1, 2, 3])
+
+    @staticmethod
+    def sample(mean, logs, eps_std=None):
+        eps_std = eps_std or 1
+        eps = torch.normal(mean=torch.zeros_like(mean),
+                           std=torch.ones_like(logs) * eps_std)
+        return mean + torch.exp(logs) * eps
+
+    @staticmethod
+    def sample_eps(shape, eps_std, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+        eps = torch.normal(mean=torch.zeros(shape),
+                           std=torch.ones(shape) * eps_std)
+        return eps
+
+
+def squeeze2d(input, factor=2):
+    assert factor >= 1 and isinstance(factor, int)
+    if factor == 1:
+        return input
+    size = input.size()
+    B = size[0]
+    C = size[1]
+    H = size[2]
+    W = size[3]
+    assert H % factor == 0 and W % factor == 0, "{}".format((H, W, factor))
+    x = input.view(B, C, H // factor, factor, W // factor, factor)
+    x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
+    x = x.view(B, C * factor * factor, H // factor, W // factor)
+    return x
+
+
+def unsqueeze2d(input, factor=2):
+    assert factor >= 1 and isinstance(factor, int)
+    factor2 = factor ** 2
+    if factor == 1:
+        return input
+    size = input.size()
+    B = size[0]
+    C = size[1]
+    H = size[2]
+    W = size[3]
+    assert C % (factor2) == 0, "{}".format(C)
+    x = input.view(B, C // factor2, factor, factor, H, W)
+    x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
+    x = x.view(B, C // (factor2), H * factor, W * factor)
+    return x
+
+
+class SqueezeLayer(nn.Module):
+    def __init__(self, factor):
+        super().__init__()
+        self.factor = factor
+
+    def forward(self, input, logdet=None, reverse=False):
+        if not reverse:
+            output = squeeze2d(input, self.factor)  # Squeeze in forward
+            return output, logdet
+        else:
+            output = unsqueeze2d(input, self.factor)
+            return output, logdet

models/SRFlow/code/models/modules/glow_arch.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch.nn as nn
+
+
+def f_conv2d_bias(in_channels, out_channels):
+    def padding_same(kernel, stride):
+        return [((k - 1) * s + 1) // 2 for k, s in zip(kernel, stride)]
+
+    padding = padding_same([3, 3], [1, 1])
+    assert padding == [1, 1], padding
+    return nn.Sequential(
+        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=[3, 3], stride=1, padding=1,
+                  bias=True))

models/SRFlow/code/models/modules/loss.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+import torch.nn as nn
+
+
+class CharbonnierLoss(nn.Module):
+    """Charbonnier Loss (L1)"""
+
+    def __init__(self, eps=1e-6):
+        super(CharbonnierLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, x, y):
+        diff = x - y
+        loss = torch.sum(torch.sqrt(diff * diff + self.eps))
+        return loss
+
+
+# Define GAN loss: [vanilla | lsgan | wgan-gp]
+class GANLoss(nn.Module):
+    def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0):
+        super(GANLoss, self).__init__()
+        self.gan_type = gan_type.lower()
+        self.real_label_val = real_label_val
+        self.fake_label_val = fake_label_val
+
+        if self.gan_type == 'gan' or self.gan_type == 'ragan':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif self.gan_type == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif self.gan_type == 'wgan-gp':
+
+            def wgan_loss(input, target):
+                # target is boolean
+                return -1 * input.mean() if target else input.mean()
+
+            self.loss = wgan_loss
+        else:
+            raise NotImplementedError('GAN type [{:s}] is not found'.format(self.gan_type))
+
+    def get_target_label(self, input, target_is_real):
+        if self.gan_type == 'wgan-gp':
+            return target_is_real
+        if target_is_real:
+            return torch.empty_like(input).fill_(self.real_label_val)
+        else:
+            return torch.empty_like(input).fill_(self.fake_label_val)
+
+    def forward(self, input, target_is_real):
+        target_label = self.get_target_label(input, target_is_real)
+        loss = self.loss(input, target_label)
+        return loss
+
+
+class GradientPenaltyLoss(nn.Module):
+    def __init__(self, device=torch.device('cpu')):
+        super(GradientPenaltyLoss, self).__init__()
+        self.register_buffer('grad_outputs', torch.Tensor())
+        self.grad_outputs = self.grad_outputs.to(device)
+
+    def get_grad_outputs(self, input):
+        if self.grad_outputs.size() != input.size():
+            self.grad_outputs.resize_(input.size()).fill_(1.0)
+        return self.grad_outputs
+
+    def forward(self, interp, interp_crit):
+        grad_outputs = self.get_grad_outputs(interp_crit)
+        grad_interp = torch.autograd.grad(outputs=interp_crit, inputs=interp,
+                                          grad_outputs=grad_outputs, create_graph=True,
+                                          retain_graph=True, only_inputs=True)[0]
+        grad_interp = grad_interp.view(grad_interp.size(0), -1)
+        grad_interp_norm = grad_interp.norm(2, dim=1)
+
+        loss = ((grad_interp_norm - 1)**2).mean()
+        return loss

models/SRFlow/code/models/modules/module_util.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+
+
+def initialize_weights(net_l, scale=1):
+    if not isinstance(net_l, list):
+        net_l = [net_l]
+    for net in net_l:
+        for m in net.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+                m.weight.data *= scale  # for residual block
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias.data, 0.0)
+
+
+def make_layer(block, n_layers):
+    layers = []
+    for _ in range(n_layers):
+        layers.append(block())
+    return nn.Sequential(*layers)
+
+
+class ResidualBlock_noBN(nn.Module):
+    '''Residual block w/o BN
+    ---Conv-ReLU-Conv-+-
+     |________________|
+    '''
+
+    def __init__(self, nf=64):
+        super(ResidualBlock_noBN, self).__init__()
+        self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+
+        # initialization
+        initialize_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = F.relu(self.conv1(x), inplace=True)
+        out = self.conv2(out)
+        return identity + out
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros'):
+    """Warp an image or feature map with optical flow
+    Args:
+        x (Tensor): size (N, C, H, W)
+        flow (Tensor): size (N, H, W, 2), normal value
+        interp_mode (str): 'nearest' or 'bilinear'
+        padding_mode (str): 'zeros' or 'border' or 'reflection'
+
+    Returns:
+        Tensor: warped image or feature map
+    """
+    assert x.size()[-2:] == flow.size()[1:3]
+    B, C, H, W = x.size()
+    # mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W))
+    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2
+    grid.requires_grad = False
+    grid = grid.type_as(x)
+    vgrid = grid + flow
+    # scale grid to [-1,1]
+    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(W - 1, 1) - 1.0
+    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(H - 1, 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+    output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode)
+    return output

models/SRFlow/code/models/modules/thops.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+
+
+def sum(tensor, dim=None, keepdim=False):
+    if dim is None:
+        # sum up all dim
+        return torch.sum(tensor)
+    else:
+        if isinstance(dim, int):
+            dim = [dim]
+        dim = sorted(dim)
+        for d in dim:
+            tensor = tensor.sum(dim=d, keepdim=True)
+        if not keepdim:
+            for i, d in enumerate(dim):
+                tensor.squeeze_(d-i)
+        return tensor
+
+
+def mean(tensor, dim=None, keepdim=False):
+    if dim is None:
+        # mean all dim
+        return torch.mean(tensor)
+    else:
+        if isinstance(dim, int):
+            dim = [dim]
+        dim = sorted(dim)
+        for d in dim:
+            tensor = tensor.mean(dim=d, keepdim=True)
+        if not keepdim:
+            for i, d in enumerate(dim):
+                tensor.squeeze_(d-i)
+        return tensor
+
+
+def split_feature(tensor, type="split"):
+    """
+    type = ["split", "cross"]
+    """
+    C = tensor.size(1)
+    if type == "split":
+        return tensor[:, :C // 2, ...], tensor[:, C // 2:, ...]
+    elif type == "cross":
+        return tensor[:, 0::2, ...], tensor[:, 1::2, ...]
+
+
+def cat_feature(tensor_a, tensor_b):
+    return torch.cat((tensor_a, tensor_b), dim=1)
+
+
+def pixels(tensor):
+    return int(tensor.size(2) * tensor.size(3))

models/SRFlow/code/models/networks.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import importlib
+
+import torch
+import logging
+import models.modules.RRDBNet_arch as RRDBNet_arch
+
+logger = logging.getLogger('base')
+
+
+def find_model_using_name(model_name):
+    model_filename = "models.modules." + model_name + "_arch"
+    modellib = importlib.import_module(model_filename)
+
+    model = None
+    target_model_name = model_name.replace('_Net', '')
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower():
+            model = cls
+
+    if model is None:
+        print(
+            "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s." % (
+                model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+####################
+# define network
+####################
+#### Generator
+def define_G(opt):
+    opt_net = opt['network_G']
+    which_model = opt_net['which_model_G']
+
+    if which_model == 'RRDBNet':
+        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=opt['scale'], opt=opt)
+    elif which_model == 'EDSRNet':
+        Arch = find_model_using_name(which_model)
+        netG = Arch(scale=opt['scale'])
+    elif which_model == 'rankSRGAN':
+        Arch = find_model_using_name(which_model)
+        netG = Arch(upscale=opt['scale'])
+    # elif which_model == 'sft_arch':  # SFT-GAN
+    #     netG = sft_arch.SFT_Net()
+    else:
+        raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))
+    return netG
+
+
+def define_Flow(opt, step):
+    opt_net = opt['network_G']
+    which_model = opt_net['which_model_G']
+
+    Arch = find_model_using_name(which_model)
+    netG = Arch(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
+                nf=opt_net['nf'], nb=opt_net['nb'], scale=opt['scale'], K=opt_net['flow']['K'], opt=opt, step=step)
+
+    return netG
+
+
+#### Discriminator
+def define_D(opt):
+    opt_net = opt['network_D']
+    which_model = opt_net['which_model_D']
+
+    if which_model == 'discriminator_vgg_128':
+        hidden_units = opt_net.get('hidden_units', 8192)
+        netD = SRGAN_arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], nf=opt_net['nf'], hidden_units=hidden_units)
+    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.get('gpu_ids', None)
+    device = torch.device('cuda' if gpu_ids else 'cpu')
+    # PyTorch pretrained_models VGG19-54, before ReLU.
+    if use_bn:
+        feature_layer = 49
+    else:
+        feature_layer = 34
+    netF = SRGAN_arch.VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn,
+                                          use_input_norm=True, device=device)
+    netF.eval()  # No need to train
+    return netF

models/SRFlow/code/options/options.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+import os.path as osp
+import logging
+import yaml
+from utils.util import OrderedYaml
+
+Loader, Dumper = OrderedYaml()
+
+
+def parse(opt_path, is_train=True):
+    with open(opt_path, mode='r') as f:
+        opt = yaml.load(f, Loader=Loader)
+    # export CUDA_VISIBLE_DEVICES
+    gpu_list = ','.join(str(x) for x in opt.get('gpu_ids', []))
+    # os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+    # print('export CUDA_VISIBLE_DEVICES=' + gpu_list)
+    opt['is_train'] = is_train
+    if opt['distortion'] == 'sr':
+        scale = opt['scale']
+
+    # datasets
+    for phase, dataset in opt['datasets'].items():
+        phase = phase.split('_')[0]
+        dataset['phase'] = phase
+        if opt['distortion'] == 'sr':
+            dataset['scale'] = scale
+        is_lmdb = False
+        if dataset.get('dataroot_GT', None) is not None:
+            dataset['dataroot_GT'] = osp.expanduser(dataset['dataroot_GT'])
+            if dataset['dataroot_GT'].endswith('lmdb'):
+                is_lmdb = True
+        if dataset.get('dataroot_LQ', None) is not None:
+            dataset['dataroot_LQ'] = osp.expanduser(dataset['dataroot_LQ'])
+            if dataset['dataroot_LQ'].endswith('lmdb'):
+                is_lmdb = True
+        dataset['data_type'] = 'lmdb' if is_lmdb else 'img'
+        if dataset['mode'].endswith('mc'):  # for memcached
+            dataset['data_type'] = 'mc'
+            dataset['mode'] = dataset['mode'].replace('_mc', '')
+
+    # path
+    for key, path in opt['path'].items():
+        if path and key in opt['path'] and key != 'strict_load':
+            opt['path'][key] = osp.expanduser(path)
+    opt['path']['root'] = '/kaggle/working/'
+    if is_train:
+        experiments_root = osp.join(opt['path']['root'], 'experiments', opt['name'])
+        opt['path']['experiments_root'] = experiments_root
+        opt['path']['models'] = osp.join(experiments_root, 'models')
+        opt['path']['training_state'] = osp.join(experiments_root, 'training_state')
+        opt['path']['log'] = experiments_root
+        opt['path']['val_images'] = osp.join(experiments_root, 'val_images')
+
+        # change some options for debug mode
+        if 'debug' in opt['name']:
+            opt['train']['val_freq'] = 8
+            opt['logger']['print_freq'] = 1
+            opt['logger']['save_checkpoint_freq'] = 8
+    else:  # test
+        if not opt['path'].get('results_root', None):
+            results_root = osp.join(opt['path']['root'], 'results', opt['name'])
+            opt['path']['results_root'] = results_root
+        opt['path']['log'] = opt['path']['results_root']
+
+    # network
+    if opt['distortion'] == 'sr':
+        opt['network_G']['scale'] = scale
+
+    # relative learning rate
+    if 'train' in opt:
+        niter = opt['train']['niter']
+        if 'T_period_rel' in opt['train']:
+            opt['train']['T_period'] = [int(x * niter) for x in opt['train']['T_period_rel']]
+        if 'restarts_rel' in opt['train']:
+            opt['train']['restarts'] = [int(x * niter) for x in opt['train']['restarts_rel']]
+        if 'lr_steps_rel' in opt['train']:
+            opt['train']['lr_steps'] = [int(x * niter) for x in opt['train']['lr_steps_rel']]
+        if 'lr_steps_inverse_rel' in opt['train']:
+            opt['train']['lr_steps_inverse'] = [int(x * niter) for x in opt['train']['lr_steps_inverse_rel']]
+        print(opt['train'])
+
+    return opt
+
+
+def dict2str(opt, indent_l=1):
+    '''dict to string for logger'''
+    msg = ''
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_l * 2) + k + ':[\n'
+            msg += dict2str(v, indent_l + 1)
+            msg += ' ' * (indent_l * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+
+class NoneDict(dict):
+    def __missing__(self, key):
+        return None
+
+
+# convert to NoneDict, which return None for missing key.
+def dict_to_nonedict(opt):
+    if isinstance(opt, dict):
+        new_opt = dict()
+        for key, sub_opt in opt.items():
+            new_opt[key] = dict_to_nonedict(sub_opt)
+        return NoneDict(**new_opt)
+    elif isinstance(opt, list):
+        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+    else:
+        return opt
+
+
+def check_resume(opt, resume_iter):
+    '''Check resume states and pretrain_model paths'''
+    logger = logging.getLogger('base')
+    if opt['path']['resume_state']:
+        if opt['path'].get('pretrain_model_G', None) is not None or opt['path'].get(
+                'pretrain_model_D', None) is not None:
+            logger.warning('pretrain_model path will be ignored when resuming training.')
+
+        opt['path']['pretrain_model_G'] = osp.join(opt['path']['models'],
+                                                   '{}_G.pth'.format(resume_iter))
+        logger.info('Set [pretrain_model_G] to ' + opt['path']['pretrain_model_G'])
+        if 'gan' in opt['model']:
+            opt['path']['pretrain_model_D'] = osp.join(opt['path']['models'],
+                                                       '{}_D.pth'.format(resume_iter))
+            logger.info('Set [pretrain_model_D] to ' + opt['path']['pretrain_model_D'])

models/SRFlow/code/prepare_data.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import sys
+
+import numpy as np
+import random
+import imageio
+import pickle
+
+from natsort import natsort
+from tqdm import tqdm
+
+def get_img_paths(dir_path, wildcard='*.png'):
+    return natsort.natsorted(glob.glob(dir_path + '/' + wildcard))
+
+def create_all_dirs(path):
+    if "." in path.split("/")[-1]:
+        dirs = os.path.dirname(path)
+    else:
+        dirs = path
+    os.makedirs(dirs, exist_ok=True)
+
+def to_pklv4(obj, path, vebose=False):
+    create_all_dirs(path)
+    with open(path, 'wb') as f:
+        pickle.dump(obj, f, protocol=4)
+    if vebose:
+        print("Wrote {}".format(path))
+
+
+from imresize import imresize
+
+def random_crop(img, size):
+    h, w, c = img.shape
+
+    h_start = np.random.randint(0, h - size)
+    h_end = h_start + size
+
+    w_start = np.random.randint(0, w - size)
+    w_end = w_start + size
+
+    return img[h_start:h_end, w_start:w_end]
+
+
+def imread(img_path):
+    img = imageio.imread(img_path)
+    if len(img.shape) == 2:
+        img = np.stack([img, ] * 3, axis=2)
+    return img
+
+
+def to_pklv4_1pct(obj, path, vebose):
+    n = int(round(len(obj) * 0.01))
+    path = path.replace(".", "_1pct.")
+    to_pklv4(obj[:n], path, vebose=True)
+
+
+def main(dir_path):
+    hrs = []
+    lqs = []
+
+    img_paths = get_img_paths(dir_path)
+    for img_path in tqdm(img_paths):
+        img = imread(img_path)
+
+        for i in range(47):
+            crop = random_crop(img, 256)
+            cropX4 = imresize(crop, scalar_scale=0.25)
+            hrs.append(crop)
+            lqs.append(cropX4)
+
+    shuffle_combined(hrs, lqs)
+
+    hrs_path = get_hrs_path(dir_path)
+    to_pklv4(hrs, hrs_path, vebose=True)
+
+    lqs_path = get_lqs_path(dir_path)
+    to_pklv4(lqs, lqs_path, vebose=True)
+
+
+def get_hrs_path(dir_path):
+    base_dir = '/kaggle/working/'
+    name = os.path.basename(dir_path)
+    hrs_path = os.path.join(base_dir, 'pkls', name + '.pklv4')
+    return hrs_path
+
+
+def get_lqs_path(dir_path):
+    base_dir = '/kaggle/working/'
+    name = os.path.basename(dir_path)
+    hrs_path = os.path.join(base_dir, 'pkls', name + '_X4.pklv4')
+    return hrs_path
+
+
+def shuffle_combined(hrs, lqs):
+    combined = list(zip(hrs, lqs))
+    random.shuffle(combined)
+    hrs[:], lqs[:] = zip(*combined)
+
+
+if __name__ == "__main__":
+    dir_path = sys.argv[1]
+    assert os.path.isdir(dir_path)
+    main(dir_path)

models/SRFlow/code/test.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+
+import glob 
+import sys
+from collections import OrderedDict
+
+from natsort import natsort
+
+import options.options as option
+from Measure import Measure, psnr
+from imresize import imresize
+from models import create_model
+import torch
+from utils.util import opt_get
+import numpy as np
+import pandas as pd
+import os
+import cv2
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+def load_model(conf_path):
+    opt = option.parse(conf_path, is_train=False)
+    opt['gpu_ids'] = None
+    opt = option.dict_to_nonedict(opt)
+    model = create_model(opt)
+
+    model_path = opt_get(opt, ['model_path'], None)
+    model.load_network(load_path=model_path, network=model.netG)
+    return model, opt
+
+
+def predict(model, lr):
+    model.feed_data({"LQ": t(lr)}, need_GT=False)
+    model.test()
+    visuals = model.get_current_visuals(need_GT=False)
+    return visuals.get('rlt', visuals.get("SR"))
+
+
+def t(array): return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255
+
+
+def rgb(t): return (
+        np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype(
+    np.uint8)
+
+
+def imread(path):
+    return cv2.imread(path)[:, :, [2, 1, 0]]
+
+
+def imwrite(path, img):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    cv2.imwrite(path, img[:, :, [2, 1, 0]])
+
+
+def imCropCenter(img, size):
+    h, w, c = img.shape
+
+    h_start = max(h // 2 - size // 2, 0)
+    h_end = min(h_start + size, h)
+
+    w_start = max(w // 2 - size // 2, 0)
+    w_end = min(w_start + size, w)
+
+    return img[h_start:h_end, w_start:w_end]
+
+
+def impad(img, top=0, bottom=0, left=0, right=0, color=255):
+    return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect')
+
+
+def main():
+    conf_path = sys.argv[1]
+    conf = conf_path.split('/')[-1].replace('.yml', '')
+    model, opt = load_model(conf_path)
+
+    data_dir = opt['dataroot']
+
+    # this_dir = os.path.dirname(os.path.realpath(__file__))
+    test_dir = os.path.join('/kaggle/working/', 'results', conf)
+    print(f"Out dir: {test_dir}")
+
+    measure = Measure(use_gpu=False)
+
+    fname = f'measure_full.csv'
+    fname_tmp = fname + "_"
+    path_out_measures = os.path.join(test_dir, fname_tmp)
+    path_out_measures_final = os.path.join(test_dir, fname)
+
+    if os.path.isfile(path_out_measures_final):
+        df = pd.read_csv(path_out_measures_final)
+    elif os.path.isfile(path_out_measures):
+        df = pd.read_csv(path_out_measures)
+    else:
+        df = None
+
+    scale = opt['scale']
+
+    pad_factor = 2
+
+    data_sets = [
+        'Set5',
+        'Set14',
+        'Urban100',
+        'BSD100'
+        ]
+    
+    final_df = pd.DataFrame()
+
+    for data_set in data_sets:
+        lr_paths = fiFindByWildcard(os.path.join(data_dir, data_set, '*LR.png'))
+        hr_paths = fiFindByWildcard(os.path.join(data_dir, data_set, '*HR.png'))
+        
+        df = pd.DataFrame(columns=['conf', 'heat', 'data_set', 'name', 'PSNR', 'SSIM', 'LPIPS', 'LRC PSNR'])
+
+        for lr_path, hr_path, idx_test in zip(lr_paths, hr_paths, range(len(lr_paths))):
+            with torch.no_grad(), torch.cuda.amp.autocast():
+                lr = imread(lr_path)
+                hr = imread(hr_path)
+
+                # Pad image to be % 2
+                h, w, c = lr.shape
+                lq_orig = lr.copy()
+                lr = impad(lr, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
+                        right=int(np.ceil(w / pad_factor) * pad_factor - w))
+
+                lr_t = t(lr)
+
+                heat = opt['heat']
+
+                if df is not None and len(df[(df['heat'] == heat) & (df['name'] == idx_test)]) == 1:
+                    continue
+
+                sr_t = model.get_sr(lq=lr_t, heat=heat)
+
+                sr = rgb(torch.clamp(sr_t, 0, 1))
+                sr = sr[:h * scale, :w * scale]
+
+                path_out_sr = os.path.join(test_dir, data_set, "{:0.2f}".format(heat).replace('.', ''), "{:06d}.png".format(idx_test))
+                imwrite(path_out_sr, sr)
+
+                meas = OrderedDict(conf=conf, heat=heat, data_set=data_set, name=idx_test)
+                meas['PSNR'], meas['SSIM'], meas['LPIPS'] = measure.measure(sr, hr)
+
+                lr_reconstruct_rgb = imresize(sr, 1 / opt['scale'])
+                meas['LRC PSNR'] = psnr(lq_orig, lr_reconstruct_rgb)
+
+                str_out = format_measurements(meas)
+                print(str_out)
+
+                df = df._append(pd.DataFrame([meas]), ignore_index=True)
+            
+        final_df = pd.concat([final_df, df])
+
+    final_df.to_csv(path_out_measures, index=False)
+    os.rename(path_out_measures, path_out_measures_final)
+
+    # str_out = format_measurements(df.mean())
+    # print(f"Results in: {path_out_measures_final}")
+    # print('Mean: ' + str_out)
+
+
+def format_measurements(meas):
+    s_out = []
+    for k, v in meas.items():
+        v = f"{v:0.2f}" if isinstance(v, float) else v
+        s_out.append(f"{k}: {v}")
+    str_out = ", ".join(s_out)
+    return str_out
+
+
+if __name__ == "__main__":
+    main()

models/SRFlow/code/train.py

@@ -0,0 +1,328 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+from os.path import basename
+import math
+import argparse
+import random
+import logging
+import cv2
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+import options.options as option
+from utils import util
+from data import create_dataloader, create_dataset
+from models import create_model
+from utils.timer import Timer, TickTock
+from utils.util import get_resume_paths
+
+import wandb
+
+def getEnv(name): import os; return True if name in os.environ.keys() else False
+
+
+def init_dist(backend='nccl', **kwargs):
+    ''' initialization for distributed training'''
+    # if mp.get_start_method(allow_none=True) is None:
+    if mp.get_start_method(allow_none=True) != 'spawn':
+        mp.set_start_method('spawn')
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_deviceDistIterSampler(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def main():
+    wandb.init(project='srflow')
+    #### options
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, help='Path to option YMAL file.')
+    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
+                        help='job launcher')
+    parser.add_argument('--local_rank', type=int, default=0)
+    args = parser.parse_args()
+    opt = option.parse(args.opt, is_train=True)
+
+    #### distributed training settings
+    opt['dist'] = False
+    rank = -1
+    print('Disabled distributed training.')
+
+    #### loading resume state if exists
+    if opt['path'].get('resume_state', None):
+        resume_state_path, _ = get_resume_paths(opt)
+
+        # distributed resuming: all load into default GPU
+        if resume_state_path is None:
+            resume_state = None
+        else:
+            device_id = torch.cuda.current_device()
+            resume_state = torch.load(resume_state_path,
+                                      map_location=lambda storage, loc: storage.cuda(device_id))
+            option.check_resume(opt, resume_state['iter'])  # check resume options
+    else:
+        resume_state = None
+
+    #### mkdir and loggers
+    if rank <= 0:  # normal training (rank -1) OR distributed training (rank 0)
+        if resume_state is None:
+            util.mkdir_and_rename(
+                opt['path']['experiments_root'])  # rename experiment folder if exists
+            util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
+                         and 'pretrain_model' not in key and 'resume' not in key))
+
+        # config loggers. Before it, the log will not work
+        util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        logger = logging.getLogger('base')
+        logger.info(option.dict2str(opt))
+
+        # tensorboard logger
+        if opt.get('use_tb_logger', False) and 'debug' not in opt['name']:
+            version = float(torch.__version__[0:3])
+            if version >= 1.1:  # PyTorch 1.1
+                from torch.utils.tensorboard import SummaryWriter
+            else:
+                logger.info(
+                    'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
+                from tensorboardX import SummaryWriter
+            conf_name = basename(args.opt).replace(".yml", "")
+            exp_dir = opt['path']['experiments_root']
+            log_dir_train = os.path.join(exp_dir, 'tb', conf_name, 'train')
+            log_dir_valid = os.path.join(exp_dir, 'tb', conf_name, 'valid')
+            tb_logger_train = SummaryWriter(log_dir=log_dir_train)
+            tb_logger_valid = SummaryWriter(log_dir=log_dir_valid)
+    else:
+        util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
+        logger = logging.getLogger('base')
+
+    # convert to NoneDict, which returns None for missing keys
+    opt = option.dict_to_nonedict(opt)
+
+    #### random seed
+    seed = opt['train']['manual_seed']
+    if seed is None:
+        seed = random.randint(1, 10000)
+    if rank <= 0:
+        logger.info('Random seed: {}'.format(seed))
+    util.set_random_seed(seed)
+
+    torch.backends.cudnn.benchmark = True
+    # torch.backends.cudnn.deterministic = True
+
+    #### create train and val dataloader
+    dataset_ratio = 200  # enlarge the size of each epoch
+    for phase, dataset_opt in opt['datasets'].items():
+        if phase == 'train':
+            full_dataset = create_dataset(dataset_opt)
+            print('Dataset created')
+            train_len = int(len(full_dataset) * 0.95)
+            val_len = len(full_dataset) - train_len
+            train_set, val_set = torch.utils.data.random_split(full_dataset, [train_len, val_len])
+            train_size = int(math.ceil(train_len / dataset_opt['batch_size']))
+            total_iters = int(opt['train']['niter'])
+            total_epochs = int(math.ceil(total_iters / train_size))
+            train_sampler = None
+            train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
+            if rank <= 0:
+                logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
+                    len(train_set), train_size))
+                logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
+                    total_epochs, total_iters))
+            val_loader = torch.utils.data.DataLoader(val_set, batch_size=1, shuffle=False, num_workers=1,
+                                           pin_memory=True)
+        elif phase == 'val':
+            continue
+        else:
+            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
+    assert train_loader is not None
+
+    #### create model
+    current_step = 0 if resume_state is None else resume_state['iter']
+    model = create_model(opt, current_step)
+
+    #### resume training
+    if resume_state:
+        logger.info('Resuming training from epoch: {}, iter: {}.'.format(
+            resume_state['epoch'], resume_state['iter']))
+
+        start_epoch = resume_state['epoch']
+        current_step = resume_state['iter']
+        model.resume_training(resume_state)  # handle optimizers and schedulers
+    else:
+        current_step = 0
+        start_epoch = 0
+
+    #### training
+    timer = Timer()
+    logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
+    timerData = TickTock()
+
+    for epoch in range(start_epoch, total_epochs + 1):
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+
+        timerData.tick()
+        for _, train_data in enumerate(train_loader):
+            timerData.tock()
+            current_step += 1
+            if current_step > total_iters:
+                break
+
+            #### training
+            model.feed_data(train_data)
+
+            #### update learning rate
+            model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])
+
+            try:
+                nll = model.optimize_parameters(current_step)
+            except RuntimeError as e:
+                print("Skipping ERROR caught in nll = model.optimize_parameters(current_step): ")
+                print(e)
+
+            if nll is None:
+                nll = 0
+
+            wandb.log({"loss": nll})
+            #### log
+            def eta(t_iter):
+                return (t_iter * (opt['train']['niter'] - current_step)) / 3600
+
+            if current_step % opt['logger']['print_freq'] == 0 \
+                    or current_step - (resume_state['iter'] if resume_state else 0) < 25:
+                avg_time = timer.get_average_and_reset()
+                avg_data_time = timerData.get_average_and_reset()
+                message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}, t:{:.2e}, td:{:.2e}, eta:{:.2e}, nll:{:.3e}> '.format(
+                    epoch, current_step, model.get_current_learning_rate(), avg_time, avg_data_time,
+                    eta(avg_time), nll)
+                print(message)
+            timer.tick()
+            # Reduce number of logs
+            if current_step % 5 == 0:
+                tb_logger_train.add_scalar('loss/nll', nll, current_step)
+                tb_logger_train.add_scalar('lr/base', model.get_current_learning_rate(), current_step)
+                tb_logger_train.add_scalar('time/iteration', timer.get_last_iteration(), current_step)
+                tb_logger_train.add_scalar('time/data', timerData.get_last_iteration(), current_step)
+                tb_logger_train.add_scalar('time/eta', eta(timer.get_last_iteration()), current_step)
+                for k, v in model.get_current_log().items():
+                    tb_logger_train.add_scalar(k, v, current_step)
+
+            # validation
+            if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
+                avg_psnr = 0.0
+                idx = 0
+                nlls = []
+                for val_data in val_loader:
+                    idx += 1
+                    img_name = os.path.splitext(os.path.basename(val_data['LQ_path'][0]))[0]
+                    img_dir = os.path.join(opt['path']['val_images'], img_name)
+                    util.mkdir(img_dir)
+
+                    model.feed_data(val_data)
+
+                    nll = model.test()
+                    if nll is None:
+                        nll = 0
+                    nlls.append(nll)
+
+                    visuals = model.get_current_visuals()
+
+                    sr_img = None
+                    # Save SR images for reference
+                    if hasattr(model, 'heats'):
+                        for heat in model.heats:
+                            for i in range(model.n_sample):
+                                sr_img = util.tensor2img(visuals['SR', heat, i])  # uint8
+                                save_img_path = os.path.join(img_dir,
+                                                             '{:s}_{:09d}_h{:03d}_s{:d}.png'.format(img_name,
+                                                                                                    current_step,
+                                                                                                    int(heat * 100), i))
+                                util.save_img(sr_img, save_img_path)
+                    else:
+                        sr_img = util.tensor2img(visuals['SR'])  # uint8
+                        save_img_path = os.path.join(img_dir,
+                                                     '{:s}_{:d}.png'.format(img_name, current_step))
+                        util.save_img(sr_img, save_img_path)
+                    assert sr_img is not None
+
+                    # Save LQ images for reference
+                    save_img_path_lq = os.path.join(img_dir,
+                                                    '{:s}_LQ.png'.format(img_name))
+                    if not os.path.isfile(save_img_path_lq):
+                        lq_img = util.tensor2img(visuals['LQ'])  # uint8
+                        util.save_img(
+                            cv2.resize(lq_img, dsize=None, fx=opt['scale'], fy=opt['scale'],
+                                       interpolation=cv2.INTER_NEAREST),
+                            save_img_path_lq)
+
+                    # Save GT images for reference
+                    gt_img = util.tensor2img(visuals['GT'])  # uint8
+                    save_img_path_gt = os.path.join(img_dir,
+                                                    '{:s}_GT.png'.format(img_name))
+                    if not os.path.isfile(save_img_path_gt):
+                        util.save_img(gt_img, save_img_path_gt)
+
+                    # calculate PSNR
+                    crop_size = opt['scale']
+                    gt_img = gt_img / 255.
+                    sr_img = sr_img / 255.
+                    cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :]
+                    cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :]
+                    avg_psnr += util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)
+
+                avg_psnr = avg_psnr / idx
+                avg_nll = sum(nlls) / len(nlls)
+
+                # log
+                logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
+                logger_val = logging.getLogger('val')  # validation logger
+                logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format(
+                    epoch, current_step, avg_psnr))
+
+                # tensorboard logger
+                tb_logger_valid.add_scalar('loss/psnr', avg_psnr, current_step)
+                tb_logger_valid.add_scalar('loss/nll', avg_nll, current_step)
+
+                tb_logger_train.flush()
+                tb_logger_valid.flush()
+
+            #### save models and training states
+            if current_step % opt['logger']['save_checkpoint_freq'] == 0:
+                if rank <= 0:
+                    logger.info('Saving models and training states.')
+                    model.save(current_step)
+                    model.save_training_state(epoch, current_step)
+
+            timerData.tick()
+
+    with open(os.path.join(opt['path']['root'], "TRAIN_DONE"), 'w') as f:
+        f.write("TRAIN_DONE")
+
+    if rank <= 0:
+        logger.info('Saving the final model.')
+        model.save('latest')
+        logger.info('End of training.')
+
+
+if __name__ == '__main__':
+    main()

models/SRFlow/code/utils/timer.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import time
+
+
+class ScopeTimer:
+    def __init__(self, name):
+        self.name = name
+
+    def __enter__(self):
+        self.start = time.time()
+        return self
+
+    def __exit__(self, *args):
+        self.end = time.time()
+        self.interval = self.end - self.start
+        print("{} {:.3E}".format(self.name, self.interval))
+
+
+class Timer:
+    def __init__(self):
+        self.times = []
+
+    def tick(self):
+        self.times.append(time.time())
+
+    def get_average_and_reset(self):
+        if len(self.times) < 2:
+            return -1
+        avg = (self.times[-1] - self.times[0]) / (len(self.times) - 1)
+        self.times = [self.times[-1]]
+        return avg
+
+    def get_last_iteration(self):
+        if len(self.times) < 2:
+            return 0
+        return self.times[-1] - self.times[-2]
+
+
+class TickTock:
+    def __init__(self):
+        self.time_pairs = []
+        self.current_time = None
+
+    def tick(self):
+        self.current_time = time.time()
+
+    def tock(self):
+        assert self.current_time is not None, self.current_time
+        self.time_pairs.append([self.current_time, time.time()])
+        self.current_time = None
+
+    def get_average_and_reset(self):
+        if len(self.time_pairs) == 0:
+            return -1
+        deltas = [t2 - t1 for t1, t2 in self.time_pairs]
+        avg = sum(deltas) / len(deltas)
+        self.time_pairs = []
+        return avg
+
+    def get_last_iteration(self):
+        if len(self.time_pairs) == 0:
+            return -1
+        return self.time_pairs[-1][1] - self.time_pairs[-1][0]

models/SRFlow/code/utils/util.py

@@ -0,0 +1,174 @@
+import glob
+import os
+import sys
+import time
+import math
+from datetime import datetime
+import random
+import logging
+from collections import OrderedDict
+
+import natsort
+import numpy as np
+import cv2
+import torch
+from torchvision.utils import make_grid
+from shutil import get_terminal_size
+
+import yaml
+
+try:
+    from yaml import CLoader as Loader, CDumper as Dumper
+except ImportError:
+    from yaml import Loader, Dumper
+
+
+def OrderedYaml():
+    '''yaml orderedDict support'''
+    _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG
+
+    def dict_representer(dumper, data):
+        return dumper.represent_dict(data.items())
+
+    def dict_constructor(loader, node):
+        return OrderedDict(loader.construct_pairs(node))
+
+    Dumper.add_representer(OrderedDict, dict_representer)
+    Loader.add_constructor(_mapping_tag, dict_constructor)
+    return Loader, Dumper
+
+
+####################
+# miscellaneous
+####################
+
+
+def get_timestamp():
+    return datetime.now().strftime('%y%m%d-%H%M%S')
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def mkdirs(paths):
+    if isinstance(paths, str):
+        mkdir(paths)
+    else:
+        for path in paths:
+            mkdir(path)
+
+
+def mkdir_and_rename(path):
+    if os.path.exists(path):
+        new_name = path + '_archived_' + get_timestamp()
+        print('Path already exists. Rename it to [{:s}]'.format(new_name))
+        logger = logging.getLogger('base')
+        logger.info('Path already exists. Rename it to [{:s}]'.format(new_name))
+        os.rename(path, new_name)
+    os.makedirs(path)
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def setup_logger(logger_name, root, phase, level=logging.INFO, screen=False, tofile=False):
+    '''set up logger'''
+    lg = logging.getLogger(logger_name)
+    formatter = logging.Formatter('%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s',
+                                  datefmt='%y-%m-%d %H:%M:%S')
+    lg.setLevel(level)
+    if tofile:
+        log_file = os.path.join(root, phase + '_{}.log'.format(get_timestamp()))
+        fh = logging.FileHandler(log_file, mode='w')
+        fh.setFormatter(formatter)
+        lg.addHandler(fh)
+    if screen:
+        sh = logging.StreamHandler()
+        sh.setFormatter(formatter)
+        lg.addHandler(sh)
+
+
+####################
+# image convert
+####################
+
+
+def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
+    '''
+    Converts a torch Tensor into an image Numpy array
+    Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
+    Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
+    '''
+    if hasattr(tensor, 'detach'):
+        tensor = tensor.detach()
+    tensor = tensor.squeeze().float().cpu().clamp_(*min_max)  # clamp
+    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])  # to range [0,1]
+    n_dim = tensor.dim()
+    if n_dim == 4:
+        n_img = len(tensor)
+        img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
+        img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
+    elif n_dim == 3:
+        img_np = tensor.numpy()
+        img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
+    elif n_dim == 2:
+        img_np = tensor.numpy()
+    else:
+        raise TypeError(
+            'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
+    if out_type == np.uint8:
+        img_np = (img_np * 255.0).round()
+        # Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
+    return img_np.astype(out_type)
+
+
+def save_img(img, img_path, mode='RGB'):
+    cv2.imwrite(img_path, img)
+
+
+####################
+# metric
+####################
+
+
+def calculate_psnr(img1, img2):
+    # img1 and img2 have range [0, 255]
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2) ** 2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))
+
+
+def get_resume_paths(opt):
+    resume_state_path = None
+    resume_model_path = None
+    ts = opt_get(opt, ['path', 'training_state'])
+    if opt.get('path', {}).get('resume_state', None) == "auto" and ts is not None:
+        wildcard = os.path.join(ts, "*")
+        paths = natsort.natsorted(glob.glob(wildcard))
+        if len(paths) > 0:
+            resume_state_path = paths[-1]
+            resume_model_path = resume_state_path.replace('training_state', 'models').replace('.state', '_G.pth')
+    else:
+        resume_state_path = opt.get('path', {}).get('resume_state')
+    return resume_state_path, resume_model_path
+
+
+def opt_get(opt, keys, default=None):
+    if opt is None:
+        return default
+    ret = opt
+    for k in keys:
+        ret = ret.get(k, None)
+        if ret is None:
+            return default
+    return ret
+

models/SRFlow/srflow.py

@@ -0,0 +1,27 @@
+import numpy as np
+import torch
+
+from code.test import imread, impad, t, load_model, rgb
+
+def return_SRFlow_result(lr_path, conf_path='/models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.6):
+    model, opt = load_model(conf_path)
+    lr = imread(lr_path)
+
+    scale = opt['scale']
+    pad_factor = 2
+
+    h, w, c = lr.shape
+    lr = impad(lr, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
+                        right=int(np.ceil(w / pad_factor) * pad_factor - w))
+    
+    lr_t = t(lr)
+    heat = opt[heat]
+    
+    sr_t = model.get_sr(lq=lr_t, heat=heat)
+
+    sr = rgb(torch.clamp(sr_t, 0, 1))
+    sr = sr[:h * scale, :w * scale]
+
+    return sr
+
+