def __init__(self, args):
self.args = args
torch.manual_seed(self.args.seed)
if self.args.cuda:
torch.cuda.manual_seed(self.args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.ToTensor()),
batch_size=self.args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=False,
transform=transforms.ToTensor()),
batch_size=self.args.batch_size,
shuffle=True,
**kwargs)
self.train_loader = train_loader
self.test_loader = test_loader
self.model = VAE()
if self.args.cuda:
self.model.cuda()
self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)
def train_vae():
model = VAE(h,w,256,50,device)
optimizer = optim.Adadelta(model.parameters())
l2 = lambda epoch: pow((1.-1.*epoch/epochs),0.9)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=l2)
for epoch in range(epochs):
scheduler.step()
model.train()
tr_recon_loss = 0
for batch_idx, (data, target) in enumerate(tr):
if batch_idx >= tr_size:
break
data = data.to(device)
optimizer.zero_grad()
recon_batch, zmu, zvar,_= model(data)
recon_loss, kl = loss_V(recon_batch, data, zmu,torch.exp(0.5*zvar))
loss = recon_loss + kl
loss.backward()
tr_recon_loss += recon_loss.item()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tReconstruction-Loss: {:.4f}, KL: {:.4f}'.format(
epoch, batch_idx * len(data), len(mnist_tr),
100. * batch_idx / len(tr),
recon_loss / len(data),kl/len(data)))
print('====> Epoch: {} Reconstruction loss: {:.4f}'.format(
epoch, tr_recon_loss / (tr_size*mb_size)))
test(epoch,model)
return model
def slave(comm):
mus, logvars = load_init_z()
vae = VAE()
vae.load_state_dict(torch.load(cfg.vae_save_ckpt, map_location=lambda storage, loc: storage)['model'])
model = RNNModel()
model.load_state_dict(torch.load(cfg.rnn_save_ckpt, map_location=lambda storage, loc: storage)['model'])
count = 1
status = MPI.Status()
gpuid = comm.rank % 4
# device = torch.device('cuda:{}'.format(gpuid))
# vae.to(device)
# model.to(device)
print('Worker {} Started, model on GPU {}'.format(comm.rank, gpuid))
while True:
solution = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
tag = status.Get_tag()
if tag == 1:
print('Worker {} received solution {}'.format(comm.rank, count))
zs = [sample_init_z(mus, logvars) for _ in range(cfg.trials_per_pop)]
controller = deflatten_controller(solution)
reward = rollout(model, controller, zs)
print('Worker {} finished solution {}, reward: mean {} | max {} | min {} | std {}'.format(
comm.rank, count, reward.mean(), reward.max(), reward.min(), reward.std()))
comm.send(reward.mean(), dest=0, tag=2)
count += 1
elif tag == 3:
print('Worker {} evaluate current solution'.format(comm.rank))
controller = deflatten_controller(solution)
reward = evaluate(model, vae, controller)
comm.send(reward, dest=0, tag=2)
def test_real(epi):
vae = VAE()
vae.load_state_dict(torch.load(cfg.vae_save_ckpt)['model'])
model = RNNModel()
model.load_state_dict(torch.load(cfg.rnn_save_ckpt)['model'])
controller = Controller()
controller.load_state_dict(torch.load(cfg.ctrl_save_ckpt)['model'])
env = DoomTakeCover(True)
obs = env.reset()
model.reset()
frames = []
for step in range(cfg.max_steps):
frames.append(cv2.resize(obs, (256, 256)))
obs = torch.from_numpy(obs.transpose(2, 0,
1)).unsqueeze(0).float() / 255.0
mu, logvar, _, z = vae(obs)
inp = torch.cat((model.hx.detach(), model.cx.detach(), z), dim=1)
y = controller(inp)
y = y.item()
action = encode_action(y)
model.step(z.unsqueeze(0), action.unsqueeze(0))
obs_next, reward, done, _ = env.step(action.item())
obs = obs_next
if done:
break
print('Episode {}: Real Reward {}'.format(epi, step))
write_video(frames, 'real_{}.avi'.format(epi), (256, 256))
os.system('mv real_{}.avi /home/bzhou/Dropbox/share'.format(epi))
class Solver(utils.BaseSolver):
def build(self):
self.model = VAE(self.config)
if self.config.cuda:
self.model.cuda()
print(self.model)
# Build Optimizer (Training Only)
if self.config.mode == 'train':
self.optimizer = self.config.optimizer(
self.model.parameters(),
lr=self.config.learning_rate,
betas=(self.config.beta1, 0.999))
self.loss_function = layers.VAELoss(self.config.recon_loss)
def train_step(self, images):
# Reconstruct Images
recon_images, mu, log_variance = self.model(images)
# Calculate loss
recon_loss, kl_div = self.loss_function(images, recon_images, mu,
log_variance)
return recon_loss, kl_div
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--train', help='train variational autoencoder', action='store_true')
parser.add_argument('--generate',help='generate images', action='store_true')
parser.add_argument('--traverse',help='traverse latent space', action='store_true')
args = parser.parse_args()
sess = make_session()
if(args.train):
#prepare_tfrecords()
training_dataset = prepare_dataset(Config.tfrecord_dir+'train.tfrecord')
iterator = tf.compat.v1.data.make_one_shot_iterator(training_dataset)
vae = VAE(sess, training=True, training_iterator = iterator)
train(sess, vae)
if(args.generate) or (args.traverse):
vae = VAE(sess, training=False, restore=True)
if(args.generate):
random_vector = tf.random.normal(shape=[Config.num_gen_imgs, vae.latent_dim])
filename = Config.results_dir + 'randomFakeGrid.jpg'
gen_imgs = generate_fake_images(sess, vae, random_vector, filename)
if(args.traverse):
traverse_latents(sess, vae)
def test_vae(self):
dae = DAE()
vae = VAE(dae)
vars = vae.get_vars()
# Check size of optimizing vars
self.assertEqual(len(vars), 14 + 12)
def __init__(self,
hparams,
train_size: int,
class_weight: Optional[Tensor] = None):
# model, criterion
self.model = VAE()
# optimizer and scheduler
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=hparams.learning_rate,
eps=hparams.eps,
weight_decay=hparams.weight_decay)
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, **hparams.scheduler)
self.bce = nn.BCEWithLogitsLoss(reduction='none')
# self.kld = nn.KLDivLoss(reduction='sum')
# device
device_for_summary = self.__init_device(hparams.device,
hparams.out_device)
# summary
self.writer = SummaryWriter(logdir=hparams.logdir)
# TODO: fill in ~~DUMMY~~INPUT~~SIZE~~
path_summary = Path(self.writer.logdir, 'summary.txt')
if not path_summary.exists():
print_to_file(path_summary, summary, (self.model, (40, 11)),
dict(device=device_for_summary))
# save hyperparameters
path_hparam = Path(self.writer.logdir, 'hparams.txt')
if not path_hparam.exists():
print_to_file(path_hparam, hparams.print_params)
def run(config, seed, device):
set_global_seeds(seed)
logdir = Path(config['log.dir']) / str(config['ID']) / str(seed)
train_loader, test_loader = make_dataset(config)
if config['nn.type'] == 'VAE':
model = VAE(config, device)
elif config['nn.type'] == 'ConvVAE':
model = ConvVAE(config, device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
engine = Engine(config,
model=model,
optimizer=optimizer,
train_loader=train_loader,
test_loader=test_loader)
train_logs = []
eval_logs = []
for epoch in range(config['train.num_epoch']):
train_logger = engine.train(epoch, logdir=logdir)
train_logs.append(train_logger.logs)
eval_logger = engine.eval(epoch, logdir=logdir)
eval_logs.append(eval_logger.logs)
pickle_dump(obj=train_logs, f=logdir / 'train_logs', ext='.pkl')
pickle_dump(obj=eval_logs, f=logdir / 'eval_logs', ext='.pkl')
return None
def train(args):
data_path = args.data
data = pd.read_csv(data_path, index_col=0)
data = prepare_data(data)
batch_size = args.batch_size
lr = args.lr
n_epochs = args.n_epochs
verbose = args.verbose
model = VAE(input_size=data.shape[0], batch_size=batch_size)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
model.train()
for epoch in range(n_epochs):
loss_sum = 0
for i in range(data.shape[1] // batch_size):
batch = Variable(data[:, i * batch_size:(i + 1) * batch_size],
requires_grad=False)
optimizer.zero_grad()
recon_batch = model(batch)
loss = criterion(recon_batch, batch)
loss.backward()
optimizer.step()
loss_sum += loss.data.numpy()
if verbose:
print('Epoch: {} - loss: {:.2f}'.format(epoch, loss_sum))
def extract(fs, idx, N):
model = VAE()
model.load_state_dict(
torch.load(cfg.vae_save_ckpt,
map_location=lambda storage, loc: storage)['model'])
model = model.cuda(idx)
for n, f in enumerate(fs):
data = np.load(f)
imgs = data['sx'].transpose(0, 3, 1, 2)
actions = data['ax']
rewards = data['rx']
dones = data['dx']
x = torch.from_numpy(imgs).float().cuda(idx) / 255.0
mu, logvar, _, z = model(x)
save_path = "{}/{}".format(cfg.seq_extract_dir, f.split('/')[-1])
np.savez_compressed(save_path,
mu=mu.detach().cpu().numpy(),
logvar=logvar.detach().cpu().numpy(),
dones=dones,
rewards=rewards,
actions=actions)
if n % 10 == 0:
print('Process %d: %5d / %5d' % (idx, n, N))
def train():
use_cuda = True
epochs = 100
batch_size = 32
hidden_dim = 500
z_dim = 20
lr = 0.0001
compressed = transforms.Compose([dataset.ToTensor()])
variable = dataset.ToVariable(use_cuda=use_cuda)
kwargs = {'num_workers': 8, 'pin_memory': True}
loaders = dataset.setup_data_loaders(dataset.LastFMCSVDataset,
use_cuda,
batch_size,
transform=compressed,
**kwargs)
print('{} steps for all data / 1 epoch'.format(len(loaders['X'])))
vae = VAE(1, hidden_dim, z_dim, use_cuda=use_cuda)
adagrad_param = {'lr': lr}
optimizer = optim.Adam(vae.parameters(), **adagrad_param)
for epoch in range(epochs):
loss = train_epoch(vae, optimizer, loaders, variable)
print('Epoch{}:{}'.format(epoch, loss))
def build_ae_model(model_config, train_config):
"""
Builds an auto-encoder model given a configuration. Built model can be initialized later
with a previous state_dict.
:param model_config: model global attribute from the config.py module
:param train_config: train attributes (a few are required, e.g. dropout probability)
:return: Tuple: encoder, decoder, full AE model
"""
encoder_model, decoder_model = build_encoder_and_decoder_models(
model_config, train_config)
# AE model
if model_config.latent_flow_arch is None:
ae_model = VAE.BasicVAE(encoder_model, model_config.dim_z,
decoder_model, train_config.normalize_losses,
train_config.latent_loss)
else:
# TODO test latent flow dropout (in all but the last flow layers)
ae_model = VAE.FlowVAE(encoder_model,
model_config.dim_z,
decoder_model,
train_config.normalize_losses,
model_config.latent_flow_arch,
concat_midi_to_z0=model_config.concat_midi_to_z)
return encoder_model, decoder_model, ae_model
def __init__(self, dataset):
self.model = VAE()
if config.USE_GPU:
self.model.cuda()
self, dataset = dataset
self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)
self.train_loader = DataLoader(self.dataset(train=True))
self.test_loader = DataLoader(self.dataset(train=False))
def main(args):
### VAE on MNIST
n_transform = transforms.Compose([transforms.ToTensor()])
dataset = MNIST('data', transform=n_transform, download=True)
### CVAE on MNIST
# n_transform = transforms.Compose([transforms.ToTensor()])
# dataset = MNIST('data', transform=n_transform)
### CVAE on facescrub-5
# n_transform = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
# dataset = ImageFolder('facescrub-5', transform=n_transform)
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
vae = VAE(args.latent_size).cuda()
### CVAE
# vae = CVAE(args.latent_size, num_labels=args.num_labels).cuda()
optimizer = torch.optim.Adam(vae.parameters(), lr=args.learning_rate)
# Decide if you want to use fixed noise or not
fix_noise = Variable(torch.randn((50, args.latent_size)).cuda())
num_iter = 0
for epoch in range(args.epochs * 10):
for _, batch in enumerate(data_loader, 0):
img = Variable(batch[0].cuda())
label = Variable(batch[1].cuda())
recon_img, mean, log_var, z = vae(img)
### CVAE
# recon_img, mean, log_var, z = vae(img, label)
loss = loss_fn(recon_img, img, mean, log_var)
optimizer.zero_grad()
loss.backward()
optimizer.step()
num_iter += 1
if num_iter % args.print_every == 0:
print("Batch %04d/%i, Loss %9.4f" %
(num_iter, len(data_loader) - 1, loss.data.item()))
if num_iter % args.save_test_sample == 0:
x = vae.inference(fix_noise)
save_img(args, x.detach(), num_iter)
if num_iter % args.save_recon_img == 0:
save_img(args, recon_img.detach(), num_iter, recon=True)
def main(args):
print("Loading data")
dataset = args.data.rstrip('/').split('/')[-1]
torch.cuda.set_device(args.cuda)
device = args.device
if dataset == 'mnist':
train_loader, test_loader = get_mnist(args.batch_size, 'data/mnist')
num = 10
elif dataset == 'fashion':
train_loader, test_loader = get_fashion_mnist(args.batch_size,
'data/fashion')
num = 10
elif dataset == 'svhn':
train_loader, test_loader, _ = get_svhn(args.batch_size, 'data/svhn')
num = 10
elif dataset == 'stl':
train_loader, test_loader, _ = get_stl10(args.batch_size, 'data/stl10')
elif dataset == 'cifar':
train_loader, test_loader = get_cifar(args.batch_size, 'data/cifar')
num = 10
elif dataset == 'chair':
train_loader, test_loader = get_chair(args.batch_size,
'~/data/rendered_chairs')
num = 1393
elif dataset == 'yale':
train_loader, test_loader = get_yale(args.batch_size, 'data/yale')
num = 38
model = VAE(28 * 28, args.code_dim, args.batch_size, 10,
dataset).to(device)
phi = nn.Sequential(
nn.Linear(args.code_dim, args.phi_dim),
nn.LeakyReLU(0.2, True),
).to(device)
model.load_state_dict(torch.load(args.fname))
if args.tsne:
datas, targets = [], []
for i, (data, target) in enumerate(test_loader):
datas.append(data), targets.append(target)
if i >= 5:
break
data, target = torch.cat(datas, dim=0), torch.cat(targets, dim=0)
c = F.one_hot(target.long(), num_classes=num).float()
_, _, _, z = model(data.to(args.device), c.to(args.device))
z, target = z.detach().cpu().numpy(), target.cpu().numpy()
tsne = TSNE(n_components=2, init='pca', random_state=0)
z_2d = tsne.fit_transform(z)
plt.figure(figsize=(6, 5))
for a in range(8):
for b in range(a + 1, 10):
plot_embedding(
z_2d,
target,
a,
b,
)
plt.savefig('tsne_c{}_{}_{}{}.png'.format(
int(args.c), dataset, a, b))
def load_checkpoint(path):
checkpoint = torch.load(path)
model = VAE(**checkpoint["model"])
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = nn.Adam(model.parameters())
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
return model, optimizer
def train_VAE():
model = VAE()
model.summary()
history = []
x_train = import_image(folder='train')
x_test = import_image(folder='test')
#x_train = import_image('test')
arbitrary = np.zeros([x_train.shape[0], 1024 * 2])
model.summary()
history = model.fit(x_train, [x_train, arbitrary],
epochs=200,
batch_size=128,
shuffle=True,
verbose=1,
validation_data=(x_test,
[x_test,
arbitrary[:x_test.shape[0]]]))
print("history =", history.history.keys())
output_loss1 = history.history['recons_loss']
output_loss2 = history.history['KLD_loss']
np.save('VAE2_recons_loss', output_loss1)
np.save('VAE2_KLD_loss', output_loss2)
model.save_weights('saved_model/VAE2_epochs60_weights.h5')
def infer():
args = parse_args()
num_layers = args.num_layers
src_vocab_size = args.vocab_size
tar_vocab_size = args.vocab_size
batch_size = args.batch_size
init_scale = args.init_scale
max_grad_norm = args.max_grad_norm
hidden_size = args.hidden_size
attr_init = args.attr_init
latent_size = 32
if args.enable_ce:
fluid.default_main_program().random_seed = 102
framework.default_startup_program().random_seed = 102
model = VAE(hidden_size,
latent_size,
src_vocab_size,
tar_vocab_size,
batch_size,
num_layers=num_layers,
init_scale=init_scale,
attr_init=attr_init)
beam_size = args.beam_size
trans_res = model.build_graph(mode='sampling', beam_size=beam_size)
# clone from default main program and use it as the validation program
main_program = fluid.default_main_program()
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
exe = Executor(place)
exe.run(framework.default_startup_program())
dir_name = args.reload_model
print("dir name", dir_name)
dir_name = os.path.join(dir_name, "checkpoint")
fluid.load(main_program, dir_name, exe)
vocab, tar_id2vocab = get_vocab(args.dataset_prefix)
infer_output = np.ones((batch_size, 1), dtype='int64') * BOS_ID
fetch_outs = exe.run(feed={'tar': infer_output},
fetch_list=[trans_res.name],
use_program_cache=False)
with io.open(args.infer_output_file, 'w', encoding='utf-8') as out_file:
for line in fetch_outs[0]:
end_id = -1
if EOS_ID in line:
end_id = np.where(line == EOS_ID)[0][0]
new_line = [tar_id2vocab[e[0]] for e in line[1:end_id]]
out_file.write(space_tok.join(new_line))
out_file.write(line_tok)
def main(args):
# Check if the output folder is exist
if not os.path.exists('./vae_results/'):
os.mkdir('./vae_results/')
# Load data
torch.manual_seed(args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data', train=True, download=True,
transform=transforms.ToTensor()),
batch_size=args.batch_size, shuffle=True, **kwargs)
# Load model
model = VAE().cuda() if torch.cuda.is_available() else VAE()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Train and generate sample every epoch
loss_list = []
for epoch in range(1, args.epochs + 1):
model.train()
_loss = train(epoch, model, train_loader, optimizer)
loss_list.append(_loss)
model.eval()
sample = torch.randn(64, 20)
sample = Variable(sample).cuda() if torch.cuda.is_available() else Variable(sample)
sample = model.decode(sample).cpu()
save_image(sample.view(64, 1, 28, 28).data, 'vae_results/sample_' + str(epoch) + '.png')
plt.plot(range(len(loss_list)), loss_list, '-o')
plt.savefig('vae_results/vae_loss_curve.png')
def main():
parser = argparse.ArgumentParser(description='Train VAE.')
parser.add_argument('-c', '--config', help='Config file.')
args = parser.parse_args()
print(args)
c = json.load(open(args.config))
print(c)
lookback = 50 # 160
input_dim = 1
returns_lookback = 20
start_date = datetime.strptime(c['start_date'],
'%Y/%m/%d') if c['start_date'] else None
end_date = datetime.strptime(c['end_date'],
'%Y/%m/%d') if c['end_date'] else None
max_n_files = None
out_path = Path(c['out_dir'])
out_path.mkdir(exist_ok=True)
# setup the VAE
vae = VAE(c['vae_series_length'], z_dim=c['z_dim'], use_cuda=c['cuda'])
if c['checkpoint_load']:
checkpoint = torch.load(c['checkpoint_load'])
vae.load_state_dict(checkpoint['model_state_dict'])
ticker_files = glob(str(Path(c['in_dir']) / '*.csv'))
if 0:
ticker_files = ticker_files[:100]
print(f'Found {len(ticker_files)} ticker files.')
extract_vae_distances(start_date=start_date,
end_date=end_date,
n_lookback_days=c['n_lookback_days'],
n_backtest_days=c['n_backtest_days'],
n_trade_days=c['n_trade_days'],
n_pairs_vae=c['n_pairs_vae'],
n_pairs_backtest=c['n_pairs_backtest'],
vae=vae,
returns_lookback=returns_lookback,
ticker_files=ticker_files,
fundamentals_file=c['fundamentals_file'],
out_dir=c['out_dir'],
r_script_exe=c['r_script_exe'],
r_backtest_script=c['r_backtest_script'],
r_trade_script=c['r_trade_script'],
backtest_sd=c['backtest_sd'],
backtest_returns_file=c['backtest_returns_file'],
backtest_plot_file=c['backtest_plot_file'],
trade_returns_file=c['trade_returns_file'],
trade_plot_file=c['trade_plot_file'],
cuda=c['cuda'])
def test_generate(self):
model = VAE()
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
# generate with z_mu with batch size 5
batch_size = 5
z_mu = np.zeros((batch_size, 10), dtype=np.float32)
xs = model.generate(sess, z_mu)
self.assertEqual((batch_size, 4096), xs.shape)
def test_transform(self):
model = VAE()
batch_size = 10
batch_xs = self.get_batch_images(batch_size)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
z_mean, z_log_sigma_sq = model.transform(sess, batch_xs)
# check shape of latent variables
self.assertEqual((batch_size, 10), z_mean.shape)
self.assertEqual((batch_size, 10), z_log_sigma_sq.shape)
def main(ARGS, device):
"""
Prepares the datasets for training, and optional, validation and
testing. Then, initializes the VAE model and runs the training (/validation)
process for a given number of epochs.
"""
data_splits = ['train', 'val']
datasets = {
split: IMDB(ARGS.data_dir, split, ARGS.max_sequence_length,
ARGS.min_word_occ, ARGS.create_data)
for split in data_splits
}
pretrained_embeddings = datasets['train'].get_pretrained_embeddings(
ARGS.embed_dim).to(device)
model = VAE(
datasets['train'].vocab_size,
ARGS.batch_size,
device,
pretrained_embeddings=pretrained_embeddings,
trainset=datasets['train'],
max_sequence_length=ARGS.max_sequence_length,
lstm_dim=ARGS.lstm_dim,
z_dim=ARGS.z_dim,
embed_dim=ARGS.embed_dim,
n_lstm_layers=ARGS.n_lstm_layers,
kl_anneal_type=ARGS.kl_anneal_type,
kl_anneal_x0=ARGS.kl_anneal_x0,
kl_anneal_k=ARGS.kl_anneal_k,
kl_fbits_lambda=ARGS.kl_fbits_lambda,
word_keep_rate=ARGS.word_keep_rate,
)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
print('Starting training process...')
amount_of_files = len(os.listdir("trained_models"))
for epoch in range(ARGS.epochs):
elbos = run_epoch(model, datasets, device, optimizer)
train_elbo, val_elbo = elbos
print(
f"[Epoch {epoch} train elbo: {train_elbo}, val_elbo: {val_elbo}]")
# Perform inference on the trained model
with torch.no_grad():
model.eval()
samples = model.inference()
print(*idx2word(samples,
i2w=datasets['train'].i2w,
pad_idx=datasets['train'].pad_idx),
sep='\n')
model.save(f"trained_models/{amount_of_files + 1}.model")
def __init__(self):
_, _, self.vocab = get_iterators(opt)
self.vae = VAE(opt)
self.vae.embedding.weight.data.copy_(self.vocab.vectors)
self.vae = get_cuda(self.vae)
checkpoint = T.load('data/saved_models/vae_model.121.pyt')
self.vae.load_state_dict(checkpoint['vae_dict'])
self.vae.eval()
del checkpoint
def test_reconstruct(self):
batch_size = 10
batch_xs = self.get_batch_images(batch_size)
model = VAE()
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
xs = model.reconstruct(sess, batch_xs)
# check reconstructed image shape
self.assertEqual((batch_size, 4096), xs.shape)
def test_partial_fit(self):
model = VAE()
batch_size = 10
batch_xs = self.get_batch_images(batch_size)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
# loss result should be float
reconstr_loss, latent_loss, _ = model.partial_fit(sess, batch_xs, 0)
self.assertEqual(np.float32, reconstr_loss.dtype)
self.assertEqual(np.float32, latent_loss.dtype)
def load_checkpoint(file_path, use_cuda=False):
"""Return EmbedNet instance"""
if use_cuda:
checkpoint = torch.load(file_path)
else:
checkpoint = torch.load(file_path,
map_location=lambda storage, location: storage)
n_latents = checkpoint['n_latents']
vae = VAE(n_latents=n_latents)
vae.load_state_dict(checkpoint['state_dict'])
return vae
def main(argv):
data_manager = DataManager()
data_manager.prepare()
dae = DAE()
vae = VAE(dae, beta=flags.vae_beta)
scan = SCAN(dae, vae, beta=flags.scan_beta, lambd=flags.scan_lambda)
scan_recomb = SCANRecombinator(dae, vae, scan)
dae_saver = CheckPointSaver(flags.checkpoint_dir, "dae", dae.get_vars())
vae_saver = CheckPointSaver(flags.checkpoint_dir, "vae", vae.get_vars())
scan_saver = CheckPointSaver(flags.checkpoint_dir, "scan", scan.get_vars())
scan_recomb_saver = CheckPointSaver(flags.checkpoint_dir, "scan_recomb",
scan_recomb.get_vars())
sess = tf.Session()
# Initialze variables
init = tf.global_variables_initializer()
sess.run(init)
# For Tensorboard log
summary_writer = tf.summary.FileWriter(flags.log_file, sess.graph)
# Load from checkpoint
dae_saver.load(sess)
vae_saver.load(sess)
scan_saver.load(sess)
scan_recomb_saver.load(sess)
# Train
if flags.train_dae:
train_dae(sess, dae, data_manager, dae_saver, summary_writer)
if flags.train_vae:
train_vae(sess, vae, data_manager, vae_saver, summary_writer)
disentangle_check(sess, vae, data_manager)
if flags.train_scan:
train_scan(sess, scan, data_manager, scan_saver, summary_writer)
sym2img_check(sess, scan, data_manager)
img2sym_check(sess, scan, data_manager)
if flags.train_scan_recomb:
train_scan_recomb(sess, scan_recomb, data_manager, scan_recomb_saver,
summary_writer)
recombination_check(sess, scan_recomb, data_manager)
sess.close()
def main(_):
if not os.path.exists(FLAGS.checkpoint_dir):
os.makedirs(FLAGS.checkpoint_dir)
if not os.path.exists(FLAGS.tensorboard_dir):
os.makedirs(FLAGS.tensorboard_dir)
#feed = webcamFeeder(color=True)
feed = videoFeeder("video.mp4", color=True)
with tf.Session() as sess:
vae = VAE(sess, feed, FLAGS)
vae.train()
def train(args):
learning_rate = args.learning_rate
batch_size = args.batch_size
training_epochs = args.training_epochs
keep_prob = args.keep_prob
checkpoint_step = args.checkpoint_step # save training results every check point step
z_dim = args.z_dim # number of latent variables.
loss_mode = args.loss_mode
diff_mode = False
if args.diff_mode == 1:
diff_mode = True
dirname = 'save'
if not os.path.exists(dirname):
os.makedirs(dirname)
with open(os.path.join(dirname, 'config.pkl'), 'w') as f:
cPickle.dump(args, f)
vae = VAE(learning_rate=learning_rate, batch_size=batch_size, z_dim = z_dim, keep_prob = keep_prob, loss_mode = loss_mode)
mnist = read_data_sets()
n_samples = mnist.num_examples
# load previously trained model if appilcable
ckpt = tf.train.get_checkpoint_state(dirname)
if ckpt:
vae.load_model(dirname)
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
avg_likelihood_loss = 0.
avg_kl_loss = 0.
mnist.shuffle_data()
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_xs = mnist.next_batch(batch_size)
if (diff_mode == True):
batch_xs = mnist.integrate_batch(batch_xs)
# Fit training using batch data
cost, likelihood_loss, kl_loss = vae.partial_fit(batch_xs)
# Compute average loss
avg_cost += cost / n_samples * batch_size
avg_likelihood_loss += likelihood_loss / n_samples * batch_size
avg_likelihood_loss += kl_loss / n_samples * batch_size
# Display logs per batch
'''
print "batch:", '%04d' % (i+1), \
"total loss =", "{:.6f}".format(cost), \
"likelihood_loss =", "{:.6f}".format(likelihood_loss), \
"kl_loss =", "{:.6f}".format(kl_loss)
'''
# Display logs per epoch step
print "Epoch:", '%04d' % (epoch+1), \
"total loss =", "{:.6f}".format(avg_cost), \
"likelihood_loss =", "{:.6f}".format(avg_likelihood_loss), \
"kl_loss =", "{:.6f}".format(avg_kl_loss)
# save model
if epoch > 0 and epoch % checkpoint_step == 0:
checkpoint_path = os.path.join('save', 'model.ckpt')
vae.save_model(checkpoint_path, epoch)
print "model saved to {}".format(checkpoint_path)
# save model one last time, under zero label to denote finish.
vae.save_model(checkpoint_path, 0)
return vae
class Sampler():
def __init__(self):
self.mnist = None
self.model = VAE()
self.model.load_model('save')
self.z = self.generate_z()
def get_random_mnist(self, with_label = False):
if self.mnist == None:
self.mnist = read_data_sets()
if with_label == True:
data, label = self.mnist.next_batch(1, with_label)
return data[0], label[0]
return self.mnist.next_batch(1)[0]
def get_random_specific_mnist(self, label = 2):
m, l = self.get_random_mnist(with_label = True)
for i in range(100):
if l == label:
break
m, l = self.get_random_mnist(with_label = True)
return m
def generate_random_label(self, label):
m = self.get_random_specific_mnist(label)
self.show_image(m)
self.show_image_from_z(self.encode(m))
def generate_z(self):
z = np.random.normal(size=(1, self.model.z_dim)).astype(np.float32)
return z
def encode(self, mnist_data):
new_shape = [1]+list(mnist_data.shape)
return self.model.transform(np.reshape(mnist_data, new_shape))
def generate(self, z=None):
if z is None:
z = self.generate_z()
else:
z = np.reshape(z, (1, self.model.z_dim))
self.z = z
return self.model.generate(z)[0]
def show_image(self, image_data):
'''
image_data is a tensor, in [height width depth]
image_data is NOT the PIL.Image class
'''
plt.subplot(1, 1, 1)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
plt.imshow(image_data, interpolation='nearest')
else:
plt.imshow(image_data.reshape(y_dim, x_dim), cmap='Greys', interpolation='nearest')
plt.axis('off')
plt.show()
def show_image_from_z(self, z):
self.show_image(self.generate(z))
def to_image(self, image_data):
# convert to PIL.Image format from np array (0, 1)
img_data = np.array(1-image_data)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
img_data = np.array(img_data.reshape((y_dim, x_dim, c_dim))*255.0, dtype=np.uint8)
else:
img_data = np.array(img_data.reshape((y_dim, x_dim))*255.0, dtype=np.uint8)
im = Image.fromarray(img_data)
return im
def diff_image(self, image_data):
# perform 2d differentiation on mnist image
m2 = np.array(image_data) # makes a copy
m2[1:,1:,:] = m2[1:,1:,:]-m2[0:-1,1:,:]
m2[1:,1:,:] = m2[1:,1:,:]-m2[1:,0:-1,:]
return m2
def integrate_image(self, image_data):
# integrates differentiated batch back to mnist image
m3 = np.array(image_data)
m3 = m3.cumsum(axis=0)
m3 = m3.cumsum(axis=1)
return m3
def __init__(self):
self.mnist = None
self.model = VAE()
self.model.load_model('save')
self.z = self.generate_z()
