import unittest import math import torch import os import tempfile import uuid import shutil import pyprob from pyprob import util, Model, LearningRateScheduler, Optimizer from pyprob.distributions import Normal, Uniform class ModelTestCase(unittest.TestCase): def __init__(self, *args, **kwargs): # http://www.robots.ox.ac.uk/~fwood/assets/pdf/Wood-AISTATS-2014.pdf class GaussianWithUnknownMeanMarsaglia(Model): def __init__(self, prior_mean=1, prior_stddev=math.sqrt(5), likelihood_stddev=math.sqrt(2)): self.prior_mean = prior_mean self.prior_stddev = prior_stddev self.likelihood_stddev = likelihood_stddev super().__init__('Gaussian with unknown mean (Marsaglia)') def marsaglia(self, mean, stddev): uniform = Uniform(-1, 1) s = 1 while float(s) >= 1: x = pyprob.sample(uniform) y = pyprob.sample(uniform) s = x*x + y*y return mean + stddev * (x * torch.sqrt(-2 * torch.log(s) / s)) def forward(self): mu = self.marsaglia(self.prior_mean, self.prior_stddev) likelihood = Normal(mu, self.likelihood_stddev) pyprob.observe(likelihood, 0, name='obs0') pyprob.observe(likelihood, 0, name='obs1') return mu self._model = GaussianWithUnknownMeanMarsaglia() super().__init__(*args, **kwargs) def test_model_train_save_load_train(self): training_traces = 128 file_name = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) self._model.reset_inference_network() self._model.learn_inference_network(num_traces=training_traces, observe_embeddings={'obs0': {'dim': 64}, 'obs1': {'dim': 64}}) self._model.save_inference_network(file_name) self._model.load_inference_network(file_name) os.remove(file_name) self._model.learn_inference_network(num_traces=training_traces, observe_embeddings={'obs0': {'dim': 64}, 'obs1': {'dim': 64}}) self._model.save_inference_network(file_name) self._model.load_inference_network(file_name) os.remove(file_name) self._model.learn_inference_network(num_traces=training_traces, observe_embeddings={'obs0': {'dim': 64}, 'obs1': {'dim': 64}}) util.eval_print('training_traces', 'file_name') self.assertTrue(True) def test_model_save_traces_load_train(self): dataset_dir = tempfile.mkdtemp() num_traces = 512 num_traces_per_file = 32 training_traces = 128 self._model.save_dataset(dataset_dir=dataset_dir, num_traces=num_traces, num_traces_per_file=num_traces_per_file) self._model.reset_inference_network() self._model.learn_inference_network(num_traces=training_traces, dataset_dir=dataset_dir, batch_size=16, valid_size=16, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}) shutil.rmtree(dataset_dir) util.eval_print('dataset_dir', 'num_traces', 'num_traces_per_file', 'training_traces') self.assertTrue(True) def test_model_train(self): num_traces = 256 self._model.reset_inference_network() self._model.learn_inference_network(num_traces=num_traces, batch_size=16, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}) util.eval_print('num_traces') self.assertTrue(True) def test_model_train_lr_scheduler_poly1(self): num_traces = 256 self._model.reset_inference_network() self._model.learn_inference_network(num_traces=num_traces, batch_size=16, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, learning_rate_scheduler_type=LearningRateScheduler.POLY1) util.eval_print('num_traces') self.assertTrue(True) def test_model_train_lr_scheduler_poly2(self): num_traces = 256 self._model.reset_inference_network() self._model.learn_inference_network(num_traces=num_traces, batch_size=16, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, learning_rate_scheduler_type=LearningRateScheduler.POLY2) util.eval_print('num_traces') self.assertTrue(True) def test_model_train_online_adam_larc_lr_poly2(self): file_name = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) file_name_2 = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) file_name_3 = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) num_traces_end = 256 batch_size = 16 learning_rate_init_correct = 0.1 learning_rate_end_correct = 0.0025 self._model.reset_inference_network() print('Training\n') self._model.learn_inference_network(num_traces=num_traces_end/2, num_traces_end=num_traces_end, batch_size=batch_size, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, optimizer_type=Optimizer.ADAM_LARC, learning_rate_scheduler_type=LearningRateScheduler.POLY2, learning_rate_init=learning_rate_init_correct, learning_rate_end=learning_rate_end_correct, log_file_name=file_name_2) print('Saving\n') # print(self._model._inference_network._optimizer) optimizer_state_step_before_save = list(self._model._inference_network._optimizer.state_dict()['state'].values())[0]['step'] self._model.save_inference_network(file_name) print('\nlog_file_name contents') with open(file_name_2) as file: for line in file: print(line, end='') self._model.reset_inference_network() print('Loading\n') self._model.load_inference_network(file_name) # print(self._model._inference_network._optimizer) optimizer_state_step_after_load = list(self._model._inference_network._optimizer.state_dict()['state'].values())[0]['step'] print('Training\n') self._model.learn_inference_network(num_traces=num_traces_end/2, num_traces_end=num_traces_end, batch_size=batch_size, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, optimizer_type=Optimizer.ADAM_LARC, learning_rate_scheduler_type=LearningRateScheduler.POLY2, learning_rate_init=learning_rate_init_correct, learning_rate_end=learning_rate_end_correct, log_file_name=file_name_3) learning_rate_end = self._model._inference_network._optimizer.param_groups[0]['lr'] print('\nlog_file_name contents') with open(file_name_3) as file: for line in file: print(line, end='') os.remove(file_name) os.remove(file_name_2) os.remove(file_name_3) util.eval_print('file_name', 'file_name_2', 'file_name_3', 'num_traces_end', 'batch_size', 'learning_rate_init_correct', 'optimizer_state_step_before_save', 'optimizer_state_step_after_load', 'learning_rate_end', 'learning_rate_end_correct') # util.eval_print('file_name', 'file_name_2', 'file_name_3', 'num_traces_end', 'batch_size', 'learning_rate_init_correct', 'learning_rate_end', 'learning_rate_end_correct') self.assertEqual(optimizer_state_step_before_save, optimizer_state_step_after_load) self.assertAlmostEqual(learning_rate_end, learning_rate_end_correct, delta=learning_rate_end_correct) def test_model_train_offline_adam_larc_lr_poly2(self): dataset_dir = tempfile.mkdtemp() dataset_num_traces = 128 dataset_num_traces_per_file = 32 file_name = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) file_name_2 = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) file_name_3 = os.path.join(tempfile.mkdtemp(), str(uuid.uuid4())) num_traces_end = 256 batch_size = 16 learning_rate_init_correct = 0.1 learning_rate_end_correct = 0.0025 print('Saving dataset\n') self._model.save_dataset(dataset_dir=dataset_dir, num_traces=dataset_num_traces, num_traces_per_file=dataset_num_traces_per_file) self._model.reset_inference_network() print('Training\n') self._model.learn_inference_network(dataset_dir=dataset_dir, num_traces=num_traces_end/2, num_traces_end=num_traces_end, batch_size=batch_size, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, optimizer_type=Optimizer.ADAM_LARC, learning_rate_scheduler_type=LearningRateScheduler.POLY2, learning_rate_init=learning_rate_init_correct, learning_rate_end=learning_rate_end_correct, log_file_name=file_name_2) print('Saving\n') # print(self._model._inference_network._optimizer) optimizer_state_step_before_save = list(self._model._inference_network._optimizer.state_dict()['state'].values())[0]['step'] self._model.save_inference_network(file_name) print('\nlog_file_name contents') with open(file_name_2) as file: for line in file: print(line, end='') self._model.reset_inference_network() print('Loading\n') self._model.load_inference_network(file_name) # print(self._model._inference_network._optimizer) optimizer_state_step_after_load = list(self._model._inference_network._optimizer.state_dict()['state'].values())[0]['step'] print('Training\n') self._model.learn_inference_network(dataset_dir=dataset_dir, num_traces=num_traces_end/2, num_traces_end=num_traces_end, batch_size=batch_size, observe_embeddings={'obs0': {'dim': 16}, 'obs1': {'dim': 16}}, optimizer_type=Optimizer.ADAM_LARC, learning_rate_scheduler_type=LearningRateScheduler.POLY2, learning_rate_init=learning_rate_init_correct, learning_rate_end=learning_rate_end_correct, log_file_name=file_name_3) learning_rate_end = self._model._inference_network._optimizer.param_groups[0]['lr'] print('\nlog_file_name contents') with open(file_name_3) as file: for line in file: print(line, end='') os.remove(file_name) os.remove(file_name_2) os.remove(file_name_3) shutil.rmtree(dataset_dir) util.eval_print('dataset_dir', 'dataset_num_traces', 'dataset_num_traces_per_file', 'file_name', 'file_name_2', 'file_name_3', 'num_traces_end', 'batch_size', 'learning_rate_init_correct', 'optimizer_state_step_before_save', 'optimizer_state_step_after_load', 'learning_rate_end', 'learning_rate_end_correct') # util.eval_print('file_name', 'file_name_2', 'file_name_3', 'num_traces_end', 'batch_size', 'learning_rate_init_correct', 'learning_rate_end', 'learning_rate_end_correct') self.assertEqual(optimizer_state_step_before_save, optimizer_state_step_after_load) self.assertAlmostEqual(learning_rate_end, learning_rate_end_correct, delta=learning_rate_end_correct) if __name__ == '__main__': pyprob.set_random_seed(123) pyprob.set_verbosity(1) unittest.main(verbosity=2)