template-nvae/vae/log_examples.py at 48377dadfcab3147ca619e01c8584c5ee39fc752 · nextml-code/template-nvae · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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_