1

core/joblib/SubprocessorBase.py

@@ -12,9 +12,11 @@ class Subprocessor(object):
 
     class Cli(object):
         def __init__ ( self, client_dict ):
-            self.s2c = multiprocessing.Queue()
-            self.c2s = multiprocessing.Queue()
-            self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,) )
+            s2c = multiprocessing.Queue()
+            c2s = multiprocessing.Queue()
+            self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,s2c,c2s) )
+            self.s2c = s2c
+            self.c2s = c2s            
             self.p.daemon = True
             self.p.start()
 
@@ -52,9 +54,8 @@ class Subprocessor(object):
         def log_err(self, msg): self.c2s.put ( {'op': 'log_err' , 'msg':msg } )
         def progress_bar_inc(self, c): self.c2s.put ( {'op': 'progress_bar_inc' , 'c':c } )
 
-        def _subprocess_run(self, client_dict):
+        def _subprocess_run(self, client_dict, s2c, c2s):
             data = None
-            s2c, c2s = self.s2c, self.c2s
             try:
                 self.on_initialize(client_dict)
 
@@ -86,6 +87,12 @@ class Subprocessor(object):
 
             c2s.put ( {'op': 'error', 'data' : data} )
         
+        # disable pickling
+        def __getstate__(self):
+            return dict()
+        def __setstate__(self, d):
+            self.__dict__.update(d)
+            
     #overridable
     def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=False):
         if not issubclass(SubprocessorCli_class, Subprocessor.Cli):
@@ -179,7 +186,7 @@ class Subprocessor(object):
                             break
                         io.process_messages(0.005)
             except:
-                raise Exception ("Unable to start subprocess %s" % (name))
+                raise Exception (f"Unable to start subprocess {name}. Error: {traceback.format_exc()}")
 
         if len(self.clis) == 0:
             raise Exception ("Unable to start Subprocessor '%s' " % (self.name))

core/leras/initializers.py

@@ -1,5 +1,10 @@
+import multiprocessing
+
 import numpy as np
 
+from core.joblib import Subprocessor
+
+
 def initialize_initializers(nn):
     tf = nn.tf
     from tensorflow.python.ops import init_ops
@@ -13,7 +18,15 @@ def initialize_initializers(nn):
                 return tf.zeros( shape, name="_cai_")
 
             @staticmethod
-            def generate(shape, eps_std=0.05, dtype=np.float32):
+            def generate_batch( data_list, eps_std=0.05 ):
+                # list of (shape, np.dtype)
+                return CAInitializerSubprocessor (data_list).run()
+  
+    nn.initializers = initializers
+
+class CAInitializerSubprocessor(Subprocessor):
+    @staticmethod
+    def generate(shape, dtype=np.float32, eps_std=0.05):
         """
         Super fast implementation of Convolution Aware Initialization for 4D shapes
         Convolution Aware Initialization https://arxiv.org/abs/1702.06295
@@ -49,4 +62,43 @@ def initialize_initializers(nn):
         x = x * np.sqrt( (2/fan_in) / np.var(x) )
         x = np.transpose( x, (2, 3, 1, 0) )
         return x.astype(dtype)
-    nn.initializers = initializers
+                
+    class Cli(Subprocessor.Cli):
+        #override
+        def process_data(self, data):
+            idx, shape, dtype = data
+            weights = CAInitializerSubprocessor.generate (shape, dtype)
+            return idx, weights
+
+    #override
+    def __init__(self, data_list):
+        self.data_list = data_list
+        self.data_list_idxs = [*range(len(data_list))]
+        self.result = [None]*len(data_list)
+        super().__init__('CAInitializerSubprocessor', CAInitializerSubprocessor.Cli)
+
+    #override
+    def process_info_generator(self):
+        for i in range( min(multiprocessing.cpu_count(), len(self.data_list)) ):
+            yield 'CPU%d' % (i), {}, {}
+
+    #override
+    def get_data(self, host_dict):
+        if len (self.data_list_idxs) > 0:
+            idx = self.data_list_idxs.pop(0)
+            shape, dtype = self.data_list[idx]
+            return idx, shape, dtype
+        return None
+
+    #override
+    def on_data_return (self, host_dict, data):
+        self.data_list_idxs.insert(0, data)
+
+    #override
+    def on_result (self, host_dict, data, result):
+        idx, weights = result
+        self.result[idx] = weights
+
+    #override
+    def get_result(self):
+        return self.result

core/leras/layers.py

@@ -77,18 +77,25 @@ def initialize_layers(nn):
 
         def init_weights(self):
             ops = []
-            tuples = []
+            
+            ca_tuples_w = []            
+            ca_tuples = []
             for w in self.get_weights():
                 initializer = w.initializer
                 for input in initializer.inputs:
                     if "_cai_" in input.name:
-                        tuples.append ( (w, nn.initializers.ca.generate(w.shape.as_list(), dtype= w.dtype.as_numpy_dtype) ) )
+                        ca_tuples_w.append (w)
+                        ca_tuples.append ( (w.shape.as_list(), w.dtype.as_numpy_dtype) )
                         break
                 else:
                     ops.append (initializer)
 
+            if len(ops) != 0:
                 nn.tf_sess.run (ops)
-            nn.tf_batch_set_value(tuples)
+            
+            if len(ca_tuples) != 0:
+                nn.tf_batch_set_value( [*zip(ca_tuples_w, nn.initializers.ca.generate_batch (ca_tuples))] )
+            
     nn.Saveable = Saveable
     
     class LayerBase():

models/Model_Quick96/Model.py

@@ -231,7 +231,7 @@ class QModel(ModelBase):
                                           [self.decoder_dst, 'decoder_dst.npy']  ]
 
             if self.is_training:
-                self.src_dst_trainable_weights = self.encoder.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
+                self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
                 
                 # Initialize optimizers
                 self.src_dst_opt = nn.TFRMSpropOptimizer(lr=2e-4, lr_dropout=0.3, name='src_dst_opt')

models/Model_SAEHD/Model.py

@@ -279,7 +279,6 @@ class SAEHDModel(ModelBase):
                                 
                 if self.is_hd:
                     x, upx = self.res0(z)                                                   
-                                                  
                     x = self.upscale0(x)
                     x = tf.nn.leaky_relu(x + upx, 0.2)                    
                     x, upx = self.res1(x)
@@ -410,8 +409,8 @@ class SAEHDModel(ModelBase):
                 self.src_dst_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
                 self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
                 if 'df' in archi:
-                    self.src_dst_all_trainable_weights = self.encoder.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
-                    self.src_dst_trainable_weights = self.encoder.get_weights() + self.decoder_src.get_weights_ex(learn_mask) + self.decoder_dst.get_weights_ex(learn_mask)
+                    self.src_dst_all_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
+                    self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights_ex(learn_mask) + self.decoder_dst.get_weights_ex(learn_mask)
 
                 elif 'liae' in archi:
                     self.src_dst_all_trainable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights()