def create_model(self, input_dim, autoenc_model):
pre_trained_autoenc = keras.models.load_model(
f'hotexamples_com/models/trained_models/{autoenc_model}')
left_autoencoder = Autoencoder(input_dim)
left_input = left_autoencoder.model.input
left_autoencoder.model.set_weights(pre_trained_autoenc.get_weights())
# Share weights for left and right autoencoder
right_autoencoder = Autoencoder(input_dim, name='right')
right_input = right_autoencoder.model.input
right_autoencoder.model.set_weights(pre_trained_autoenc.get_weights())
left_embed_layer = left_autoencoder.model.layers[-2].output
right_embed_layer = right_autoencoder.model.layers[-2].output
merge_layer = keras.layers.Concatenate()(
[left_embed_layer, right_embed_layer])
dnn_layer = keras.layers.Dense(100, activation='relu')(merge_layer)
dnn_layer = keras.layers.Dense(100, activation='relu')(dnn_layer)
output = keras.layers.Dense(2, activation='softmax')(dnn_layer)
model = keras.models.Model(inputs=[left_input, right_input],
outputs=output)
print(model.summary())
model.compile(optimizer=keras.optimizers.Adam(),
loss=keras.losses.BinaryCrossentropy(),
metrics=['accuracy'])
return model
def __init__(self, *args):
super().__init__(*args)
self.model = Autoencoder(encoder=self.encode,
shape=self.shape,
beta=1.)
if len(self.prior):
self.model.fit(*zip(*self.prior.items()),
epochs=initial_epochs)
def train_heinsfeld_autoencoder():
X, y = read_data('rois_cc200')
clf = Autoencoder(num_classes=2,
dropout=(0.6, 0.8),
learning_rate=(0.0001, 0.0001, 0.0005),
momentum=0.9,
noise=(0.2, 0.3),
batch_size=(100, 10, 10),
num_epochs=(700, 2000, 100))
clf.train(X, y)
def __init__(self, reconstruction_loss_factor: float, cycle_loss_factor: float):
# share weights of the upper encoder & lower decoder
encoder_upper, decoder_lower = UpperEncoder(), LowerDecoder()
self.ae_day = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower, UpperDecoder())
self.ae_night = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower, UpperDecoder())
self.loss_fn = nn.L1Loss()
self.reconstruction_loss_factor = reconstruction_loss_factor
self.cycle_loss_factor = cycle_loss_factor
self.optimizer = None
self.scheduler = None
def predict_heinsfeld_autoencoder():
trn_x, trn_y = read_data('rois_cc200')
tst_x, tst_y = read_data('rois_cc200', training=False)
clf = Autoencoder(num_classes=2,
dropout=(0.6, 0.8),
learning_rate=(0.0001, 0.0001, 0.0005),
momentum=0.9,
noise=(0.2, 0.3),
batch_size=(100, 10, 10),
num_epochs=(700, 2000, 100))
clf.predict(trn_x, trn_y, tst_x, tst_y)
def __init__(self):
encoder_upper, decoder_lower = UpperEncoder(), LowerDecoder()
self.ae_day = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower,
UpperDecoder())
self.ae_night = Autoencoder(LowerEncoder(), encoder_upper,
decoder_lower, UpperDecoder())
self.loss_fn = nn.L1Loss()
self.optimizer_day = None
self.optimizer_night = None
self.scheduler_day = None
self.scheduler_night = None
def __init__(self, params: dict):
self.params = params
# share weights of the upper encoder & lower decoder
encoder_upper, decoder_lower = UpperEncoder(), LowerDecoder()
self.ae_day = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower,
UpperDecoder())
self.ae_night = Autoencoder(LowerEncoder(), encoder_upper,
decoder_lower, UpperDecoder())
self.reconst_loss = nn.L1Loss()
self.optimizer = None
self.scheduler = None
def test_autoencoder_training(self):
# create a one sample dataframe for the test
train_ds = TensorDataset(self.input.squeeze(0), self.input.squeeze(0)) # need two dimensional (1, 7) shape for input
train_dl = DataLoader(train_ds, batch_size=1, shuffle=True)
# create models
embed_model = PretrainedTransformerGenerator(self.args)
decoder = GRUDecoder(embed_model.config.d_model, self.tokenizer.vocab_size, embed_model.config.d_model, n_layers=1, dropout=0)
decoder = decoder.to(self.args.device) # warning: device is cpu for CI, slow
autoencoder = Autoencoder(embed_model, decoder, self.args.device, tokenizer=self.tokenizer).to(self.args.device)
# create needed params
autoencoder_optimizer = optim.Adam(autoencoder.parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss(ignore_index=0)
loss_df = pd.DataFrame(columns=['batch_num', 'loss'])
# see if it works
autoencoder = train_autoencoder(self.args, autoencoder, train_dl, train_dl, autoencoder_optimizer, criterion, 1, loss_df, num_epochs=2)
class Agent(agents.random.RandomAgent(epochs)):
def __init__(self, *args):
super().__init__(*args)
self.model = Autoencoder(encoder=self.encode,
shape=self.shape,
beta=1.)
def act(self, seqs):
return list(
zip(*sorted(zip(self.model.predict(seqs), seqs))
[-self.batch:]))[1]
def observe(self, data):
super().observe(data)
self.model.fit(*zip(*self.seen.items()), epochs=epochs)
def load_pretrained_model(args):
"""
Load the pretrained model
:param args: Command line arguments passed to this file, including class of pretrained model
:return: A pretrained model object
"""
pretrained_ckpt = torch.load(args.pretrained_model_path, map_location=config.device)
if args.encoder_pruned:
pretrained_state_dict = pretrained_ckpt
else:
pretrained_state_dict = pretrained_ckpt['state_dict']
# Create pretrained model object
# For 'framework' concept, later more model selection
if args.model is Model.VGG.value:
pretrained_model = VGGNet(hparams=hparams)
elif args.model is Model.Autoencoder.value:
if args.encoder_pruned:
# Pretrained model's encoder is already pruned
params = config.compute_pruned_autoencoder_params(config.remove_ratio, encoder=True, decoder=False)
else:
# Pretrained model is unpruned
params = config.autoencoder_params
pretrained_model = Autoencoder(hparams=hparams, model_params=params)
pretrained_model.load_state_dict(pretrained_state_dict)
return pretrained_model
def test_can_decode_basic_more_layers(self):
embed_model = PretrainedTransformerGenerator(self.args)
gru_decoder = GRUDecoder(embed_model.config.d_model, self.tokenizer.vocab_size, embed_model.config.d_model, n_layers=4, dropout=.2)
autoencoder = Autoencoder(embed_model, gru_decoder, "cpu:0").to("cpu:0")
output = autoencoder(self.input, self.target)
criterion = nn.CrossEntropyLoss(ignore_index=0)
output = output.permute((1, 2, 0)) # swap for loss
loss = criterion(output, self.target)
assert type(loss.item()) == float, "could not get loss value: type {}".format(type(loss.item()))
def test_decode_to_text(self):
embed_model = PretrainedTransformerGenerator(self.args)
gru_decoder = GRUDecoder(embed_model.config.d_model, self.tokenizer.vocab_size, embed_model.config.d_model, n_layers=4, dropout=.2)
autoencoder = Autoencoder(embed_model, gru_decoder, "cpu:0").to("cpu:0")
output = autoencoder(self.input, self.target)
# get text output
_, best_guess = torch.max(output, dim=2)
predicted = self.tokenizer.convert_ids_to_tokens(best_guess.permute(1, 0).flatten().tolist())
string_pred = " ".join(predicted)
print('Predicted: ', string_pred)
print('Actual: ', self.tokenizer.convert_ids_to_tokens(self.input.flatten().tolist()))
assert type(string_pred) == str, "predicted value was not a string, was a {}".format(type(string_pred))
def main(unused_argv):
# load test images
test_list = list_image(FLAGS.test_folder)
# load model
assert (FLAGS.snapshot_dir != ""
or FLAGS.model_fname != ""), 'No pretrained model specified'
model = Autoencoder(cfgs.patch_size * cfgs.patch_size, cfgs, log_dir=None)
snapshot_fname = FLAGS.model_fname if FLAGS.model_fname != "" \
else tf.train.latest_checkpoint(FLAGS.snapshot_dir)
model.restore(snapshot_fname)
print('Restored from %s' % snapshot_fname)
sum_psnr = 0.0
stride = FLAGS.stride
for img_fname in test_list:
orig_img = load_image('%s/%s' % (FLAGS.test_folder, img_fname))
# pre-process image
gray_img = toGrayscale(orig_img)
img = gray_img.astype(np.float32)
img -= cfgs.mean_value
img *= cfgs.scale
# make measurement and reconstruct image
recon_img = overlap_inference(model,
img,
bs=cfgs.batch_size,
stride=stride)
recon_img /= cfgs.scale
recon_img += cfgs.mean_value
# save reconstruction
cv.imwrite(
'%s/%sOI_%d_%s' %
(FLAGS.reconstruction_folder, FLAGS.prefix, stride, img_fname),
recon_img.astype(np.uint8))
psnr_ = psnr(gray_img.astype(np.float32), recon_img)
print('Image %s, psnr: %f' % (img_fname, psnr_))
sum_psnr += psnr_
mean_psnr = sum_psnr / len(test_list)
print('---------------------------')
print('Mean PSNR: %f' % mean_psnr)
def __init__(self, config):
"""
Construct a new GAN trainer
:param Config config: The parsed network configuration.
"""
self.config = config
LOG.info("CUDA version: {0}".format(version.cuda))
LOG.info("Creating data loader from path {0}".format(config.FILENAME))
self.data_loader = Data(
config.FILENAME,
config.BATCH_SIZE,
polarisations=config.POLARISATIONS, # Polarisations to use
frequencies=config.FREQUENCIES, # Frequencies to use
max_inputs=config.
MAX_SAMPLES, # Max inputs per polarisation and frequency
normalise=config.NORMALISE) # Normalise inputs
shape = self.data_loader.get_input_shape()
width = shape[1]
LOG.info("Creating models with input shape {0}".format(shape))
self._autoencoder = Autoencoder(width)
self._discriminator = Discriminator(width)
# TODO: Get correct input and output widths for generator
self._generator = Generator(width, width)
if config.USE_CUDA:
LOG.info("Using CUDA")
self.autoencoder = self._autoencoder.cuda()
self.discriminator = self._discriminator.cuda()
self.generator = self._generator.cuda()
else:
LOG.info("Using CPU")
self.autoencoder = self._autoencoder
self.discriminator = self._discriminator
self.generator = self._generator
def train_autoencoder():
# Instantiate the model
# Normalization function: We will scale the input image pixels within 0-1 range by dividing all input value by 255.
autoencoder_definition = Autoencoder(input_dim = features, num_output_classes = num_output_classes, transformation = normalization)
autoencoder_model = autoencoder_definition.create_autoencoder()
reader_train = create_reader(train_file, True, input_dim, num_output_classes)
# Train Autoencoder
# Map the data streams to the input.
# Instantiate the loss and error function.
loss_function = mse(autoencoder_model, normalization(features))
error_function = mse(autoencoder_model, normalization(features))
input_map={
features : reader_train.streams.features
}
train(reader=reader_train, model=autoencoder_model, loss_function=loss_function, error_function=error_function, input_map=input_map,
num_sweeps_to_train_with = 100, num_samples_per_sweep = 2000, minibatch_size = 10, learning_rate = 0.02)
autoencoder_model.save('autoencoder.model')
return autoencoder_definition
def __init__(self,
encoder,
dim,
shape,
beta=0.,
alpha=5e-4,
zeta=1e-2,
lam=1e-6,
mu=0.5,
itr=200,
M=1000,
eps=1e-4,
minibatch=100,
gpbatch=2000):
'''encoder: convert sequences to one-hot arrays.
alpha: embedding learning rate.
zeta: induced point ascent learning rate
shape: sequence shape (len, channels).
beta: embedding score weighting.
dim: embedding dimensionality.
lam: l2 regularization constant.
mu: GP prior mean.
M: max number of induced points.
itr: gradient ascent iterations for induced pseudo-inputs.
eps: numerical stability
'''
super().__init__()
self.X, self.Y = (), ()
self.minibatch = gpbatch
self.embed = Autoencoder(encoder,
dim=dim,
alpha=alpha,
shape=shape,
lam=lam,
beta=beta,
minibatch=minibatch)
self.mu = mu
self.dim = dim
self.alpha = alpha
self.itr = itr
self.eps = eps
self.M = M
self.zeta = zeta
def main(unused_argv):
val_losses = []
assert FLAGS.output_dir, "--output_dir is required"
# Create training directory.
output_dir = FLAGS.output_dir
if not tf.gfile.IsDirectory(output_dir):
tf.gfile.MakeDirs(output_dir)
dl = DataLoader(FLAGS.db_fname, mean=cfgs.mean_value, scale=cfgs.scale, n_vals=FLAGS.n_vals)
dl.prepare()
x_dim = dl.get_data_dim()
model = Autoencoder(x_dim, cfgs, log_dir=FLAGS.log_dir)
model.quantize_weights()
txt_log_fname = FLAGS.log_dir + 'text_log.txt'
log_fout = open(txt_log_fname, 'w')
if FLAGS.pretrained_fname:
try:
log_train(log_fout, 'Resume from %s' %(FLAGS.pretrained_fname))
model.restore(FLAGS.pretrained_fname)
except:
log_train(log_fout, 'Cannot restore from %s' %(FLAGS.pretrained_fname))
pass
lr = cfgs.initial_lr
epoch_counter = 0
ite = 0
while True:
start = time.time()
x, flag = dl.next_batch(cfgs.batch_size, 'train')
load_data_time = time.time() - start
if flag:
epoch_counter += 1
do_log = (ite % FLAGS.log_every_n_steps == 0) or flag
do_snapshot = flag and epoch_counter > 0 and epoch_counter % FLAGS.save_every_n_epochs == 0
val_loss = -1
# train one step
start = time.time()
loss, _, summary, ite = model.partial_fit(x, lr, do_log)
one_iter_time = time.time() - start
# writing outs
if do_log:
log_train(log_fout, 'Iteration %d, (lr=%f) training loss : %f' %(ite, lr, loss))
if FLAGS.log_time:
log_train(log_fout, 'Iteration %d, data loading: %f(s) ; one iteration: %f(s)'
%(ite, load_data_time, one_iter_time))
model.log(summary)
if flag:
val_loss = val(model, dl)
val_losses.append(val_loss)
log_train(log_fout, '----------------------------------------------------')
if ite == 0:
log_train(log_fout, 'Initial validation loss: %f' %(val_loss))
else:
log_train(log_fout, 'Epoch %d, validation loss: %f' %(epoch_counter, val_loss))
log_train(log_fout, '----------------------------------------------------')
model.log(summary)
if do_snapshot:
log_train(log_fout, 'Snapshotting')
model.save(FLAGS.output_dir)
if flag:
if cfgs.lr_update == 'val' and len(val_losses) >= 5 and val_loss >= max(val_losses[-5:-1]):
lr = lr * cfgs.lr_decay_factor
log_train(log_fout, 'Decay learning rate to %f' %lr)
elif cfgs.lr_update == 'step' and epoch_counter % cfgs.num_epochs_per_decay == 0:
lr = lr * cfgs.lr_decay_factor
log_train(log_fout, 'Decay learning rate to %f' %lr)
if epoch_counter == FLAGS.n_epochs:
if not do_snapshot:
log_train(log_fout, 'Final snapshotting')
model.save(FLAGS.output_dir)
break
log_fout.close()
class Train(object):
"""
Main GAN trainer. Responsible for training the GAN and pre-training the generator autoencoder.
"""
def __init__(self, config):
"""
Construct a new GAN trainer
:param Config config: The parsed network configuration.
"""
self.config = config
LOG.info("CUDA version: {0}".format(version.cuda))
LOG.info("Creating data loader from path {0}".format(config.FILENAME))
self.data_loader = Data(
config.FILENAME,
config.BATCH_SIZE,
polarisations=config.POLARISATIONS, # Polarisations to use
frequencies=config.FREQUENCIES, # Frequencies to use
max_inputs=config.
MAX_SAMPLES, # Max inputs per polarisation and frequency
normalise=config.NORMALISE) # Normalise inputs
shape = self.data_loader.get_input_shape()
width = shape[1]
LOG.info("Creating models with input shape {0}".format(shape))
self._autoencoder = Autoencoder(width)
self._discriminator = Discriminator(width)
# TODO: Get correct input and output widths for generator
self._generator = Generator(width, width)
if config.USE_CUDA:
LOG.info("Using CUDA")
self.autoencoder = self._autoencoder.cuda()
self.discriminator = self._discriminator.cuda()
self.generator = self._generator.cuda()
else:
LOG.info("Using CPU")
self.autoencoder = self._autoencoder
self.discriminator = self._discriminator
self.generator = self._generator
def check_requeue(self, epochs_complete):
"""
Check and re-queue the training script if it has completed the desired number of training epochs per session
:param int epochs_complete: Number of epochs completed
:return: True if the script has been requeued, False if not
:rtype bool
"""
if self.config.REQUEUE_EPOCHS > 0:
if epochs_complete >= self.config.REQUEUE_EPOCHS:
# We've completed enough epochs for this instance. We need to kill it and requeue
LOG.info(
"REQUEUE_EPOCHS of {0} met, calling REQUEUE_SCRIPT".format(
self.config.REQUEUE_EPOCHS))
subprocess.call(self.config.REQUEUE_SCRIPT,
shell=True,
cwd=os.path.dirname(
self.config.REQUEUE_SCRIPT))
return True # Requeue performed
return False # No requeue needed
def load_state(self, checkpoint, module, optimiser=None):
"""
Load the provided checkpoint into the provided module and optimiser.
This function checks whether the load threw an exception and logs it to the user.
:param Checkpoint checkpoint: The checkpoint to load
:param module: The pytorch module to load the checkpoint into.
:param optimiser: The pytorch optimiser to load the checkpoint into.
:return: None if the load failed, int number of epochs in the checkpoint if load succeeded
"""
try:
module.load_state_dict(checkpoint.module_state)
if optimiser is not None:
optimiser.load_state_dict(checkpoint.optimiser_state)
return checkpoint.epoch
except RuntimeError as e:
LOG.exception(
"Error loading module state. This is most likely an input size mismatch. Please delete the old module saved state, or change the input size"
)
return None
def close(self):
"""
Close the data loader used by the trainer.
"""
self.data_loader.close()
def generate_labels(self, num_samples, pattern):
"""
Generate labels for the discriminator.
:param int num_samples: Number of input samples to generate labels for.
:param list pattern: Pattern to generator. Should be either [1, 0], or [0, 1]
:return: New labels for the discriminator
"""
var = torch.FloatTensor([pattern] * num_samples)
return var.cuda() if self.config.USE_CUDA else var
def _train_autoencoder(self):
"""
Main training loop for the autencoder.
This function will return False if:
- Loading the autoencoder succeeded, but the NN model did not load the state dicts correctly.
- The script needs to be re-queued because the NN has been trained for REQUEUE_EPOCHS
:return: True if training was completed, False if training needs to continue.
:rtype bool
"""
criterion = nn.SmoothL1Loss()
optimiser = optim.Adam(self.generator.parameters(),
lr=0.00003,
betas=(0.5, 0.999))
checkpoint = Checkpoint("autoencoder")
epoch = 0
if checkpoint.load():
epoch = self.load_state(checkpoint, self.autoencoder, optimiser)
if epoch is not None and epoch >= self.config.MAX_AUTOENCODER_EPOCHS:
LOG.info("Autoencoder already trained")
return True
else:
LOG.info(
"Autoencoder training beginning from epoch {0}".format(
epoch))
else:
LOG.info('Autoencoder checkpoint not found. Training from start')
# Train autoencoder
self._autoencoder.set_mode(Autoencoder.Mode.AUTOENCODER)
vis_path = os.path.join(
os.path.splitext(self.config.FILENAME)[0], "autoencoder",
str(datetime.now()))
with Visualiser(vis_path) as vis:
epochs_complete = 0
while epoch < self.config.MAX_AUTOENCODER_EPOCHS:
if self.check_requeue(epochs_complete):
return False # Requeue needed and training not complete
for step, (data, _, _) in enumerate(self.data_loader):
if self.config.USE_CUDA:
data = data.cuda()
if self.config.ADD_DROPOUT:
# Drop out parts of the input, but compute loss on the full input.
out = self.autoencoder(nn.functional.dropout(
data, 0.5))
else:
out = self.autoencoder(data)
loss = criterion(out.cpu(), data.cpu())
self.autoencoder.zero_grad()
loss.backward()
optimiser.step()
vis.step_autoencoder(loss.item())
# Report data and save checkpoint
fmt = "Epoch [{0}/{1}], Step[{2}/{3}], loss: {4:.4f}"
LOG.info(
fmt.format(epoch + 1,
self.config.MAX_AUTOENCODER_EPOCHS, step,
len(self.data_loader), loss))
epoch += 1
epochs_complete += 1
checkpoint.set(self.autoencoder.state_dict(),
optimiser.state_dict(), epoch).save()
LOG.info("Plotting autoencoder progress")
vis.plot_training(epoch)
data, _, _ = iter(self.data_loader).__next__()
vis.test_autoencoder(epoch, self.autoencoder, data.cuda())
LOG.info("Autoencoder training complete")
return True # Training complete
def _train_gan(self):
"""
TODO: Add in autoencoder to perform dimensionality reduction on data
TODO: Not working yet - trying to work out good autoencoder model first
:return:
"""
criterion = nn.BCELoss()
discriminator_optimiser = optim.Adam(self.discriminator.parameters(),
lr=0.003,
betas=(0.5, 0.999))
discriminator_scheduler = optim.lr_scheduler.LambdaLR(
discriminator_optimiser, lambda epoch: 0.97**epoch)
discriminator_checkpoint = Checkpoint("discriminator")
discriminator_epoch = 0
if discriminator_checkpoint.load():
discriminator_epoch = self.load_state(discriminator_checkpoint,
self.discriminator,
discriminator_optimiser)
else:
LOG.info('Discriminator checkpoint not found')
generator_optimiser = optim.Adam(self.generator.parameters(),
lr=0.003,
betas=(0.5, 0.999))
generator_scheduler = optim.lr_scheduler.LambdaLR(
generator_optimiser, lambda epoch: 0.97**epoch)
generator_checkpoint = Checkpoint("generator")
generator_epoch = 0
if generator_checkpoint.load():
generator_epoch = self.load_state(generator_checkpoint,
self.generator,
generator_optimiser)
else:
LOG.info('Generator checkpoint not found')
if discriminator_epoch is None or generator_epoch is None:
epoch = 0
LOG.info(
"Discriminator or generator failed to load, training from start"
)
else:
epoch = min(generator_epoch, discriminator_epoch)
LOG.info("Generator loaded at epoch {0}".format(generator_epoch))
LOG.info("Discriminator loaded at epoch {0}".format(
discriminator_epoch))
LOG.info("Training from lowest epoch {0}".format(epoch))
vis_path = os.path.join(
os.path.splitext(self.config.FILENAME)[0], "gan",
str(datetime.now()))
with Visualiser(vis_path) as vis:
real_labels = None # all 1s
fake_labels = None # all 0s
epochs_complete = 0
while epoch < self.config.MAX_EPOCHS:
if self.check_requeue(epochs_complete):
return # Requeue needed and training not complete
for step, (data, noise1,
noise2) in enumerate(self.data_loader):
batch_size = data.size(0)
if real_labels is None or real_labels.size(
0) != batch_size:
real_labels = self.generate_labels(batch_size, [1.0])
if fake_labels is None or fake_labels.size(
0) != batch_size:
fake_labels = self.generate_labels(batch_size, [0.0])
if self.config.USE_CUDA:
data = data.cuda()
noise1 = noise1.cuda()
noise2 = noise2.cuda()
# ============= Train the discriminator =============
# Pass real noise through first - ideally the discriminator will return 1 #[1, 0]
d_output_real = self.discriminator(data)
# Pass generated noise through - ideally the discriminator will return 0 #[0, 1]
d_output_fake1 = self.discriminator(self.generator(noise1))
# Determine the loss of the discriminator by adding up the real and fake loss and backpropagate
d_loss_real = criterion(
d_output_real, real_labels
) # How good the discriminator is on real input
d_loss_fake = criterion(
d_output_fake1, fake_labels
) # How good the discriminator is on fake input
d_loss = d_loss_real + d_loss_fake
self.discriminator.zero_grad()
d_loss.backward()
discriminator_optimiser.step()
# =============== Train the generator ===============
# Pass in fake noise to the generator and get it to generate "real" noise
# Judge how good this noise is with the discriminator
d_output_fake2 = self.discriminator(self.generator(noise2))
# Determine the loss of the generator using the discriminator and backpropagate
g_loss = criterion(d_output_fake2, real_labels)
self.discriminator.zero_grad()
self.generator.zero_grad()
g_loss.backward()
generator_optimiser.step()
vis.step(d_loss_real.item(), d_loss_fake.item(),
g_loss.item())
# Report data and save checkpoint
fmt = "Epoch [{0}/{1}], Step[{2}/{3}], d_loss_real: {4:.4f}, d_loss_fake: {5:.4f}, g_loss: {6:.4f}"
LOG.info(
fmt.format(epoch + 1, self.config.MAX_EPOCHS, step + 1,
len(self.data_loader), d_loss_real,
d_loss_fake, g_loss))
epoch += 1
epochs_complete += 1
discriminator_checkpoint.set(
self.discriminator.state_dict(),
discriminator_optimiser.state_dict(), epoch).save()
generator_checkpoint.set(self.generator.state_dict(),
generator_optimiser.state_dict(),
epoch).save()
vis.plot_training(epoch)
data, noise1, _ = iter(self.data_loader).__next__()
if self.config.USE_CUDA:
data = data.cuda()
noise1 = noise1.cuda()
vis.test(epoch, self.data_loader.get_input_size_first(),
self.discriminator, self.generator, noise1, data)
generator_scheduler.step(epoch)
discriminator_scheduler.step(epoch)
LOG.info("Learning rates: d {0} g {1}".format(
discriminator_optimiser.param_groups[0]["lr"],
generator_optimiser.param_groups[0]["lr"]))
LOG.info("GAN Training complete")
def __call__(self):
"""
Main training loop for the GAN.
The training process is interruptable; the model and optimiser states are saved to disk each epoch, and the
latest states are restored when the trainer is resumed.
If the script is not able to load the generator's saved state, it will attempt to load the pre-trained generator
autoencoder from the generator_decoder_complete checkpoint (if it exists). If this also fails, the generator is
pre-trained as an autoencoder. This training is also interruptable, and will produce the
generator_decoder_complete checkpoint on completion.
On successfully restoring generator and discriminator state, the trainer will proceed from the earliest restored
epoch. For example, if the generator is restored from epoch 7 and the discriminator is restored from epoch 5,
training will proceed from epoch 5.
Visualisation plots are produces each epoch and stored in
/path_to_input_file_directory/{gan/generator_auto_encoder}/{timestamp}/{epoch}
Each time the trainer is run, it creates a new timestamp directory using the current time.
"""
# Load the autoencoder, and train it if needed.
if not self._train_autoencoder():
# Autoencoder training incomplete
return
# Setup environment and environment state builder
state_builder = RoomsStateBuilder(width=opt.width,
height=opt.height,
grayscale=opt.grayscale)
env = RoomsEnvironment(action_space=action_space,
mission_name=mission_name,
mission_xml=mission_xml,
remotes=clients,
state_builder=state_builder,
role=e,
recording_path=None)
vqae = None
if args.use_vqae:
# Initialize VQAE
vqae = Autoencoder(plot_class=plot_class)
# Initialize processor
processor = MalmoProcessor(autoencoder=vqae,
plot_class=plot_class,
action_space=action_space)
if args.agent_type == 'Random':
# Use Random agent to train the VQAE
agent = RandomAgent(num_actions=env.available_actions,
processor=processor)
elif args.agent_type == 'DDQN':
# Setup exploration policy
policy = LinearAnnealedPolicy(EpsGreedyQPolicy(),
attr='eps',
value_max=opt.eps_value_max,
if __name__ == '__main__':
hparams = {
'in_channels': 3,
'train_batch_size': config.train_batch_size,
'val_batch_size': config.val_batch_size,
'lr': config.lr,
'type': 'rgb',
'train_original': True,
'map_location': "cuda:0" if torch.cuda.is_available() else "cpu"
}
checkpoint = os.path.join(
config.ckpts_dir, 'Autoencoder/original_epoch=437-val_loss=1.40.ckpt')
pretrained_ckpt = torch.load(checkpoint, map_location=config.device)
pretrained_state_dict = pretrained_ckpt['state_dict']
pretrained_model = Autoencoder(hparams=hparams)
pretrained_model.load_state_dict(pretrained_state_dict)
list_of_files = glob.glob(
os.path.join(config.ckpts_dir,
'Autoencoder/Step/pruned-train-original/*.ckpt'))
for i, path in enumerate(list_of_files):
print(i)
print(path)
alpha = re.search('=(.+?)_', path)
if alpha:
alpha = alpha.group(1)
print("Alpha: " + alpha)
hparams['alpha'] = float(alpha)
test_model = PrunedModel(hparams=hparams,
class CycleModel(CustomModule):
ae_day: Autoencoder
ae_night: Autoencoder
reconstruction_loss_factor: float
cycle_loss_factor: float
def __init__(self, reconstruction_loss_factor: float, cycle_loss_factor: float):
# share weights of the upper encoder & lower decoder
encoder_upper, decoder_lower = UpperEncoder(), LowerDecoder()
self.ae_day = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower, UpperDecoder())
self.ae_night = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower, UpperDecoder())
self.loss_fn = nn.L1Loss()
self.reconstruction_loss_factor = reconstruction_loss_factor
self.cycle_loss_factor = cycle_loss_factor
self.optimizer = None
self.scheduler = None
def __call__(self, input):
raise NotImplementedError
def init_optimizers(self):
"""
Is called right before training and after model has been moved to GPU.
Supposed to initialize optimizers and schedulers.
"""
parameters = set()
parameters |= set(self.ae_day.parameters())
parameters |= set(self.ae_night.parameters())
self.optimizer = Adam(parameters)
# initialize scheduler
self.scheduler = ReduceLROnPlateau(self.optimizer, patience=15, verbose=True)
def train_epoch(self, train_loader, epoch, use_cuda, log_path, **kwargs):
loss_day2night2day_sum, loss_night2day2night_sum, loss_day2day_sum, loss_night2night_sum = 0, 0, 0, 0
for day_img, night_img in train_loader:
if use_cuda:
day_img, night_img = day_img.cuda(), night_img.cuda()
# Day -> Night -> Day
self.optimizer.zero_grad()
loss_day2night2day, loss_day2day = self.cycle_plus_reconstruction_loss(day_img, self.ae_day, self.ae_night)
loss = loss_day2night2day * self.cycle_loss_factor + loss_day2day * self.reconstruction_loss_factor
loss.backward()
self.optimizer.step()
# Night -> Day -> Night
self.optimizer.zero_grad()
loss_night2day2night, loss_night2night \
= self.cycle_plus_reconstruction_loss(night_img, self.ae_night, self.ae_day)
loss = loss_night2day2night * self.cycle_loss_factor + loss_night2night * self.reconstruction_loss_factor
loss.backward()
self.optimizer.step()
loss_day2night2day_sum += loss_day2night2day
loss_day2day_sum += loss_day2day
loss_night2day2night_sum += loss_night2day2night
loss_night2night_sum += loss_night2night
loss_day2night2day_mean = loss_day2night2day_sum / len(train_loader)
loss_day2day_mean = loss_day2day_sum / len(train_loader)
loss_night2day2night_mean = loss_night2day2night_sum / len(train_loader)
loss_night2night_mean = loss_night2night_sum / len(train_loader)
loss_mean = (loss_day2night2day_mean + loss_day2day_mean + loss_night2day2night_mean + loss_night2night_mean)/4
self.scheduler.step(loss_mean, epoch)
# log losses
log_str = f'[Epoch {epoch}] ' \
f'Train loss day -> night -> day: {loss_day2night2day_mean} ' \
f'Train loss night -> day -> night: {loss_night2day2night_mean} ' \
f'Train loss day -> day: {loss_day2day_mean} ' \
f'Train loss night -> night: {loss_night2night_mean}'
print(log_str)
with open(os.path.join(log_path, 'log.txt'), 'a+') as f:
f.write(log_str + '\n')
def validate(self, val_loader, epoch, use_cuda, log_path, **kwargs):
loss_day2night2day_sum, loss_night2day2night_sum, loss_day2day_sum, loss_night2night_sum = 0, 0, 0, 0
day_img, night_img = None, None
with torch.no_grad():
for day_img, night_img in val_loader:
if use_cuda:
day_img, night_img = day_img.cuda(), night_img.cuda()
# Day -> Night -> Day and Day -> Day
loss_day2night2day, loss_day2day = \
self.cycle_plus_reconstruction_loss(day_img, self.ae_day, self.ae_night)
# Night -> Day -> Night and Night -> Night
loss_night2day2night, loss_night2night = \
self.cycle_plus_reconstruction_loss(night_img, self.ae_night, self.ae_day)
loss_day2night2day_sum += loss_day2night2day
loss_day2day_sum += loss_day2day
loss_night2day2night_sum += loss_night2day2night
loss_night2night_sum += loss_night2night
loss_day2night2day_mean = loss_day2night2day_sum / len(val_loader)
loss_night2day2night_mean = loss_night2day2night_sum / len(val_loader)
loss_day2day_mean = loss_day2day_sum / len(val_loader)
loss_night2night_mean = loss_night2night_sum / len(val_loader)
# log losses
log_str = f'[Epoch {epoch}] ' \
f'Val loss day -> night -> day: {loss_day2night2day_mean} ' \
f'Val loss night -> day -> night: {loss_night2day2night_mean} ' \
f'Val loss day -> day: {loss_day2day_mean} ' \
f'Val loss night -> night: {loss_night2night_mean}'
print(log_str)
with open(os.path.join(log_path, 'log.txt'), 'a+') as f:
f.write(log_str + '\n')
# create sample images
latent_day = self.ae_day.encode(day_img[0].unsqueeze(0))
latent_night = self.ae_night.encode(night_img[0].unsqueeze(0))
# reconstruction
day2day = self.ae_day.decode(latent_day)
night2night = self.ae_night.decode(latent_night)
# domain translation
day2night = self.ae_night.decode(latent_day)
night2day = self.ae_day.decode(latent_night)
# cycle
day2night2day = self.ae_day.decode(self.ae_night.encode(day2night))
night2day2night = self.ae_night.decode(self.ae_day.encode(night2day))
# save sample images
samples = {
'day_img': day_img[0],
'night_img': night_img[0],
'day2day': day2day[0],
'night2night': night2night[0],
'day2night': day2night[0],
'night2day': night2day[0],
'day2night2day': day2night2day[0],
'night2day2night': night2day2night[0],
}
for name, img in samples.items():
ToPILImage()(img.cpu()).save(os.path.join(log_path, f'{epoch}_{name}.jpeg'), 'JPEG')
def cycle_plus_reconstruction_loss(self, image, autoencoder1, autoencoder2):
# send the image through the cycle
intermediate_latent_1 = autoencoder1.encode(image)
intermediate_opposite = autoencoder2.decode(intermediate_latent_1)
intermediate_latent_2 = autoencoder2.encode(intermediate_opposite)
cycle_img = autoencoder1.decode(intermediate_latent_2)
# do simple reconstruction
reconstructed_img = autoencoder1.decode(intermediate_latent_1)
cycle_loss = self.loss_fn(cycle_img, image)
reconstruction_loss = self.loss_fn(reconstructed_img, image)
return cycle_loss, reconstruction_loss
def train(self):
self.ae_day.train()
self.ae_night.train()
def eval(self):
self.ae_day.eval()
self.ae_night.eval()
def cuda(self):
self.ae_day.cuda()
self.ae_night.cuda()
def state_dict(self):
return {
'encoder_lower_day': self.ae_day.encoder_lower.state_dict(),
'encoder_lower_night': self.ae_night.encoder_lower.state_dict(),
'encoder_upper': self.ae_day.encoder_upper.state_dict(),
'decoder_day': self.ae_day.decoder.state_dict(),
'decoder_night': self.ae_night.decoder.state_dict()
}
def optim_state_dict(self):
return {
'optimizer': self.optimizer.state_dict(),
}
def load_state_dict(self, state):
self.ae_day.encoder_lower.load_state_dict(state['encoder_lower_day'])
self.ae_night.encoder_lower.load_state_dict(state['encoder_lower_night'])
self.ae_day.encoder_upper.load_state_dict(state['encoder_upper'])
self.ae_day.decoder.load_state_dict(state['decoder_day'])
self.ae_night.decoder.load_state_dict(state['decoder_night'])
}
tb_logger = pl_loggers.TensorBoardLogger(save_dir=os.path.join(
config.logs_dir, args.model),
name='original')
early_stopping = EarlyStopping(monitor='val_loss',
patience=config.patience,
mode='min')
checkpoint_callback = ModelCheckpoint(filepath=os.path.join(
config.ckpts_dir, args.model, 'original_{epoch:02d}-{val_loss:.2f}'),
monitor='val_loss',
mode='min')
# Create model object to train
if args.model is Model.VGG.value:
model = VGGNet(hparams=hparams)
elif args.model is Model.Autoencoder.value:
hparams['type'] = 'rgb'
model = Autoencoder(hparams=hparams)
model.apply(utils.weights_init)
model = model.to(torch.device(hparams['map_location']))
utils.checkParams(model)
trainer = pl.Trainer(logger=tb_logger,
gpus=1,
max_epochs=config.max_epochs,
callbacks=[checkpoint_callback])
trainer.fit(model)
z = np.random.normal(size=(dim_z, 1)) samples = subset3.sample(z, noise=True) """ N = 3 projection = MatrixProjection(dim_z, dim_input, N, noise_variance=1e-2) randomprojection = RandomMatrixProjection(dim_z, dim_input, N, noise_variance=1e-2) """ #### PyTorch #%% PyTorch net = Autoencoder(dim_input, 10, 3, 10, dim_input) lr = 0.0001 criterion = nn.MSELoss() optimizer = optim.Adam(net.parameters(), lr=lr) n_iter = int(1e3) n_epoch = 20 idx_in, idx_target = 0, 0 #launcher_predict(net, subset3, criterion, optimizer, n_epoch, n_iter, (idx_in, idx_target), plot=True) #%% Pytorch auto encoder multiple modalities encoder_layer_1_size = 3 encoder_layer_2_size = 3
def __init__(self, *args):
super().__init__(*args)
self.model = Autoencoder(encoder=self.encode,
shape=self.shape,
beta=1.)
if __name__ == '__main__':
torch.manual_seed(0)
data = read("./dataset-1")
features = np.array([
np.append(np.unpackbits(state), i % 2) for game in data
for i, state in enumerate(game[0])
])
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f'Using {device}')
model = Autoencoder().to(device)
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.98)
dataset_size = features.shape[0]
train_size = int(dataset_size * 0.95)
test_size = dataset_size - train_size
batch_size = 256
train_tensor = torch.Tensor(features[:train_size])
test_tensor = torch.Tensor(features[train_size:])
train_dataloader = DataLoader(TensorDataset(train_tensor, train_tensor),
batch_size=batch_size)
test_dataloader = DataLoader(TensorDataset(test_tensor, test_tensor),
class SimpleModel(CustomModule, EmbeddingGenerator):
ae_day: Autoencoder
ae_night: Autoencoder
def __init__(self):
encoder_upper, decoder_lower = UpperEncoder(), LowerDecoder()
self.ae_day = Autoencoder(LowerEncoder(), encoder_upper, decoder_lower,
UpperDecoder())
self.ae_night = Autoencoder(LowerEncoder(), encoder_upper,
decoder_lower, UpperDecoder())
self.loss_fn = nn.L1Loss()
self.optimizer_day = None
self.optimizer_night = None
self.scheduler_day = None
self.scheduler_night = None
def __call__(self, input):
raise NotImplementedError
def init_optimizers(self):
"""
Is called right before training and after model has been moved to GPU.
Supposed to initialize optimizers and schedulers.
"""
self.optimizer_day = Adam(self.ae_day.parameters(), lr=1e-4)
self.optimizer_night = Adam(self.ae_night.parameters(), lr=1e-4)
self.scheduler_day = ReduceLROnPlateau(self.optimizer_day,
patience=15,
verbose=True)
self.scheduler_night = ReduceLROnPlateau(self.optimizer_night,
patience=15,
verbose=True)
def train_epoch(self, train_loader, epoch, use_cuda, log_path, **kwargs):
loss_day_sum, loss_night_sum = 0, 0
for day_img, night_img in train_loader:
if use_cuda:
day_img, night_img = day_img.cuda(), night_img.cuda()
# zero day gradients
self.optimizer_day.zero_grad()
# train day autoencoder
out_day = self.ae_day(day_img)
loss_day = self.loss_fn(out_day, day_img)
# optimize
loss_day.backward()
self.optimizer_day.step()
# zero night gradients
self.optimizer_night.zero_grad()
# train night autoencoder
out_night = self.ae_night(night_img)
loss_night = self.loss_fn(out_night, night_img)
# optimize
loss_night.backward()
self.optimizer_night.step()
loss_day_sum += loss_day
loss_night_sum += loss_night
loss_day_mean = loss_day_sum / len(train_loader)
loss_night_mean = loss_night_sum / len(train_loader)
self.scheduler_day.step(loss_day_mean, epoch)
self.scheduler_night.step(loss_night_mean, epoch)
# log losses
log_str = f'[Epoch {epoch}] Train day loss: {loss_day_mean} Train night loss: {loss_night_mean}'
print(log_str)
with open(os.path.join(log_path, 'log.txt'), 'a+') as f:
f.write(log_str + '\n')
def validate(self, val_loader, epoch, use_cuda, log_path, **kwargs):
loss_day_sum, loss_night_sum = 0, 0
day_img, night_img, out_day, out_night = (None, ) * 4
with torch.no_grad():
for day_img, night_img in val_loader:
if use_cuda:
day_img, night_img = day_img.cuda(), night_img.cuda()
out_day = self.ae_day(day_img)
loss_day = self.loss_fn(out_day, day_img)
out_night = self.ae_night(night_img)
loss_night = self.loss_fn(out_night, night_img)
loss_day_sum += loss_day
loss_night_sum += loss_night
loss_day_mean = loss_day_sum / len(val_loader)
loss_night_mean = loss_night_sum / len(val_loader)
# domain translation
day_to_night = self.ae_night.decode(
self.ae_day.encode(day_img[0].unsqueeze(0)))
night_to_day = self.ae_day.decode(
self.ae_night.encode(night_img[0].unsqueeze(0)))
# log losses
log_str = f'[Epoch {epoch}] Val day loss: {loss_day_mean} Val night loss: {loss_night_mean}'
print(log_str)
with open(os.path.join(log_path, 'log.txt'), 'a+') as f:
f.write(log_str + '\n')
# save sample images
samples = {
'day_img': day_img[0],
'night_img': night_img[0],
'out_day': out_day[0],
'out_night': out_night[0],
'day_to_night': day_to_night[0],
'night_to_day': night_to_day[0]
}
for name, img in samples.items():
ToPILImage()(img.cpu()).save(
os.path.join(log_path, f'{epoch}_{name}.jpeg'), 'JPEG')
def register_hooks(self, layers):
"""
This function is not supposed to be called from outside the class.
"""
handles = []
embedding_dict = {}
def get_hook(name, embedding_dict):
def hook(model, input, output):
embedding_dict[name] = output.detach()
return hook
for layer in layers:
hook = get_hook(layer, embedding_dict)
handles.append(
getattr(self.ae_day.encoder_upper,
layer).register_forward_hook(hook))
return handles, embedding_dict
def deregister_hooks(self, handles):
"""
This function is not supposed to be called from outside the class.
"""
for handle in handles:
handle.remove()
def get_day_embeddings(self, img, layers):
"""
Returns deep embeddings for the passed layers inside the upper encoder.
"""
handles, embedding_dict = self.register_hooks(layers)
# forward pass
self.ae_day.encode(img)
self.deregister_hooks(handles)
return embedding_dict
def get_night_embeddings(self, img, layers):
"""
Returns deep embeddings for the passed layers inside the upper encoder.
"""
handles, embedding_dict = self.register_hooks(layers)
# forward pass
self.ae_night.encode(img)
self.deregister_hooks(handles)
return embedding_dict
def train(self):
self.ae_day.train()
self.ae_night.train()
def eval(self):
self.ae_day.eval()
self.ae_night.eval()
def cuda(self):
self.ae_day.cuda()
self.ae_night.cuda()
def state_dict(self):
return {
'encoder_lower_day': self.ae_day.encoder_lower.state_dict(),
'encoder_lower_night': self.ae_night.encoder_lower.state_dict(),
'encoder_upper': self.ae_day.encoder_upper.state_dict(),
'decoder_day': self.ae_day.decoder.state_dict(),
'decoder_night': self.ae_night.decoder.state_dict()
}
def optim_state_dict(self):
return {
'optimizer_day': self.optimizer_day.state_dict(),
'optimizer_night': self.optimizer_night.state_dict()
}
def load_state_dict(self, state):
self.ae_day.encoder_lower.load_state_dict(state['encoder_lower_day'])
self.ae_night.encoder_lower.load_state_dict(
state['encoder_lower_night'])
self.ae_day.encoder_upper.load_state_dict(state['encoder_upper'])
self.ae_day.decoder.load_state_dict(state['decoder_day'])
self.ae_night.decoder.load_state_dict(state['decoder_night'])
def load_optim_state_dict(self, state):
self.optimizer_day.load_state_dict(state['optimizer_day'])
self.optimizer_night.load_state_dict(state['optimizer_night'])
sampled_text = gen.sample_text(2)
print("Starting off, sampled text is ", sampled_text)
# TODO: clean up the hardcoded amount
if args.gen_model_type in ["gpt2", "ctrl"]:
decoder = GRUDecoder(gen.config.n_embd, gen_tokenizer.vocab_size,
args.decoder_hidden, args.decoder_layers,
args.decoder_dropout) # FOR GPT2
else:
decoder = GRUDecoder(gen.config.d_model, gen_tokenizer.vocab_size,
args.decoder_hidden, args.decoder_layers,
args.decoder_dropout)
autoencoder = Autoencoder(gen,
decoder,
args.device,
tokenizer=gen_tokenizer,
model_type=args.gen_model_type)
if args.record_run:
wandb.init(project="humorgan", config=args, dir="~/wandb")
wandb.watch((gen, dis))
# prepare main datasets
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
print("Batch size is {}".format(args.train_batch_size))
if args.gen_model_type in ["gpt2", "ctrl"]:
# don't need all types of input as the other models
train_dataset = load_and_cache_examples_generator(args,
gen_tokenizer,
def create_model(dataset):
final_model = Autoencoder(intermediate_dim=512,
original_dim=9408,
dataset=dataset)
return final_model
# threads = 4)
dataset = MNISTDataSet('../MNIST_data',
batch_size = 96)
test_dataset = MNISTDataSet('../MNIST_data',
batch_size = 96)
network = Autoencoder(
sess = sess,
n_classes = 2,
zed_dim = 8,
n_kernels = 16,
bayesian = False,
dataset = dataset,
input_channel = 1,
log_dir = log_dir,
variational = True,
save_dir = save_dir,
input_dims = [28,28],
load_snapshot = False,
learning_rate = 1e-3,
encoder_type = 'small'
test_dataset = test_dataset,
adversarial_training = True)
## Has to come after init_op ???
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
'train_original': args.train_original,
'lr': config.lr,
'map_location': map_location
}
# Load models
if args.model is Model.VGG.value:
smaller_checkpoint = torch.load(args.smaller_model_path, map_location=torch.device(hparams['map_location']))
smaller_model = VGGNet(hparams=hparams, model_params=config.compute_pruned_vgg_params(config.remove_ratio))
smaller_model.load_state_dict(smaller_checkpoint)
original_checkpoint = torch.load(args.original_model_path, map_location=torch.device(hparams['map_location']))
original_model = VGGNet(hparams=hparams)
original_model.load_state_dict(original_checkpoint['state_dict'])
elif args.model is Model.Autoencoder.value:
hparams['type'] = 'rgb'
hparams['prune_encoder'] = True
hparams['prune_decoder'] = True
smaller_checkpoint = torch.load(args.smaller_model_path, map_location=torch.device(hparams['map_location']))
smaller_model = Autoencoder(hparams=hparams, model_params=config.compute_pruned_autoencoder_params(config.remove_ratio, True, True))
smaller_model.load_state_dict(smaller_checkpoint)
original_checkpoint = torch.load(args.original_model_path, map_location=torch.device(hparams['map_location']))
original_model = Autoencoder(hparams=hparams)
original_model.load_state_dict(original_checkpoint['state_dict'])
csv_path = os.path.join(config.data_dir, args.model)
if not os.path.exists(csv_path):
os.mkdir(csv_path)
compare(original_model, smaller_model, csv_path, args.use_gpu)