def test_spec2im_filters(efilt, dfilt, filters_e_sorted, filters_d_sorted): from atomai.nets import SignalDecoder, SignalEncoder X_train, X_test = gen_spectra() y_train, y_test = gen_image_data() in_dim, out_dim = (16,), (8, 8) t = ImSpecTrainer( in_dim, out_dim, nbfilters_encoder=efilt, nbfilters_decoder=dfilt) t.compile_trainer((X_train, y_train, X_test, y_test), batch_size=4, loss="mse", training_cycles=1) filters_e, filters_d = [], [] for child in t.net.children(): if isinstance(child, SignalEncoder): filt_ = [] for p in child.state_dict().values(): if len(p.shape) == 3: filt_.append(p.shape[0]) filters_e.extend(filt_) elif isinstance(child, SignalDecoder): filt_ = [] for p in child.state_dict().values(): if len(p.shape) == 4: filt_.append(p.shape[0]) filters_d.extend(filt_) assert_equal(sorted(np.unique(filters_e).tolist()), filters_e_sorted) assert_equal(sorted(np.unique(filters_d).tolist()), filters_d_sorted)
def test_init_imspec_bn(bn, encoder_layers, decoder_layers, n_bn): X_train, X_test = gen_image_data() y_train, y_test = gen_spectra() in_dim, out_dim = (8, 8), (16,) t = ImSpecTrainer(in_dim, out_dim, nblayers_encoder=encoder_layers, nblayers_decoder=decoder_layers, batch_norm=bn) t.compile_trainer((X_train, y_train, X_test, y_test), batch_size=4, loss="mse", training_cycles=1) n_bn_ = len([k for k in t.net.state_dict().keys() if 'running_mean' in k]) assert_equal(n_bn_, n_bn)
def test_init_spec2im_dupsample(upsample, dim_): from atomai.utils import Hook, mock_forward X_train, X_test = gen_spectra() y_train, y_test = gen_image_data() in_dim, out_dim = (16,), (8, 8) t = ImSpecTrainer(in_dim, out_dim, decoder_upsampling=upsample) t.compile_trainer((X_train, y_train, X_test, y_test), batch_size=4, loss="mse", training_cycles=1) hookF = [Hook(layer[1]) for layer in list(t.net.decoder._modules.items())] mock_forward(t.net, dims=(1, 16)) assert_equal([h.output.shape[-1] for h in hookF][1:4], dim_)
def test_init_im2spec_edownsample(downsample_factor, output_dim): from atomai.utils import Hook, mock_forward X_train, X_test = gen_image_data() y_train, y_test = gen_spectra() in_dim, out_dim = (8, 8), (16,) t = ImSpecTrainer(in_dim, out_dim, encoder_downsampling=downsample_factor) t.compile_trainer((X_train, y_train, X_test, y_test), batch_size=4, loss="mse", training_cycles=1) hookF = [Hook(layer[1]) for layer in list(t.net.encoder._modules.items())] mock_forward(t.net, dims=(1, 8, 8)) assert_equal(tuple([h.output.shape[-2:] for h in hookF][0]), output_dim)
def test_spec2imtrainer_dataloader(): X_train, X_test = gen_spectra() y_train, y_test = gen_image_data() in_dim, out_dim = (16,), (8, 8) t = ImSpecTrainer(in_dim, out_dim, seed=1) t.compile_trainer((X_train, y_train, X_test, y_test), loss="mse", training_cycles=5, batch_size=4) X_train_, y_train_ = t.dataloader(0) assert_equal(X_train_.dtype, torch.float32) assert_equal(y_train_.dtype, torch.float32) assert_equal(len(X_train), len(y_train)) assert_equal(X_train_.is_cuda, torch.cuda.is_available())
def test_im2spec_trainer_determinism(latent_dim): X_train, X_test = gen_image_data() y_train, y_test = gen_spectra() in_dim = (8, 8) out_dim = (16,) t1 = ImSpecTrainer(in_dim, out_dim, latent_dim, seed=1) t1.compile_trainer((X_train, y_train, X_test, y_test), loss="mse", training_cycles=5, batch_size=4) _ = t1.run() loss1 = t1.loss_acc["train_loss"][-1] t2 = ImSpecTrainer(in_dim, out_dim, latent_dim, seed=1) t2.compile_trainer((X_train, y_train, X_test, y_test), loss="mse", training_cycles=5, batch_size=4) _ = t2.run() loss2 = t2.loss_acc["train_loss"][-1] assert_allclose(loss1, loss2) for p1, p2 in zip(t1.net.parameters(), t2.net.parameters()): assert_allclose(p1.detach().cpu().numpy(), p2.detach().cpu().numpy())