def main(args):
if args.seed >= 0:
random_seed(args.seed)
# create output directory for saving result images
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# prepare and read dataset
train_set, _, test_set = mnist(args.data_dir)
train_x, train_y = train_set
test_x, test_y = test_set
# specify the encoder and decoder net structure
encoder_net = Net([
Dense(256),
ReLU(),
Dense(64)
])
decoder_net = Net([
ReLU(),
Dense(256),
Tanh(),
Dense(784),
Tanh()
])
nets = (encoder_net, decoder_net)
optimizers = (Adam(args.lr), Adam(args.lr))
model = AutoEncoder(nets, loss=MSE(), optimizer=optimizers)
# for pre-trained model, test generated images from latent space
if args.load_model is not None:
# load pre-trained model
model.load(os.path.join(args.output_dir, args.load_model))
print("Loaded model fom %s" % args.load_model)
# transition from test[from_idx] to test[to_idx] in n steps
idx_arr, n = [2, 4, 32, 12, 82], 160
print("Transition in numbers", [test_y[i] for i in idx_arr],
"in %d steps ..." % n)
stops = [model.en_net.forward(test_x[i]) for i in idx_arr]
k = int(n / (len(idx_arr) - 1)) # number of code per transition
# generate all transition codes
code_arr = []
for i in range(len(stops) - 1):
t = [c.copy() for c in transition(stops[i], stops[i+1], k)]
code_arr += t
# apply decoding all n "code" from latent space...
batch = None
for code in code_arr:
# translate latent space to image
genn = model.de_net.forward(code)
# save decoded results in a batch
if batch is None:
batch = np.array(genn)
else:
batch = np.concatenate((batch, genn))
output_path = os.path.join(args.output_dir, "genn-latent.png")
save_batch_as_images(output_path, batch)
quit()
# train the auto-encoder
iterator = BatchIterator(batch_size=args.batch_size)
for epoch in range(args.num_ep):
for batch in iterator(train_x, train_y):
origin_in = batch.inputs
# make noisy inputs
m = origin_in.shape[0] # batch size
mu = args.gaussian_mean # mean
sigma = args.gaussian_std # standard deviation
noises = np.random.normal(mu, sigma, (m, 784))
noises_in = origin_in + noises # noisy inputs
# forward
genn = model.forward(noises_in)
# back-propagate
loss, grads = model.backward(genn, origin_in)
# apply gradients
model.apply_grads(grads)
print("Epoch: %d Loss: %.3f" % (epoch, loss))
# save all the generated images and original inputs for this batch
noises_in_path = os.path.join(
args.output_dir, "ep%d-input.png" % epoch)
genn_path = os.path.join(
args.output_dir, "ep%d-genn.png" % epoch)
save_batch_as_images(noises_in_path, noises_in, titles=batch.targets)
save_batch_as_images(genn_path, genn, titles=batch.targets)
# save the model after training
model.save(os.path.join(args.output_dir, args.save_model))
