def pre_train_EN(self, out_z, out_x):
self.net_mode(True)
out = False
pre_EN_cnt = 0
pbar = tqdm(total=self.niter_preEN)
pbar.update(pre_EN_cnt)
loss_fun = torch.nn.MSELoss()
loss_table = []
while not out:
for z, x in zip(out_z, out_x):
x = Variable(cuda(x.float(), self.use_cuda))
z = Variable(cuda(z.float(), self.use_cuda))
distributions = self.net._encode(x)
mu = distributions[:, :self.z_dim]
logvar = distributions[:, self.z_dim:]
z_hat = reparametrize(mu, logvar)
loss = loss_fun(z_hat, z)
loss_table.append(loss.data.item())
self.optim_EN.zero_grad()
loss.backward()
self.optim_EN.step()
if pre_EN_cnt >= self.niter_preEN:
out = True
break
pre_EN_cnt += 1
pbar.update(1)
pbar.write("[Pretrain Encoder Finished]")
pbar.close()
sys.stdout.flush()
return loss_table
def latent_traversal(BVAE, DAE_net, data, channel_mean, channel_std,
SAVE_PATH):
images = []
hsv_images = []
pertubs = np.linspace(-3.0, 3.0, 20)
batch = data.type(torch.cuda.FloatTensor).cuda()
reconstruction, mean, std = BVAE(batch)
newz = reparametrize(mean, std)
#newz[0][0]+=pertubs[0]
#try1=torch.rand(100,32)
#try1=try1.type(torch.cuda.FloatTensor).cuda()
#print(newz.shape)
for i in range(32):
for j in range(20):
y = torch.empty_like(newz).copy_(newz)
y[0][i] = pertubs[j]
reconstruct = BVAE.decoder(y)
recon = DAE_net(reconstruct)
new = recon.cpu().data.numpy()
denormalize(new[0], channel_mean, channel_std, sigmoid=False)
#changed=np.transpose(new[0],(1,2,0))
a, hsv = tensor_to_im(new[0])
images.append(a)
#hsv_images.append(hsv)
#images.append(a)
path = SAVE_PATH
#hsv_path="HSV_traversal"
#os.mkdir(hsv_path)
for i in range(32):
#hsv_index=os.path.join(hsv_path,str(i))
index = os.path.join(path, str(i))
os.mkdir(index)
#os.mkdir(hsv_index)
for j in range(20):
tosave = images[i * 20 + j]
#tosave_hsv=hsv_images[i*20+j]
imagename = str(j) + ".jpg"
save_path = os.path.join(index, imagename)
#hsv_save_path=os.path.join(hsv_index,imagename)
cv2.imwrite(save_path, tosave)
#cv2.imwrite(hsv_save_path,tosave_hsv)
make_gif(index, index)
def decoder_curves(args, data_loader, flag=False):
seed = args.seed
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
net = Solver(args)
net.net_mode(True)
out = False
pre_DE_cnt = 0
pbar = tqdm(total=net.niter_preDE)
pbar.update(pre_DE_cnt)
loss_table = []
while not out:
for x, y, mz in data_loader:
x = Variable(cuda(x.float(), net.use_cuda))
y = Variable(cuda(y.float(), net.use_cuda))
mz = Variable(cuda(mz.float(), net.use_cuda))
#z = mz[:,:,:args.z_dim]
z = reparametrize(mz[:, :, :args.z_dim], mz[:, :, args.z_dim:])
if flag:
y_zeros = torch.zeros_like(y)
z = torch.cat((y, y_zeros), dim=2)
x_recon = net.net._decode(z)
loss = reconstruction_loss(x, x_recon, 'bernoulli')
loss_table.append(loss.data.item())
net.optim_DE.zero_grad()
loss.backward()
net.optim_DE.step()
if pre_DE_cnt >= net.niter_preDE:
out = True
break
pre_DE_cnt += 1
pbar.update(1)
pbar.write("[Pretrain Encoder Finished]")
pbar.close()
sys.stdout.flush()
return loss_table
def gen_z(self, gen_size=10, fullz=False):
'''
Randomly sample x from dataloader, feed it to encoder, generate z
Return z and true latent value
@ out_z should be a list with length equals gen_size, each object has
size B*z_dim
@ out_y should be a list with length equals gen_size, each object has
size B*6
'''
self.net_mode(train=False)
out = False
gen_cnt = 0
#pbar = tqdm(total=gen_size)
#pbar.update(gen_cnt)
out_z = []
out_y = []
out_x = []
out_distr_list = []
while not out:
for x, y in self.data_loader:
out_y.append(y.squeeze(1)[:, 1:])
out_x.append(x)
#pbar.update(1)
gen_cnt += 1
x = Variable(cuda(x.float(), self.use_cuda))
out_distri = self.net.encoder(x).data
mu = out_distri[:, :self.z_dim]
logvar = out_distri[:, self.z_dim:]
out_z.append(reparametrize(mu, logvar))
out_distr_list.append(out_distri)
if gen_cnt >= gen_size:
out = True
break
self.net_mode(train=True)
if fullz == True:
return out_distr_list, out_y, out_x
else:
return out_z, out_y, out_x
def main(CLASS="None"):
if CLASS == "None": exit()
USE_CUDA = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
bvae_net = CapsuleBVAE(z_dim=opt.z_dim)
if opt.model != '':
bvae_net.load_state_dict(torch.load(opt.model))
else:
print('pls set the model path')
if USE_CUDA:
print("Let's use", torch.cuda.device_count(), "GPUs!")
bvae_net = torch.nn.DataParallel(bvae_net)
bvae_net.to(device)
# process for 'shapenet_part' or 'shapenet_core13'
bvae_net.eval()
test_loss_sum = 0
mu = torch.load("tmp_checkpoints_cbvae/airplane_latent_mu.pt")
logvar = torch.load("tmp_checkpoints_cbvae/airplane_latent_var.pt")
if USE_CUDA:
mu = mu.cuda()
logvar = logvar.cuda()
print(mu.size())
z = reparametrize(mu, logvar)
latent_caps = bvae_net.module._decode(z)
print(z)
for i in range(opt.batch_size):
torch.save(latent_caps[i, :],
"tmp_lcs/cbvae_latcaps_%s_%03d.pt" % (CLASS.lower(), i))