refactor: Adjust image augmentation parameters in training script

Author: GreenWizard2015

Core/CModelTrainer.py

@@ -28,7 +28,7 @@ def __init__(self, timesteps, model='simple', **kwargs):
     return
 
   def compile(self):
-    self._model.compile(optimizer=NNU.createOptimizer())
+    self._optimizer = NNU.createOptimizer()
     return
 
   def _pointLoss(self, ytrue, ypred):
@@ -46,22 +46,22 @@ def _trainStep(self, Data):
     x, (y, ) = Data
     y = y[..., 0, :]
     losses = {}
-    parts = list(self._embeddings.values()) + [self._model]
-    TV = sum([p.trainable_variables for p in parts], [])
-    with tf.GradientTape(watch_accessed_variables=False) as tape:
-      tape.watch(TV)
+    with tf.GradientTape() as tape:
       data = x['augmented']
       data = self._replaceByEmbeddings(data)
       predictions = self._model(data, training=True)
-      predictions = dict(**predictions['intermediate'], final=predictions['result'])
-      for name, pts in predictions.items():
+      intermediate = predictions['intermediate']
+      losses['final'] = tf.reduce_mean(self._pointLoss(y, predictions['result']))
+      for name, encoder in self._intermediateEncoders.items():
+        latent = intermediate[name]
+        pts = encoder(latent, training=True)
         loss = self._pointLoss(y, pts)
         losses['loss-%s' % name] = tf.reduce_mean(loss)
         continue
       loss = sum(losses.values())
       losses['loss'] = loss
 
-    self._model.optimizer.minimize(loss, TV, tape=tape)
+    self._optimizer.minimize(loss, tape.watched_variables(), tape=tape)
     ###############
     return losses
 

Core/CModelWrapper.py

@@ -5,7 +5,7 @@
 from tensorflow.keras import layers as L
 
 class CModelWrapper:
-  def __init__(self, timesteps, model='simple', user=None, stats=None, **kwargs):
+  def __init__(self, timesteps, model='simple', user=None, stats=None, use_encoders=True, **kwargs):
     if user is None:
       user = {
         'userId': 0,
@@ -37,6 +37,15 @@ def __init__(self, timesteps, model='simple', user=None, stats=None, **kwargs):
       'placeId': L.Embedding(len(stats['placeId']), embeddings['size']),
       'screenId': L.Embedding(len(stats['screenId']), embeddings['size']),
     }
+    self._intermediateEncoders = {}
+    if use_encoders:
+      shapes = self._modelRaw['intermediate shapes']
+      for name, shape in shapes.items():
+        enc = networks.IntermediatePredictor(name='%s-encoder' % name)
+        enc.build(shape)
+        self._intermediateEncoders[name] = enc
+        continue
+   
     if 'weights' in kwargs:
       self.load(**kwargs['weights'])
     return
@@ -80,6 +89,15 @@ def save(self, folder=None, postfix=''):
       embeddings[nm] = weights
       continue
     np.savez_compressed(path.replace('.h5', '-embeddings.npz'), **embeddings)
+    # save intermediate encoders
+    if self._intermediateEncoders:
+      encoders = {}
+      for nm, encoder in self._intermediateEncoders.items():
+        # save each variable separately
+        for ww in encoder.trainable_variables:
+          encoders['%s-%s' % (nm, ww.name)] = ww.numpy()
+        continue
+      np.savez_compressed(path.replace('.h5', '-intermediate-encoders.npz'), **encoders)
     return
 
   def load(self, folder=None, postfix='', embeddings=False):
@@ -92,6 +110,16 @@ def load(self, folder=None, postfix='', embeddings=False):
         if not emb.built: emb.build((None, w.shape[0]))
         emb.set_weights([w]) # replace embeddings
         continue
+    
+    if self._intermediateEncoders:
+      encodersName = path.replace('.h5', '-intermediate-encoders.npz')
+      if os.path.isfile(encodersName):
+        encoders = np.load(encodersName)
+        for nm, encoder in self._intermediateEncoders.items():
+          for ww in encoder.trainable_variables:
+            w = encoders['%s-%s' % (nm, ww.name)]
+            ww.assign(w)
+          continue
     return
 
   def lock(self, isLocked):
@@ -101,4 +129,7 @@ def lock(self, isLocked):
   @property
   def timesteps(self):
     return self._timesteps
