wip

Author: GreenWizard2015

Core/CBaseModel.py

@@ -0,0 +1,59 @@
+import os
+import numpy as np
+from tensorflow.keras import layers as L
+
+class CBaseModel:
+  def __init__(self, model, embeddings, submodels):
+    self._model = model
+    self._embeddings = {
+      'userId': L.Embedding(embeddings['userId'], embeddings['size']),
+      'placeId': L.Embedding(embeddings['placeId'], embeddings['size']),
+      'screenId': L.Embedding(embeddings['screenId'], embeddings['size']),
+    }
+    self._submodels = submodels
+    return  
+
+  def replaceByEmbeddings(self, data):
+    data = dict(**data) # copy
+    for name, emb in self._embeddings.items():
+      data[name] = emb(data[name][..., 0])
+      continue
+    return data
+
+  def _modelFilename(self, folder, postfix=''):
+    postfix = '-' + postfix if postfix else ''
+    return os.path.join(folder, '%s%s.h5' % (self._modelID, postfix))
+  
+  def save(self, folder=None, postfix=''):
+    path = self._modelFilename(folder, postfix)
+    if 1 < len(self._submodels):
+      for i, model in enumerate(self._submodels):
+        model.save_weights(path.replace('.h5', '-%d.h5' % i))
+    else:
+      self._submodels[0].save_weights(path)
+
+    embeddings = {}
+    for nm in self._embeddings.keys():
+      weights = self._embeddings[nm].get_weights()[0]
+      embeddings[nm] = weights
+    
+    np.savez_compressed(path.replace('.h5', '-embeddings.npz'), **embeddings)
+    
+  def load(self, folder=None, postfix='', embeddings=False):
+    path = self._modelFilename(folder, postfix) if not os.path.isfile(folder) else folder
+    if 1 < len(self._submodels):
+      for i, model in enumerate(self._submodels):
+        model.load_weights(path.replace('.h5', '-%d.h5' % i))
+    else:
+      self._submodels[0].load_weights(path)
+      
+    if embeddings:
+      embeddings = np.load(path.replace('.h5', '-embeddings.npz'))
+      for nm, emb in self._embeddings.items():
+        w = embeddings[nm]
+        if not emb.built: emb.build((None, w.shape[0]))
+        emb.set_weights([w]) # replace embeddings
+    
+  def trainable_variables(self):
+    parts = list(self._embeddings.values()) + self._submodels
+    return sum([p.trainable_variables for p in parts], [])

Core/CDataSamplerInpainting.py

@@ -3,6 +3,7 @@
 from Core.Utils import FACE_MESH_POINTS
 
 import numpy as np
+import tensorflow as tf
 
 '''
 This sampler are sample N frames from the dataset, where N is the number of timesteps.
@@ -24,8 +25,9 @@
   - The target point.
 '''
 class CDataSamplerInpainting(CBaseDataSampler):
-    def __init__(self, storage, batch_size, minFrames, defaults={}, maxT=1.0, cumulative_time=True):
+    def __init__(self, storage, batch_size, minFrames, keys, defaults={}, maxT=1.0, cumulative_time=True):
         super().__init__(storage, batch_size, minFrames, defaults, maxT, cumulative_time)
+        self._keys = keys
 
     def _stepsFor(self, mainInd, steps, stepsSampling='uniform', **_):
         if (steps is None) or (1 == steps): return [(mainInd, 0.0)]
@@ -46,7 +48,7 @@ def _stepsFor(self, mainInd, steps, stepsSampling='uniform', **_):
     def sample(self, **kwargs):
         kwargs = {**self._defaults, **kwargs}
         timesteps = kwargs.get('timesteps', None)
-        N = kwargs.get('N', self._batchSize)
+        N = kwargs.get('N', self._batchSize) // len(self._keys)
         indexes = []
         for _ in range(N):
             added = False
@@ -101,6 +103,7 @@ def sampleByIds(self, ids, **kwargs):
     def _indexes2XY(self, indexesAndTime, kwargs):
         timesteps = kwargs.get('timesteps', None)
         assert timesteps is not None, 'The number of timesteps must be defined.'
+        B = len(indexesAndTime) // timesteps
         samples = [self._storage[i] for i, _ in indexesAndTime]
         ##############
         userIds = np.unique([x['userId'] for x in samples])
@@ -130,7 +133,25 @@ def _indexes2XY(self, indexesAndTime, kwargs):
             ),
             userIds[0], placeIds[0], screenIds[0]
         )
-        X = X['clean']
+        for k in X.keys():
+            # add the target point to the X
+            targets = np.array([x['goal'] for x in samples], np.float32).reshape((B, timesteps, 2))
+            X[k]['target'] = tf.constant(targets, dtype=tf.float32)
+
+        if 1 == len(self._keys):
+            X = X[self._keys[0]]
+        else:
+            res = {}
+            k = self._keys[0]
+            subkeys = list(X[k].keys())
+            for k in subkeys:
+                values = [X[key][k] for key in self._keys]
+                res[k] = tf.concat(values, axis=0)
+                continue
+            X = res
+            indexesAndTime = indexesAndTime * len(self._keys)
+            B = len(self._keys) * B
+
         ###############
         # generate the target data
         targets = kwargs.get('targets', {'keypoints': timesteps, 'total': timesteps})
@@ -141,7 +162,7 @@ def _indexes2XY(self, indexesAndTime, kwargs):
 
         samples_indexes = np.array([ i  for i, _ in indexesAndTime], np.int32)
         samples_indexes = samples_indexes.reshape((-1, timesteps))
-        B = samples_indexes.shape[0]
+        assert samples_indexes.shape[0] == B, 'Invalid number of samples: %d != %d' % (samples_indexes.shape[0], B)
         targetsIdx = np.zeros((B, T), np.int32)
         for i in range(B):
             # sample K frames from the X
