From 055b4b369d0f85bbb92dfdeda6ea7c81e2690e7d Mon Sep 17 00:00:00 2001
From: Stanislaw Adaszewski <stanislaw.adaszewski@roche.com>
Date: Thu, 16 Jul 2020 16:49:19 +0200
Subject: [PATCH] run_epoch() progresses at a crawl.

---
 experiments/decagon_run/decagon_run.py | 23 +++++++++++++++++++++++
 src/icosagon/trainloop.py              |  3 ++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/experiments/decagon_run/decagon_run.py b/experiments/decagon_run/decagon_run.py
index d21cce0..a490628 100644
--- a/experiments/decagon_run/decagon_run.py
+++ b/experiments/decagon_run/decagon_run.py
@@ -1,6 +1,10 @@
 #!/usr/bin/env python3
 
 from icosagon.data import Data
+from icosagon.trainprep import TrainValTest, \
+    prepare_training
+from icosagon.model import Model
+from icosagon.trainloop import TrainLoop
 import os
 import pandas as pd
 from bisect import bisect_left
@@ -88,9 +92,28 @@ def load_data():
     print()
     print('OK')
 
+    return data
+
+
+def _wrap(obj, method_name):
+    orig_fn = getattr(obj, method_name)
+    def fn(*args, **kwargs):
+        print(f'{method_name}() :: ENTER')
+        res = orig_fn(*args, **kwargs)
+        print(f'{method_name}() :: EXIT')
+        return res
+    setattr(obj, method_name, fn)
+
 
 def main():
     data = load_data()
+    prep_d = prepare_training(data, TrainValTest(.8, .1, .1))
+    _wrap(Model, 'build')
+    model = Model(prep_d)
+    _wrap(TrainLoop, 'build')
+    _wrap(TrainLoop, 'run_epoch')
+    loop = TrainLoop(model, batch_size=1000000)
+    loop.run_epoch()
 
 
 if __name__ == '__main__':
diff --git a/src/icosagon/trainloop.py b/src/icosagon/trainloop.py
index 5492878..3392a40 100644
--- a/src/icosagon/trainloop.py
+++ b/src/icosagon/trainloop.py
@@ -46,7 +46,8 @@ class TrainLoop(object):
         # pred = self.model(None)
         # n = len(list(iter(batch)))
         loss_sum = 0
-        for indices in batch:
+        for i, indices in enumerate(batch):
+            print('%.2f%% (%d/%d)' % (i * batch.batch_size * 100 / batch.total_edge_count, i * batch.batch_size, batch.total_edge_count))
             self.opt.zero_grad()
             pred = self.model(None)
             pred = flatten_predictions(pred)