def test_synthetic(self):
epoch_size_list = [20, 50]
data_edge_size_list = [3, 8]
dist_type_list = ['gaussian', 'laplacian']
num_classes_list = [None, 10]
rand_state_list = [np.random.RandomState(12345)]
for epoch_size in epoch_size_list:
for data_edge_size in data_edge_size_list:
for dist_type in dist_type_list:
for num_classes in num_classes_list:
for rand_state in rand_state_list:
params = types.SimpleNamespace()
params.dataset = 'synthetic'
params.shuffle_data = True
params.epoch_size = epoch_size
params.batch_size = 10
params.data_edge_size = data_edge_size
params.dist_type = dist_type
params.num_classes = num_classes
params.rand_state = rand_state
train_loader, val_loader, test_loader, data_params = dataset_utils.load_dataset(params)
for key, value in data_params.items():
setattr(params, key, value)
assert len(train_loader.dataset) == epoch_size
for batch_idx, (data, target) in enumerate(train_loader):
assert data.numpy().shape == (params.batch_size, params.data_edge_size, params.data_edge_size, 1)
assert batch_idx + 1 == epoch_size // params.batch_size
def test_model_loading(self):
for model_type in self.model_list:
model_type = ''.join(
model_type.split('_')[:-1]) # remove '_model' at the end
model = loaders.load_model(model_type)
params = loaders.load_params(self.test_params_file,
key=model_type + '_params')
train_loader, val_loader, test_loader, data_params = datasets.load_dataset(
params)
for key, value in data_params.items():
setattr(params, key, value)
model.setup(params)
def test_mnist(self):
try: # only run the test if the dataset is already downloaded
mnist = datasets.MNIST(root=self.data_dir, train=True, download=False)
except:
return 0
standardize_data_list = [True, False]
rescale_data_list = [True, False]
for standardize_data in [True, False]:
for rescale_data_to_one in [True, False]:
params = types.SimpleNamespace()
params.standardize_data = standardize_data
params.rescale_data_to_one = rescale_data_to_one
if(params.standardize_data or params.rescale_data_to_one):
params.eps = 1e-8
params.data_dir = self.data_dir
params.dataset = 'mnist'
params.shuffle_data = True
params.batch_size = 10000
train_loader, val_loader, test_loader, data_params = dataset_utils.load_dataset(params)
for key, value in data_params.items():
setattr(params, key, value)
assert len(train_loader.dataset) == params.epoch_size
(data, target) = next(iter(train_loader))
assert data.numpy().shape == (params.batch_size, 28, 28, 1)
0.25,
0.3,
0.35,
#0.4,
0.5,
#0.8,
1.0
]
num_models = len(log_files)
for model_index in range(num_models):
logger = Logger(log_files[model_index], overwrite=False)
log_text = logger.load_file()
params = logger.read_params(log_text)[-1]
params.cp_latest_filename = cp_latest_filenames[model_index]
train_loader, val_loader, test_loader, data_params = dataset_utils.load_dataset(
params)
for key, value in data_params.items():
setattr(params, key, value)
model = loaders.load_model(params.model_type)
model.setup(params, logger)
model.params.analysis_out_dir = os.path.join(
*[model.params.model_out_dir, 'analysis', model.params.version])
model.params.analysis_save_dir = os.path.join(
model.params.analysis_out_dir, 'savefiles')
if not os.path.exists(model.params.analysis_save_dir):
os.makedirs(model.params.analysis_save_dir)
model.to(params.device)
model.load_checkpoint()
fmodel = PyTorchModel(model.eval(), bounds=(0, 1))
print('\n', '~' * 79)
num_batches = len(test_loader.dataset) // model.params.batch_size
def test_lca_ensemble_gradients(self):
params = {}
models = {}
params['lca'] = loaders.load_params(self.test_params_file,
key='lca_params')
params['lca'].train_logs_per_epoch = None
params['lca'].shuffle_data = False
train_loader, val_loader, test_loader, data_params = datasets.load_dataset(
params['lca'])
for key, value in data_params.items():
setattr(params['lca'], key, value)
models['lca'] = loaders.load_model(params['lca'].model_type)
models['lca'].setup(params['lca'])
models['lca'].to(params['lca'].device)
params['ensemble'] = loaders.load_params(self.test_params_file,
key='ensemble_params')
for key, value in data_params.items():
setattr(params['ensemble'], key, value)
err_msg = f'\ndata_shape={params["ensemble"].data_shape}'
err_msg += f'\nnum_pixels={params["ensemble"].num_pixels}'
err_msg += f'\nbatch_size={params["ensemble"].batch_size}'
err_msg += f'\nepoch_size={params["ensemble"].epoch_size}'
models['ensemble'] = loaders.load_model(params['ensemble'].model_type)
models['ensemble'].setup(params['ensemble'])
models['ensemble'].to(params['ensemble'].device)
ensemble_state_dict = models['ensemble'].state_dict()
ensemble_state_dict['lca.w'] = models['lca'].w.clone()
models['ensemble'].load_state_dict(ensemble_state_dict)
data, target = next(iter(train_loader))
train_data_batch = models['lca'].preprocess_data(
data.to(params['lca'].device))
train_target_batch = target.to(params['lca'].device)
models['lca'].optimizer.zero_grad()
for submodel in models['ensemble']:
submodel.optimizer.zero_grad()
inputs = [train_data_batch] # only the first model acts on input
for submodel in models['ensemble']:
inputs.append(submodel.get_encodings(inputs[-1]).detach())
lca_loss = models['lca'].get_total_loss(
(train_data_batch, train_target_batch))
ensemble_losses = [
models['ensemble'].get_total_loss((inputs[0], train_target_batch),
0)
]
ensemble_losses.append(models['ensemble'].get_total_loss(
(inputs[1], train_target_batch), 1))
lca_loss.backward()
ensemble_losses[0].backward()
ensemble_losses[1].backward()
lca_loss_val = lca_loss.cpu().detach().numpy()
lca_w_grad = models['lca'].w.grad.cpu().numpy()
ensemble_loss_val = ensemble_losses[0].cpu().detach().numpy()
ensemble_w_grad = models['ensemble'][0].w.grad.cpu().numpy()
assert lca_loss_val == ensemble_loss_val, (
err_msg + '\n' +
'Losses should be equal, but are lca={lca_loss_val} and ensemble={ensemble_loss_val}'
)
assert np.all(lca_w_grad == ensemble_w_grad), (
err_msg + '\nGrads should be equal, but are not.')
lca_pre_train_w = models['lca'].w.cpu().detach().numpy().copy()
ensemble_pre_train_w = models['ensemble'][0].w.cpu().detach().numpy(
).copy()
run_utils.train_epoch(1, models['lca'], train_loader)
run_utils.train_epoch(1, models['ensemble'], train_loader)
lca_w = models['lca'].w.cpu().detach().numpy().copy()
ensemble_w = models['ensemble'][0].w.cpu().detach().numpy().copy()
assert np.all(lca_pre_train_w == ensemble_pre_train_w), (
err_msg + '\n' +
"lca & ensemble weights are not equal before one epoch of training"
)
assert not np.all(lca_pre_train_w == lca_w), (
err_msg + '\n' +
"lca weights are not different from init after one epoch of training"
)
assert not np.all(ensemble_pre_train_w == ensemble_w), (
err_msg + '\n' +
"ensemble weights are not different from init after one epoch of training"
)
assert np.all(lca_w == ensemble_w), (
err_msg + '\n' +
"lca & ensemble weights are not equal after one epoch of training")