class TestVAEHVI(unittest.TestCase):
def setUp(self) -> None:
self._device = gpu_device
self._model = VAEHVI()
self._data_module = MNISTDataModule()
# noinspection PyTypeChecker
def test_training(self):
trainer = pl.Trainer(max_epochs=1, gpus=1)
trainer.fit(self._model, self._data_module)
def test_testing(self):
trainer = pl.Trainer(max_epochs=1, gpus=1)
self._data_module.setup("test")
trainer.test(self._model, self._data_module.test_dataloader())
def test_show_vae_reconstruction(self):
show_vae_reconstruction(self._model,
self._data_module,
dpi=150,
figsize=(2.5, 6))
plt.savefig("tmp/reconstruction.png")
def test_show_vae_generation(self):
show_vae_generation(self._model, dpi=150, figsize=(2.5, 6))
plt.savefig("tmp/generation.png")
def setUp(self) -> None:
self._device = gpu_device
self._model = VAEHVI()
self._data_module = MNISTDataModule()
def setUp(self) -> None:
self._device = gpu_device
self._model = VaeGaussian()
self._data_module = MNISTDataModule()
return avo.models.VaeGaussian(encoder_config={"batch_norm": True},
beta=0.5,
gamma=1 / 80.)
elif model_type == "vae_hvi":
return avo.models.VAEHVI(encoder_config={"batch_norm": True},
beta=0.5,
gamma=1 / 80.)
elif model_type == "vae_hvi_avo":
return avo.models.VAEHVIAVO(encoder_config={"batch_norm": True},
beta=0.5,
gamma=1 / 80.)
arguments = parser.parse_args()
if arguments.command == "train":
data_module = MNISTDataModule(batch_size=256)
model = make_model(arguments.model)
trainer = pl.Trainer(max_epochs=arguments.epoch,
gpus=1,
progress_bar_refresh_rate=40)
trainer.fit(model, data_module)
elif arguments.command == "test":
model = make_model(arguments.model)
model.load_state_dict(torch.load(arguments.model_checkpoint)["state_dict"])
data_module = MNISTDataModule(batch_size=256)
show_vae_reconstruction(model, data_module, dpi=150, figsize=(2.5, 6))
plt.savefig("reconstruction.png")
show_vae_generation(model, dpi=200, figsize=(2.5, 4))
plt.savefig("generation.png")
trainer = pl.Trainer(max_epochs=1, gpus=1)
data_module.setup("test")