def __init__(self, path, mode):
self.path = path
self.mode = mode
(X_train, y_train), (X_test, y_test) = multi_mnist(path)
if self.mode == 'train':
self.X, self.y = X_train, y_train
else:
self.X, self.y = (X_test, y_test)
def load_data():
inpath = get_data_directory(__file__)
(X_np, Y), _ = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_np = X_np.astype(np.float32)
X_np /= 255.0
X = torch.from_numpy(X_np)
counts = torch.FloatTensor([len(objs) for objs in Y])
return X, counts
def fetch_data():
inpath = 'data/multi_mnist/'
(X_train, y_train), (X_test, y_test) = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_train, X_test = X_train.astype(np.float32), X_test.astype(np.float32)
X_train /= 255.0
X_test /= 255.0
mnist_train = torch.from_numpy(X_train)
mnist_test = torch.from_numpy(X_test)
return mnist_train, y_train, mnist_test, y_test
def load_data():
inpath = './data'
(X_np, Y), _ = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_np = X_np.astype(np.float32)
X_np /= 255.0
X = Variable(torch.from_numpy(X_np))
# Using FloatTensor to allow comparison with values sampled from
# Bernoulli.
counts = torch.FloatTensor([len(objs) for objs in Y])
return X, counts
def load_data():
inpath = './data'
(X_np, Y), _ = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_np = X_np.astype(np.float32)
X_np /= 255.0
X = torch.from_numpy(X_np)
# Using FloatTensor to allow comparison with values sampled from
# Bernoulli.
counts = torch.FloatTensor([len(objs) for objs in Y])
return X, counts
def __init__(self,
root,
training=True,
download=True,
max_digits=2,
canvas_size=50,
seed=42,
mini_data_size=None):
self.root = os.path.expanduser(root)
# check if multi mnist already compiled
self.multi_mnist_filename = 'multi_mnist_{}_{}_{}'.format(
max_digits, canvas_size, seed)
if not self._check_processed_exists():
if self._check_raw_exists():
# process into pt file
data = np.load(
os.path.join(self.root, 'raw',
self.multi_mnist_filename + '.npz'))
train_data, train_labels, test_data, test_labels = [
data[f] for f in data.files
]
self._process_and_save(train_data, train_labels, test_data,
test_labels)
else:
if not download:
raise RuntimeError(
'Dataset not found. Use download=True to download it.')
else:
(train_data,
train_labels), (test_data, test_labels) = multi_mnist(
root, max_digits, canvas_size, seed)
self._process_and_save(train_data, train_labels, test_data,
test_labels)
else:
data = torch.load(
os.path.join(self.root, 'processed',
self.multi_mnist_filename + '.pt'))
self.train_data, self.train_labels, self.test_data, self.test_labels = \
data['train_data'], data['train_labels'], data['test_data'], data['test_labels']
if training:
self.x, self.y = self.train_data, self.train_labels
else:
self.x, self.y = self.test_data, self.test_labels
if mini_data_size != None:
self.x = self.x[:mini_data_size]
self.y = self.y[:mini_data_size]
action='store_true',
default=False,
help='write hyper parameters and network architecture to stdout')
args = parser.parse_args()
if 'save' in args:
if os.path.exists(args.save):
raise RuntimeError('Output file "{}" already exists.'.format(
args.save))
if args.seed is not None:
pyro.set_rng_seed(args.seed)
# Load data.
inpath = './data'
(X_np, _), _ = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_np = X_np.astype(np.float32)
X_np /= 255.0
X = Variable(torch.from_numpy(X_np))
X_size = X.size(0)
if args.cuda:
X = X.cuda()
# Yields the following distribution over the number of steps (when
# taking a maximum of 3 steps):
# p(0) = 0.4
# p(1) = 0.3
# p(2) = 0.2
# p(3) = 0.1
def default_z_pres_prior_p(t):
help='fudge z_pres to remove discreteness for testing')
parser.add_argument('--seed', type=int, help='random seed', default=None)
parser.add_argument('-v', '--verbose', action='store_true', default=False,
help='write hyper parameters and network architecture to stdout')
args = parser.parse_args()
if 'save' in args:
if os.path.exists(args.save):
raise RuntimeError('Output file "{}" already exists.'.format(args.save))
if args.seed is not None:
pyro.set_rng_seed(args.seed)
# Load data.
inpath = './data'
(X_np, _), _ = multi_mnist(inpath, max_digits=2, canvas_size=50, seed=42)
X_np = X_np.astype(np.float32)
X_np /= 255.0
X = Variable(torch.from_numpy(X_np))
X_size = X.size(0)
if args.cuda:
X = X.cuda()
# Yields the following distribution over the number of steps (when
# taking a maximum of 3 steps):
# p(0) = 0.4
# p(1) = 0.3
# p(2) = 0.2
# p(3) = 0.1
def default_z_pres_prior_p(t):
def experiment(exp_specs):
ptu.set_gpu_mode(exp_specs['use_gpu'])
# Set up logging ----------------------------------------------------------
exp_id = exp_specs['exp_id']
exp_prefix = exp_specs['exp_name']
seed = exp_specs['seed']
set_seed(seed)
setup_logger(exp_prefix=exp_prefix, exp_id=exp_id, variant=exp_specs)
img_save_path = 'junk_vis/debug_more_proper'
# Prep the data -----------------------------------------------------------
data_path = 'junk_vis/multi_mnist_data'
canvas_size = 36
(X_train, _), (X_test, _) = multi_mnist(data_path,
max_digits=1,
canvas_size=canvas_size,
seed=42,
use_max=True)
X_train = X_train[:, None, ...]
X_test = X_test[:, None, ...]
X_train, X_test = torch.FloatTensor(X_train) / 255.0, torch.FloatTensor(
X_test) / 255.0
# np_imgs = np.load('/u/kamyar/dsprites-dataset/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz')['imgs']
# np_imgs = None
X_train = torch.clamp(X_train, 0.05, 0.95)
X_test = torch.clamp(X_test, 0.05, 0.95)
train_ds = TensorDataset(X_train)
val_ds = TensorDataset(X_test)
# Model Definition --------------------------------------------------------
if exp_specs['masked']:
model = MaskedVAE(
[1, canvas_size, canvas_size],
exp_specs['vae_specs']['z_dim'],
exp_specs['vae_specs']['encoder_specs'],
exp_specs['vae_specs']['decoder_specs'],
)
else:
model = VAE(
[1, canvas_size, canvas_size],
exp_specs['vae_specs']['z_dim'],
exp_specs['vae_specs']['encoder_specs'],
exp_specs['vae_specs']['decoder_specs'],
)
if ptu.gpu_enabled():
model.cuda()
# Optimizer ---------------------------------------------------------------
model_optim = Adam(model.parameters(),
lr=float(exp_specs['model_lr']),
weight_decay=float(exp_specs['model_wd']))
# -------------------------------------------------------------------------
global_iter = 0
for epoch in range(exp_specs['epochs']):
train_loader = DataLoader(train_ds,
batch_size=exp_specs['batch_size'],
shuffle=True,
num_workers=4,
pin_memory=True,
drop_last=True)
for iter_num, img_batch in enumerate(train_loader):
img_batch = img_batch[0]
if ptu.gpu_enabled(): img_batch = img_batch.cuda()
z_mean, z_log_cov, recon_mean, recon_log_cov, enc_mask, dec_mask = model(
img_batch)
elbo, KL = model.compute_ELBO(z_mean,
z_log_cov,
recon_mean,
recon_log_cov,
img_batch,
average_over_batch=True)
loss = -1. * elbo
loss.backward()
model_optim.step()
if global_iter % 1000 == 0:
mse = ((recon_mean - img_batch)**2).mean()
print('\nTraining Iter %d...' % global_iter)
print('ELBO:\t%.4f' % elbo)
print('MSE:\t%.4f' % mse)
print('KL:\t%.4f' % KL)
save_pytorch_tensor_as_img(
img_batch[0].data.cpu(),
os.path.join(img_save_path,
'%d_train_img.png' % (global_iter)))
save_pytorch_tensor_as_img(
recon_mean[0].data.cpu(),
os.path.join(img_save_path,
'%d_train_recon.png' % (global_iter)))
if exp_specs['masked']:
save_pytorch_tensor_as_img(
enc_mask[0].data.cpu(),
os.path.join(img_save_path,
'%d_train_enc_mask.png' % (global_iter)))
# save_pytorch_tensor_as_img(dec_mask[0].data.cpu(), os.path.join(img_save_path, '%d_train_dec_mask.png'%(global_iter)))
if global_iter % exp_specs['freq_val'] == 0:
with torch.no_grad():
print('Validating Iter %d...' % global_iter)
model.eval()
idxs = np.random.choice(int(X_test.size(0)),
size=exp_specs['batch_size'],
replace=False)
img_batch = X_test[idxs]
if ptu.gpu_enabled(): img_batch = img_batch.cuda()
z_mean, z_log_cov, recon_mean, recon_log_cov, enc_mask, dec_mask = model(
img_batch)
elbo, KL = model.compute_ELBO(z_mean,
z_log_cov,
recon_mean,
recon_log_cov,
img_batch,
average_over_batch=True)
mse = ((recon_mean - img_batch)**2).mean()
print('ELBO:\t%.4f' % elbo)
print('MSE:\t%.4f' % mse)
print('KL:\t%.4f' % KL)
for i in range(1):
save_pytorch_tensor_as_img(
img_batch[i].data.cpu(),
os.path.join(img_save_path,
'%d_%d_img.png' % (global_iter, i)))
save_pytorch_tensor_as_img(
recon_mean[i].data.cpu(),
os.path.join(img_save_path,
'%d_%d_recon.png' % (global_iter, i)))
if exp_specs['masked']:
save_pytorch_tensor_as_img(
enc_mask[i].data.cpu(),
os.path.join(
img_save_path,
'%d_%d_enc_mask.png' % (global_iter, i)))
# save_pytorch_tensor_as_img(dec_mask[i].data.cpu(), os.path.join(img_save_path, '%d_%d_dec_mask.png'%(global_iter, i)))
model.train()
global_iter += 1
def experiment(exp_specs):
ptu.set_gpu_mode(exp_specs['use_gpu'])
# Set up logging ----------------------------------------------------------
exp_id = exp_specs['exp_id']
exp_prefix = exp_specs['exp_name']
seed = exp_specs['seed']
set_seed(seed)
setup_logger(exp_prefix=exp_prefix, exp_id=exp_id, variant=exp_specs)
# Prep the data -----------------------------------------------------------
path = 'junk_vis/debug_att_vae_shallower_48_64_dim_0p1_kl_stronger_seg_conv'
(X_train, Y_train), (X_test, Y_test) = multi_mnist(path,
max_digits=2,
canvas_size=48,
seed=42,
use_max=False)
convert_dict = {0: [0., 0.], 1: [1., 0.], 2: [1., 1.]}
Num_train = np.array([convert_dict[a.shape[0]] for a in Y_train])
Num_test = np.array([convert_dict[a.shape[0]] for a in Y_test])
X_train = X_train[:, None, ...]
X_test = X_test[:, None, ...]
X_train, X_test = torch.FloatTensor(X_train) / 255.0, torch.FloatTensor(
X_test) / 255.0
mask_train, mask_test = torch.FloatTensor(Num_train), torch.FloatTensor(
Num_test)
train_ds = TensorDataset(X_train, Num_train)
val_ds = TensorDataset(X_test, Num_test)
# Model Definition --------------------------------------------------------
model = AttentiveVAE([1, 48, 48], exp_specs['vae_specs']['z_dim'],
exp_specs['vae_specs']['x_encoder_specs'],
exp_specs['vae_specs']['z_seg_conv_specs'],
exp_specs['vae_specs']['z_seg_fc_specs'],
exp_specs['vae_specs']['z_obj_conv_specs'],
exp_specs['vae_specs']['z_obj_fc_specs'],
exp_specs['vae_specs']['z_seg_recon_fc_specs'],
exp_specs['vae_specs']['z_seg_recon_upconv_specs'],
exp_specs['vae_specs']['z_obj_recon_fc_specs'],
exp_specs['vae_specs']['z_obj_recon_upconv_specs'],
exp_specs['vae_specs']['recon_upconv_part_specs'])
if ptu.gpu_enabled():
model.cuda()
# Optimizer ---------------------------------------------------------------
model_optim = Adam(model.parameters(),
lr=float(exp_specs['model_lr']),
weight_decay=float(exp_specs['model_wd']))
# -------------------------------------------------------------------------
global_iter = 0
for epoch in range(exp_specs['epochs']):
train_loader = DataLoader(train_ds,
batch_size=exp_specs['batch_size'],
shuffle=True,
num_workers=4,
pin_memory=False,
drop_last=True)
for iter_num, img_batch in enumerate(train_loader):
img_batch, num_batch = img_batch[0], img_batch[1]
if ptu.gpu_enabled(): img_batch = img_batch.cuda()
what_means, what_log_covs, where_means, where_log_covs, masks, recon_mean, recon_log_cov = model(
img_batch, num_batch)
elbo, KL = model.compute_ELBO(what_means + where_means,
what_log_covs + where_log_covs,
recon_mean,
recon_log_cov,
img_batch,
average_over_batch=True)
loss = -1. * elbo
loss = loss + 1. * sum([m.mean() for m in masks])
loss.backward()
model_optim.step()
if global_iter % exp_specs['freq_val'] == 0:
with torch.no_grad():
print('\nValidating Iter %d...' % global_iter)
model.eval()
idxs = np.random.choice(int(X_test.size(0)),
size=exp_specs['batch_size'],
replace=False)
img_batch, num_batch = X_test[idxs], Num_test[idxs]
if ptu.gpu_enabled(): img_batch = img_batch.cuda()
what_means, what_log_covs, where_means, where_log_covs, masks, recon_mean, recon_log_cov = model(
img_batch, num_batch)
elbo, KL = model.compute_ELBO(what_means + where_means,
what_log_covs +
where_log_covs,
recon_mean,
recon_log_cov,
img_batch,
average_over_batch=True)
mse = ((recon_mean - img_batch)**2).mean()
print('ELBO:\t%.4f' % elbo)
print('MSE:\t%.4f' % mse)
print('KL:\t%.4f' % KL)
for i in range(1):
save_pytorch_tensor_as_img(
img_batch[i].data.cpu(),
os.path.join(path,
'%d_%d_img.png' % (global_iter, i)))
save_pytorch_tensor_as_img(
recon_mean[i].data.cpu(),
os.path.join(path,
'%d_%d_recon.png' % (global_iter, i)))
save_pytorch_tensor_as_img(
masks[0][i].data.cpu(),
os.path.join(path, '%d_%d_mask_0.png' %
(global_iter, i)))
# save_pytorch_tensor_as_img(masks[1][i].data.cpu(), os.path.join(path, '%d_%d_mask_1.png'%(global_iter, i)))
model.train()
global_iter += 1