def test_variational_autoencoder():
# Build data and mlmodel
data_name = "adult"
data = OnlineCatalog(data_name)
model = MLModelCatalog(data, "ann", backend="pytorch")
device = "cuda" if torch.cuda.is_available() else "cpu"
test_input = np.zeros((1, 13))
test_input = torch.Tensor(test_input).to(device)
vae = VariationalAutoencoder(
data_name, layers=[11, 512, 256, 8], mutable_mask=model.get_mutable_mask()
)
vae.fit(xtrain=data.df[model.feature_input_order])
test_reconstructed, _, _ = vae.predict(test_input)
assert test_reconstructed.shape == test_input.shape
# test loading vae
new_vae = VariationalAutoencoder(
data_name, layers=[11, 512, 256, 8], mutable_mask=model.get_mutable_mask()
)
new_vae.load(11)
def test_cs_vae():
# Build data and mlmodel
data_name = "adult"
data = OnlineCatalog(data_name)
model = MLModelCatalog(data, "ann", backend="pytorch")
test_input = np.zeros((1, 13))
test_input = torch.Tensor(test_input)
test_class = torch.Tensor(np.array([[0, 0]]))
csvae = CSVAE(data_name, layers=[11, 16, 8], mutable_mask=model.get_mutable_mask())
csvae.fit(data=data.df[model.feature_input_order + [data.target]], epochs=1)
output = csvae.predict(test_input, test_class)
test_reconstructed = output[0]
assert test_reconstructed.shape == test_input.shape
# test loading vae
new_csvae = CSVAE(
data_name, layers=[11, 16, 8], mutable_mask=model.get_mutable_mask()
)
new_csvae.load(11)