/
normal_sampler.py
85 lines (65 loc) · 2.92 KB
/
normal_sampler.py
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
import torch
from ...models import BaseAE
from ..base import BaseSampler
from .normal_config import NormalSamplerConfig
class NormalSampler(BaseSampler):
"""Samples from a Standard normal distribution in the Autoencoder's latent space.
Args:
model (BaseAE): The vae model to sample from.
sampler_config (BaseSamplerConfig): An instance of BaseSamplerConfig in which any sampler's
parameters is made available. If None a default configuration is used. Default: None
"""
def __init__(self, model: BaseAE, sampler_config: NormalSamplerConfig = None):
if sampler_config is None:
sampler_config = NormalSamplerConfig()
BaseSampler.__init__(self, model=model, sampler_config=sampler_config)
def sample(
self,
num_samples: int = 1,
batch_size: int = 500,
output_dir: str = None,
return_gen: bool = True,
save_sampler_config: bool = False,
) -> torch.Tensor:
"""Main sampling function of the sampler.
Args:
num_samples (int): The number of samples to generate
batch_size (int): The batch size to use during sampling
output_dir (str): The directory where the images will be saved. If does not exist the
folder is created. If None: the images are not saved. Defaults: None.
return_gen (bool): Whether the sampler should directly return a tensor of generated
data. Default: True.
save_sampler_config (bool): Whether to save the sampler config. It is saved in
output_dir.
Returns:
~torch.Tensor: The generated images
"""
full_batch_nbr = int(num_samples / batch_size)
last_batch_samples_nbr = num_samples % batch_size
x_gen_list = []
for i in range(full_batch_nbr):
z = torch.randn(batch_size, self.model.latent_dim).to(self.device)
x_gen = self.model.decoder(z)["reconstruction"].detach()
if output_dir is not None:
for j in range(batch_size):
self.save_img(
x_gen[j], output_dir, "%08d.png" % int(batch_size * i + j)
)
x_gen_list.append(x_gen)
if last_batch_samples_nbr > 0:
z = torch.randn(last_batch_samples_nbr, self.model.latent_dim).to(
self.device
)
x_gen = self.model.decoder(z)["reconstruction"].detach()
if output_dir is not None:
for j in range(last_batch_samples_nbr):
self.save_img(
x_gen[j],
output_dir,
"%08d.png" % int(batch_size * full_batch_nbr + j),
)
x_gen_list.append(x_gen)
if save_sampler_config:
self.save(output_dir)
if return_gen:
return torch.cat(x_gen_list, dim=0)