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())