-  
+  
+  def trainable_variables(self):
+    parts = list(self._embeddings.values()) + [self._model] + list(self._intermediateEncoders.values())
+    return sum([p.trainable_variables for p in parts], [])

NN/Utils.py

@@ -27,7 +27,7 @@ def call(self, x):
     coefs = tf.pow(self._base, powers)
     return tf.reduce_sum(x * coefs, axis=-1)
 ############################################
-SMLP_GLOBAL_DROPOUT = 0.01
+SMLP_GLOBAL_DROPOUT = 0.0
 class sMLP(tf.keras.layers.Layer):
   def __init__(self, sizes, activation='linear', dropout=None, **kwargs):
     super().__init__(**kwargs)
@@ -40,7 +40,11 @@ def __init__(self, sizes, activation='linear', dropout=None, **kwargs):
       continue
     self._F = tf.keras.Sequential(layers, name=self.name + '/F')
     return
-  
+ 
+  def build(self, input_shape):
+    self._F.build(input_shape)
+    return super().build(input_shape)
+   
   def call(self, x, **kwargs):
     return self._F(x, **kwargs)
 ############################################

NN/networks.py

@@ -25,6 +25,7 @@ def build(self, input_shape):
       sizes=[128, 64, 32], activation='relu',
       name='%s/MLP' % self.name
     )
+    self._mlp.build(input_shape)
     self._decodePoints = L.Dense(2, name='%s/DecodePoints' % self.name)
     return super().build(input_shape)
 
@@ -168,8 +169,6 @@ def Face2LatentModel(
     **stepsData['intermediate'],
     **res['intermediate'],
   }
-  # drop all intermediate outputs
-  res['intermediate'] = {}
 
   inputs = {
     'points': points,
@@ -181,13 +180,10 @@ def Face2LatentModel(
     'screenId': screenIdEmb,
   }
 
-  intermediate = res['intermediate']
-  IP = lambda x: IntermediatePredictor()(x) # own IntermediatePredictor for each output
-  res['intermediate'] = {k: IP(x) for k, x in intermediate.items()}
-  res['result'] = IP(res['latent'])
-  
+  res['result'] = IntermediatePredictor()(res['latent'])
   main = tf.keras.Model(inputs=inputs, outputs=res)
   return {
+    'intermediate shapes': {k: v.shape for k, v in res['intermediate'].items()},
     'main': main,
     'Face2Step': Face2Step,
     'Step2Latent': Step2Latent,

scripts/train.py

@@ -182,22 +182,24 @@ def _trainer_from(args):
   raise Exception('Unknown trainer: %s' % (args.trainer, ))
 
 def averageModels(folder, model, noiseStd=0.0):
-  TV = [np.zeros_like(x) for x in model._model.get_weights()]
+  TV = [np.zeros_like(x) for x in model.trainable_variables()]
   N = 0
   for nm in glob.glob(os.path.join(folder, '*.h5')):
     if not('best' in nm): continue # only the best models
     model.load(nm, embeddings=True)
     # add the weights to the total
-    weights = model._model.get_weights()
+    weights = model.trainable_variables()
     for i in range(len(TV)):
-      TV[i] += weights[i]
+      TV[i] += weights[i].numpy()
       continue
     N += 1
     continue
 
   # average the weights
   TV = [(x / N) + np.random.normal(0.0, noiseStd, x.shape) for x in TV]
-  model._model.set_weights(TV)
+  for v, new in zip(model.trainable_variables(), TV):
+    v.assign(new)
+    continue
   model.compile() # recompile the model with the new weights
   return
 
@@ -230,7 +232,7 @@ def main(args):
       ),
     )
   )
-  model = dict(timesteps=timesteps, stats=stats)
+  model = dict(timesteps=timesteps, stats=stats, use_encoders=args.with_enconders)
   if args.model is not None:
     model['weights'] = dict(folder=folder, postfix=args.model, embeddings=args.embeddings)
   if args.modelId is not None:
@@ -327,6 +329,9 @@ def performRandomSearch(epoch=0):
     '--restarts', type=int, default=1,
     help='Number of times to restart the model reinitializing the weights'
   )
+  parser.add_argument(
+    '--with-enconders', default=False, action='store_true',
+  )
 
   main(parser.parse_args())
   pass