refactor: modify structure according to new setup

Author: samedii

guild.yml

@@ -3,7 +3,6 @@
 - model: nvae-mixture-logistic
   sourcecode:
     - '*.py'
-    - 'vae/splits/*.json'
     - guild.yml
     - exclude:
         dir:

update_splits.py

@@ -1,5 +1,7 @@
 from vae import problem
 from vae import datastream
 from vae import architecture
+from vae import metrics
 
+from vae.log_examples import log_examples
 from vae.train import train

vae/__init__.py

@@ -1,16 +1,14 @@
-import numpy as np
-import pandas as pd
 from datastream import Datastream
 
-from vae import problem, splits
+from vae import problem
 
 
 def evaluate_datastreams():
-    evaluate_datasets = problem.evaluate_datasets()
+    datasets = problem.datasets()
     return {
-        split_name: Datastream(dataset)
+        split_name: Datastream(dataset).take(256)
         for split_name, dataset in dict(
-            gradient=evaluate_datasets['train'],
-            compare=evaluate_datasets['compare'],
+            gradient=datasets['train'],
+            compare=datasets['compare'],
         ).items()
     }

vae/datastream/evaluate_datastreams.py

@@ -1,14 +1,12 @@
-import numpy as np
 from datastream import Datastream
 
-from vae.datastream import (
-    evaluate_datastreams, augmenter
-)
+from vae import problem
+from vae.datastream import augmenter
 
 
-def GradientDatastream():    
+def GradientDatastream():
     augmenter_ = augmenter()
     return (
-        evaluate_datastreams()['gradient']
+        Datastream(problem.datasets()['gradient'])
         .map(lambda example: example.augment(augmenter_))
     )

vae/datastream/gradient_datastream.py

@@ -0,0 +1,86 @@
+import numpy as np
+import torch
+from workflow.torch import module_eval
+
+from vae import architecture
+
+
+def log_examples(description, trainer, model):
+    def log_examples_(engine, logger, event_name):
+        n_examples = 5
+        indices = np.random.choice(
+            len(engine.state.output['predictions']),
+            n_examples,
+            replace=False,
+        )
+
+        logger.writer.add_images(
+            f'{description}/predictions',
+            np.stack([
+                np.concatenate([
+                    np.array(
+                        engine.state.output['examples'][index]
+                        .representation()
+                    ),
+                    np.array(
+                        engine.state.output['predictions'][index]
+                        .representation()
+                    ),
+                ], axis=0) / 255
+                for index in indices
+            ]),
+            trainer.state.epoch,
+            dataformats='NHWC',
+        )
+
+        with torch.no_grad(), module_eval(model) as eval_model:
+            std_samples = [
+                eval_model.generated(16, prior_std)
+                for prior_std in np.linspace(0.4, 1.1, num=8)
+            ]
+
+        logger.writer.add_images(
+            f'{description}/samples',
+            np.stack([np.concatenate([
+                np.concatenate([
+                    np.array(sample.representation())
+                    for sample in samples
+                ], axis=1)
+                for samples in std_samples
+            ], axis=0)]) / 255,
+            trainer.state.epoch,
+            dataformats='NHWC',
+        )
+
+        with torch.no_grad(), module_eval(model) as eval_model:
+            partial_samples = [
+                eval_model.partially_generated(
+                    architecture.FeaturesBatch.from_examples(
+                        [
+                            engine.state.output['examples'][index]
+                            for index in indices
+                        ]
+                    ).image_batch,
+                    sample=[
+                        index == sample_index
+                        for index in range(model.levels)
+                    ],
+                    prior_std=0.7,
+                )
+                for sample_index in range(model.levels)
+            ]
+
+        logger.writer.add_images(
+            f'{description}/partially_sampled',
+            np.concatenate([
+                np.stack([
+                    np.array(sample.representation())
+                    for sample in samples
+                ])
+                for samples in partial_samples
+            ], axis=1) / 255,
+            trainer.state.epoch,
+            dataformats='NHWC',
+        )
+
+    return log_examples_

vae/log_examples.py

@@ -1,4 +1,2 @@
 from vae.problem.example import Example
-from vae.problem.evaluate_datasets import (
-    evaluate_datasets
-)
+from vae.problem.datasets import datasets

vae/problem/__init__.py

@@ -3,7 +3,7 @@
 import pandas as pd
 from datastream import Dataset
 
-from vae import problem, splits
+from vae import problem
 
 
 def dataframe():
@@ -30,25 +30,12 @@ def dataset(dataframe):
     )
 
 
-def evaluate_datasets(frozen=True):
-    datasets = (
+def datasets(frozen=True):
+    return (
         dataset(dataframe())
         .split(
             key_column='key',
             proportions=dict(train=0.8, compare=0.2),
-            # filepath=splits.compare,
-            # frozen=frozen,
-            # don't save split until we have a solution for remote training
-            # guild is not saving .json as sourcecode to remote
             seed=177,
         )
     )
-    # TODO: temporary
-    return dict(
-        train=datasets['train'],
-        compare=datasets['compare'].split(
-            key_column='key',
-            proportions=dict(keep=0.05, throw=0.95),
-            seed=523,
-        )['keep'],
-    )