Update for CUDA(11.4)

generator.py

@@ -2,7 +2,7 @@
 import random
 import numpy as np
 import cv2
-from keras.utils import Sequence
+from tensorflow.keras.utils import Sequence
 
 
 class NoisyImageGenerator(Sequence):
@@ -62,7 +62,7 @@ def __init__(self, image_dir, val_noise_model):
             h, w, _ = y.shape
             y = y[:(h // 16) * 16, :(w // 16) * 16]  # for stride (maximum 16)
             x = val_noise_model(y)
-            self.data.append([np.expand_dims(x, axis=0), np.expand_dims(y, axis=0)])
+            self.data.append((np.expand_dims(x, axis=0), np.expand_dims(y, axis=0)))
 
     def __len__(self):
         return self.image_num

model.py

@@ -1,9 +1,9 @@
-from keras.models import Model
-from keras.layers import Input, Add, PReLU, Conv2DTranspose, Concatenate, MaxPooling2D, UpSampling2D, Dropout
-from keras.layers.convolutional import Conv2D
-from keras.layers.normalization import BatchNormalization
-from keras.callbacks import Callback
-from keras import backend as K
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Input, Add, PReLU, Conv2DTranspose, Concatenate, MaxPooling2D, UpSampling2D, Dropout
+from tensorflow.keras.layers import Conv2D
+from tensorflow.keras.layers import BatchNormalization
+from tensorflow.keras.callbacks import Callback
+from tensorflow.keras import backend as K
 import tensorflow as tf
 
 
@@ -34,8 +34,8 @@ def on_epoch_begin(self, epoch, logs=None):
 
 
 def tf_log10(x):
-    numerator = tf.log(x)
-    denominator = tf.log(tf.constant(10, dtype=numerator.dtype))
+    numerator = tf.math.log(x)
+    denominator = tf.math.log(tf.constant(10, dtype=numerator.dtype))
     return numerator / denominator
 
 

requirements.txt

@@ -1,4 +1,3 @@
-numpy
-keras
-tensorflow
-opencv-python
+numpy==1.19.5
+tensorflow==2.6
+opencv-python==4.5.3.56

train.py

@@ -1,8 +1,8 @@
 import argparse
 import numpy as np
 from pathlib import Path
-from keras.callbacks import LearningRateScheduler, ModelCheckpoint
-from keras.optimizers import Adam
+from tensorflow.keras.callbacks import LearningRateScheduler, ModelCheckpoint
+from tensorflow.keras.optimizers import Adam
 from model import get_model, PSNR, L0Loss, UpdateAnnealingParameter
 from generator import NoisyImageGenerator, ValGenerator
 from noise_model import get_noise_model