@@ -186,20 +207,20 @@ def _indexes2XY(self, indexesAndTime, kwargs):
                 Y['right eye'][i, j] = data['right eye'][p:p+32, p:p+32]
                 Y['time'][i, j] = (data['time'] - startT) / duration
                 Y['target'][i, j] = data['goal']
-                
+        # eyes in 0..255, so we need to normalize them
+        Y['left eye'] /= 255.0
+        Y['right eye'] /= 255.0
         # check that time is between 0 and 1
         assert np.all((0 <= Y['time']) & (Y['time'] <= 1)), 'Invalid time: ' + str(Y['time'])
-        B = Y['points'].shape[0] 
         for k, v in X.items():
             assert B == v.shape[0], f'Invalid batch size for X[{k}]: {v.shape[0]} != {B} ({v.shape})'
         for k, v in Y.items():
             assert B == v.shape[0], f'Invalid batch size for Y[{k}]: {v.shape[0]} != {B} ({v.shape})'
         return (X, Y)
 
     def merge(self, samples, expected_batch_size):
-        # each dictionary contains the subkeys: points, left eye, right eye, time, userId, placeId, screenId
         X = {}
-        for subkey in ['points', 'left eye', 'right eye', 'time', 'userId', 'placeId', 'screenId']:
+        for subkey in ['points', 'left eye', 'right eye', 'time', 'userId', 'placeId', 'screenId', 'target']:
             data = [x[subkey] for x, _ in samples]
             X[subkey] = np.concatenate(data, axis=0)
             assert X[subkey].shape[0] == expected_batch_size, 'Invalid batch size: %d != %d' % (X[subkey].shape[0], expected_batch_size)
@@ -211,4 +232,4 @@ def merge(self, samples, expected_batch_size):
             Y[subkey] = np.concatenate(data, axis=0)
             assert Y[subkey].shape[0] == expected_batch_size, 'Invalid batch size: %d != %d' % (Y[subkey].shape[0], expected_batch_size)
             continue
-        return (X, (Y, ))
+        return (X, Y)

Core/CInpaintingTrainer.py

@@ -0,0 +1,101 @@
+import tensorflow as tf
+import time
+import NN.Utils as NNU
+import NN.networks as networks
+from Core.CBaseModel import CBaseModel
+
+class CInpaintingTrainer:
+  def __init__(self, timesteps, model='simple', KP=5, **kwargs):
+    stats = kwargs.get('stats', None)
+    embeddingsSize = kwargs.get('embeddingsSize', 64)
+    latentSize = kwargs.get('latentSize', 64)
+    embeddings = {
+      'userId': len(stats['userId']),
+      'placeId': len(stats['placeId']),
+      'screenId': len(stats['screenId']),
+      'size': embeddingsSize,
+    }
+
+    self._encoder = networks.InpaintingEncoderModel(
+      steps=timesteps, latentSize=latentSize,
+      embeddingsSize=embeddingsSize,
+      KP=KP,
+    )
+    self._decoder = networks.InpaintingDecoderModel(
+      latentSize=latentSize,
+      embeddingsSize=embeddingsSize,
+      KP=KP,
+    )
+    self._model = CBaseModel(
+       model=model, embeddings=embeddings, submodels=[self._encoder, self._decoder]
+    )
+    self.compile()
+    # add signatures to help tensorflow optimize the graph
+    specification = networks.InpaintingInputSpec()
+    self._trainStep = tf.function(
+      self._trainStep,
+      input_signature=[specification]
+    )
+    self._eval = tf.function(
+      self._eval,
+      input_signature=[specification]
+    )
+    return
+  
+  def compile(self):
+    self._optimizer = NNU.createOptimizer()
+
+  def _trainStep(self, Data):
+    print('Instantiate _trainStep')
+    ###############
+    x, y = Data
+    losses = {}
+    with tf.GradientTape() as tape:
+      x = self._model.replaceByEmbeddings(x)
+      latents = self._encoder(x, training=True)['latent']
+      decoderArgs = {
+        'keyPoints': latents,
+        'time': y['time'],
+        'userId': x['userId'],
+        'placeId': x['placeId'],
+        'screenId': x['screenId'],
+      }
+      predictions = self._decoder(decoderArgs, training=True)
+      losses = {}
+      for k in predictions.keys():
+        pred = predictions[k]
+        gt = y[k]
+        tf.assert_equal(tf.shape(pred), tf.shape(gt))
+        loss = tf.losses.mse(gt, pred)
+        losses[f"loss-{k}"] = tf.reduce_mean(loss)
+        
+      # calculate total loss and final loss
+      losses['loss'] = loss = sum(losses.values())
+  
+    self._optimizer.minimize(loss, tape.watched_variables(), tape=tape)
+    ###############
+    return losses
+
+  def fit(self, data):
+    t = time.time()
+    losses = self._trainStep(data)
+    losses = {k: v.numpy() for k, v in losses.items()}
+    return {'time': int((time.time() - t) * 1000), 'losses': losses}
+  
+  def _eval(self, xy):
+    print('Instantiate _eval')
+    x, (y,) = xy
+    x = self._replaceByEmbeddings(x)
+    y = y[:, :, 0]
+    predictions = self._model(x, training=False)
+    points = predictions['result'][:, :, :]
+    tf.assert_equal(tf.shape(points), tf.shape(y))
+
+    loss = self._pointLoss(y, points)
+    tf.assert_equal(tf.shape(loss), tf.shape(y)[:2])
+    _, dist = NNU.normVec(points - y)
+    return loss, points, dist
+
+  def eval(self, data):
+    loss, sampled, dist = self._eval(data)
+    return loss.numpy(), sampled.numpy(), dist.numpy()