def test_autoencoder():
"""
Test that all components of the auto-encoder work correctly by executing a
training run against generated data.
"""
input_shape = (3, )
epochs = 1000
# Generate some data
x_train = np.random.rand(100, 3)
x_test = np.random.rand(30, 3)
# Define encoder and decoder model
def create_encoder_model(input_shape):
model_input = Input(shape=input_shape)
encoder = Dense(4)(model_input)
encoder = BatchNormalization()(encoder)
encoder = Activation(activation='relu')(encoder)
return Model(model_input, encoder)
def create_decoder_model(embedding_shape):
embedding_a = Input(shape=embedding_shape)
decoder = Dense(3)(embedding_a)
decoder = BatchNormalization()(decoder)
decoder = Activation(activation='relu')(decoder)
return Model(embedding_a, decoder)
# Create auto-encoder network
encoder_model = create_encoder_model(input_shape)
decoder_model = create_decoder_model(encoder_model.output_shape)
autoencoder = AutoEncoder(encoder_model, decoder_model)
# Prepare auto-encoder for training
autoencoder.compile(loss='binary_crossentropy', optimizer='adam')
# Evaluate network before training to establish a baseline
score_before = autoencoder.evaluate(x_train, x_train)
# Train network
autoencoder.fit(x_train,
x_train,
validation_data=(x_test, x_test),
epochs=epochs)
# Evaluate network
score_after = autoencoder.evaluate(x_train, x_train)
# Ensure that the training loss score improved as a result of the training
assert (score_before > score_after)
def train_autoencoder_and_log(
autoencoder: AutoEncoder,
train_loader: DataLoader,
test_loader: DataLoader,
) -> None:
"""Train the autoencoder and log data to disc."""
if os.path.exists(config.LOG_DIR):
shutil.rmtree(config.LOG_DIR)
os.mkdir(config.LOG_DIR)
train_losses = []
test_losses = []
test_images = _get_sample_digits(test_loader)
all_encodings = []
all_reconstructions = []
for i in tqdm(range(config.STEPS)):
# Train
autoencoder.train_step(steps=1)
# collect train and test losses
train_losses.append(autoencoder.evaluate(train_loader))
test_losses.append(autoencoder.evaluate(test_loader))
# collect and save train encodings
encodings = None
for images, _ in train_loader:
encodings_ = autoencoder.encode(images)
if encodings is None:
encodings = encodings_
else:
encodings = torch.cat((encodings, encodings_), 0)
all_encodings.append(encodings)
# collect the sample reconstructions
reconstructions = autoencoder.autoencode(test_images)
all_reconstructions.append(reconstructions)
def save(filename, object):
filename = os.path.join(config.LOG_DIR, filename)
np.save(filename, object)
# Format and save data
train_losses = np.array(train_losses)
save("train_losses.npy", train_losses)
test_losses = np.array(test_losses)
save("test_losses.npy", test_losses)
test_images = (test_images.numpy() + 1) / 2
save("test_images.npy", test_images)
all_encodings = np.array([x.numpy() for x in all_encodings])
save("encodings.npy", all_encodings)
all_reconstructions = np.array([x.numpy() for x in all_reconstructions])
all_reconstructions = (all_reconstructions + 1) / 2
save("reconstructions.npy", all_reconstructions)
labels = [x.numpy() for __, x in train_loader]
labels = np.array(labels, dtype=int).flatten()
save("labels.npy", labels)
