def main(params):
feature_size = params.feature_size
epoch_num = params.epoch_num
lr = params.lr
weight_decay = params.weight_decay
hidden_list = params.hidden_list
save_name = params.name
device = torch.device('cuda:1')
features = load_cell_gene_features(params)
features = torch.tensor(features, dtype=torch.float32).to(device)
vae = VAE(embedding_size=features.shape[1],
hidden_size_list=hidden_list,
mid_hidden=feature_size).to(device)
optimizer = optim.Adam(vae.parameters(), lr=lr, weight_decay=weight_decay)
criterion = nn.MSELoss()
for i in range(epoch_num):
x_hat, kl_div = vae(features)
loss = criterion(x_hat, features)
if kl_div is not None:
loss += kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f"epoch {i}: {loss:.4f} loss.")
proj_path = Path(__file__).parent.resolve().parent.resolve()
model_path = proj_path / 'saved_model'
if not model_path.exists():
model_path.mkdir()
torch.save(vae, model_path / save_name)
def main(model_name='AE',
embedding_size=128,
n_epochs=1000,
batch_size=32,
roi_size=128):
if model_name == 'VAE':
model = VAE(embedding_size=embedding_size)
criterion = VAELoss()
elif 'AE':
model = AE(embedding_size=embedding_size)
criterion = nn.MSELoss()
details = 'L{}'.format(embedding_size)
log_dir = os.path.join(
log_dir_root, '{}_{}_{}'.format(model_name, details,
time.strftime('%Y%m%d_%H%M%S')))
#criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
if is_cuda:
print('This is CUDA!!!!')
torch.backends.cudnn.benchmark = True #useful for arrays of fix dimension
model = model.cuda()
criterion = criterion.cuda()
generator = ROIFlowBatch(mask_file,
feat_file,
is_cuda=is_cuda,
batch_size=batch_size,
roi_size=roi_size,
is_shuffle=True)
t = TrainerAutoEncoder(model, optimizer, criterion, generator, n_epochs,
log_dir)
t.fit()
class FixedRandomModel(nn.Module):
def __init__(self, alpha):
super().__init__()
#self.conv = Conv()
vae_file = join("./vae", 'best.tar')
assert exists(vae_file), "vae is untrained."
vae_state = torch.load(vae_file, map_location={'cuda:0': str(device)})
logger.info("Loading VAE at epoch {} with test loss {}".format(
vae_state['epoch'], vae_state['precision']))
self.vae = VAE(3, LSIZE).to(device).double()
self.vae.load_state_dict(vae_state['state_dict'])
self.W_in = nn.Linear(2 * 2 * 256, 512, bias=False)
self.W = nn.Linear(512, 512, bias=False)
self.W.weight.data = init_W(512, 512)
self.x_esn = None
self.alpha = alpha
def forward(self, obs):
B = obs.shape[0]
_, _, _, x_conv, _ = self.vae(obs)
x_conv_flat = x_conv.view(B, -1)
if self.x_esn is None or self.x_esn.shape[0] != B:
x_esn = torch.tanh(self.W_in(x_conv_flat))
else:
x_hat = torch.tanh(self.W_in(x_conv_flat) + self.W(self.x_esn))
x_esn = (1 - self.alpha) * self.x_esn + self.alpha * x_hat
self.x_esn = x_esn
return (x_conv_flat, x_esn)
def _latent_attack(self, src, tar):
tar_recon, _, tar_z_params = self.vae(tar)
tar_z_params = VAE.flatten_dists_params(tar_z_params).detach()
noise = Normal(loc=src, scale=self.args.noise_std)
adv = torch.clamp(src + noise.sample(), 0, 1).requires_grad_()
optimizer = optim.Adam([adv], lr=self.args.eta)
# per-pixel lower and upper bounds
_min = torch.max(torch.zeros_like(src.data), src.data - self.args.eps)
_max = torch.min(torch.ones_like(src.data), src.data + self.args.eps)
for i in range(self.args.steps):
adv_recon, _, adv_z_params = self.vae(adv)
adv_z_params = VAE.flatten_dists_params(adv_z_params)
l2pix = l2_norm(adv - src)
l2latent = l2_norm(adv_z_params - tar_z_params)
loss = self.args._lambda * l2pix + l2latent
loss.backward()
optimizer.step()
adv.data = torch.min(torch.max(adv.data, _min), _max)
return adv
def __init__(self, config):
self.config = config
self.lr = config.lr
self.batchsize = config.batch
self.dataRoot = config.data_root
self.z_dim = config.nz
self.lr = config.lr
if torch.cuda.is_available():
self.use_cuda = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.set_device(config.deviceId)
else:
self.use_cuda = False
torch.set_default_tensor_type('torch.FloatTensor')
#Transforms for the data
transformList = []
transformList.append(transforms.ToTensor())
transformSequence = transforms.Compose(transformList)
self.dataLoaderTrain_L, self.dataLoaderTrain_U, self.dataLoaderVal, self.dataLoaderTest = \
get_dataLoaderVAE(self.dataRoot, transformSequence, batch_size=self.batchsize)
self.model = VAE(zdim=self.z_dim)
if self.use_cuda: self.model = self.model.cuda()
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
def __init__(self, mdir, device, controller_model):
""" Run one step.
Load VAE and MDRNN from files
Take the controller (exp/ctrl) an an input, so we can easily change stuff inside the other file.
"""
self.controller = controller_model.to(device)
# Load controllers
vae_file, rnn_file, ctrl_file = \
[join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)
]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
print("Loading {} at epoch {} "
"with test loss {}".format(m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device)
self.vae.load_state_dict(vae_state['state_dict'])
# MDRNNCell
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device)
self.mdrnn.load_state_dict(
{k.strip('_l0'): v
for k, v in rnn_state['state_dict'].items()})
class Agent():
"""
Agent for training
"""
def __init__(self):
# Loading world model and vae
vae_file, rnn_file, ctrl_file = \
[join("./training", m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)
]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
logger.info("Loading {} at epoch {} "
"with test loss {}".format(m, s['epoch'],
s['precision']))
self.vae = VAE(3, LSIZE).to(device).double()
self.vae.load_state_dict(vae_state['state_dict'])
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device).double()
self.mdrnn.load_state_dict(
{k.strip('_l0'): v
for k, v in rnn_state['state_dict'].items()})
for p in self.vae.parameters():
p.requires_grad = False
for p in self.mdrnn.parameters():
p.requires_grad = False
self.net = Controller(LSIZE, RSIZE, ASIZE).to(device).double()
# load controller if it was previously saved
if exists(ctrl_file):
ctrl_state = torch.load(ctrl_file,
map_location={'cuda:0': str(device)})
logger.info("Loading Controller with reward {}".format(
ctrl_state['reward']))
self.net.load_state_dict(ctrl_state['state_dict'])
def select_action(self, state, hidden):
with torch.no_grad():
_, latent_mu, _ = self.vae(state)
alpha, beta = self.net(latent_mu, hidden[0])[0]
action = alpha / (alpha + beta)
_, _, _, _, _, next_hidden = self.mdrnn(action, latent_mu, hidden)
action = action.squeeze().cpu().numpy()
return action, next_hidden
def load_param(self):
self.net.load_state_dict(torch.load('param/ppo_net_params.pkl'))
def compute_cav(model: VAE, pos, neg):
"""
Computes the CAV for a single concept
"""
pos_enc = [model.reparameterize(*model.encode(x)) for x in pos]
neg_enc = [model.reparameterize(*model.encode(x)) for x in neg]
clf = train_binary_linear_classifier(pos_enc, neg_enc)
cav = tf.math.l2_normalize(clf.weights[0])
return cav, clf
def train_VAE():
transform = transforms.Compose(
[transforms.CenterCrop(127),
transforms.ToTensor()])
print("Loading ShapeNet...")
shapenet = ShapeNetDataset(root_dir='/home/svcl-oowl/dataset/shapenet',
transform=transform)
print("Initializing dataloader: {} workers allocated".format(
args.num_workers))
dataloader = DataLoader(shapenet,
batch_size=args.batch_size,
num_workers=args.num_workers,
worker_init_fn=shapenet.worker_init_fn,
pin_memory=True)
print("Initializing model...")
vae = VAE(image_channels=3).to(device)
print(vae)
# for name, param in vae.encoder.named_parameters():
# Freeze last couple conv layers and fc layer of resnet18 encoder
# if 'layer4' in name or 'fc' in name:
# param.requires_grad = True
# else:
# param.requires_grad = False
if os.path.isfile('vae.torch'):
print("Previous checkpoint found. Loading from memory...")
# vae.load_state_dict(torch.load('vae.torch', map_location=device))
optimizer = torch.optim.Adam(vae.parameters(), lr=args.lr)
print("Training...")
for epoch in range(args.epochs):
epoch_loss = 0
for i_batch, images in enumerate(dataloader):
images = images.to(device)
recon_images, mu, logvar = vae(images)
loss, bce, kld = loss_fn(recon_images, images, mu, logvar)
optimizer.zero_grad()
loss.backward()
optimizer.step()
bs = args.batch_size
to_print = "Epoch[{}/{}] Loss: {:.3f} {:.3f} {:.3f}".format(
epoch + 1, args.epochs,
loss.data.item() / bs,
bce.data.item() / bs,
kld.data.item() / bs)
epoch_loss += loss.data.item() / bs
# print(to_print)
# print("Images processed[{}/{}]".format((i_batch+1)*bs, '~26000'))
epoch_loss = (1 / (i_batch + 1)) * epoch_loss
print("Averaged loss over entire epoch[{}/{}]: {:.3f}".format(
epoch + 1, args.epochs, epoch_loss))
print("Saving model...")
torch.save(vae.state_dict(), 'vae-resnet.torch')
def plot_reconstructed_images(model: VAE, test_sample: tf.Tensor, filename=None, figsize=(4, 4)):
"""Forwards samples through the model and displays them."""
mean, log_variance = model.encode(test_sample)
z = model.reparameterize(mean, log_variance)
predictions = model.sample(z)
plt.figure(figsize=figsize)
for i in range(predictions.shape[0]):
plt.subplot(*figsize, i+1)
plt.imshow(predictions[i], vmin=0., vmax=1.)
plt.axis('off')
if filename:
plt.savefig(filename, dpi=300)
plt.show()
def __init__(self, M, eh1, eh2, dh2, ci, lr_ac=0.001, lr_cr=0.001):
## Network initializations
# Actor
self.actor = VAE(
M, eh1, eh2, dh2
) # Number of inputs, units in encoder_hidden_layer1, encoder_hidden_layer2,
#decoder_hidden_layer1
# Critic
self.critic = Critic(ci) # Length of feature vector
# Optimizers
self.optim_actor = torch.optim.Adam(self.actor.parameters(), lr=lr_ac)
self.optim_critic = torch.optim.Adam(self.critic.parameters(),
lr=lr_cr)
self.mse = torch.nn.MSELoss()
def compute_tcav(model: VAE, pos_concept_a, neg_concept_a, pos_concept_b,
neg_concept_b):
"""
Computes the modified VAE-TCAV score for two concepts as the dot-product between the CAVs
"""
pos_concept_b_enc = [
model.reparameterize(*model.encode(x)) for x in pos_concept_b
]
neg_concept_b_enc = [
model.reparameterize(*model.encode(x)) for x in neg_concept_b
]
concept_b_classifier = train_binary_linear_classifier(
pos_concept_b_enc, neg_concept_b_enc)
concept_b_gradient = tf.math.l2_normalize(concept_b_classifier.weights[0])
pos_concept_a_enc = [
model.reparameterize(*model.encode(x)) for x in pos_concept_a
]
neg_concept_a_enc = [
model.reparameterize(*model.encode(x)) for x in neg_concept_a
]
concept_a_classifier = train_binary_linear_classifier(
pos_concept_a_enc, neg_concept_a_enc)
concept_a_gradient = tf.math.l2_normalize(concept_a_classifier.weights[0])
tcav_score = tf.tensordot(tf.transpose(concept_b_gradient),
concept_a_gradient,
axes=1)
return tcav_score, concept_a_classifier, concept_b_classifier
def compute_loss(model: VAE, data: tf.Tensor):
"""
Forwards the data through the model and computes the ELBO (Loss)
"""
mean, logvar = model.encode(data)
z = model.reparameterize(mean, logvar)
data_pred = model.decode(z)
# tensorflow vae tutorial monte carlo estimate
cross_ent = tf.losses.mean_squared_error(data, data_pred)
logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2])
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
return -tf.reduce_mean(logpx_z + logpz - logqz_x)
def __init__(self, mdir, device, time_limit, explorer=False):
""" Build vae, rnn, controller and environment. """
self.explorer = explorer
# Load controllers
vae_file, rnn_file, ctrl_file = \
[join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
if self.explorer:
ctrl_file = join(mdir, 'exp', 'best.tar')
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)
]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
print("Loading {} at epoch {} "
"with test loss {}".format(m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device)
self.vae.load_state_dict(vae_state['state_dict'])
# MDRNNCell
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device)
self.mdrnn.load_state_dict(
{k.strip('_l0'): v
for k, v in rnn_state['state_dict'].items()})
self.controller = Controller(LSIZE, RSIZE, ASIZE).to(device)
# load controller if it was previously saved
if exists(ctrl_file):
ctrl_state = torch.load(ctrl_file,
map_location={'cuda:0': str(device)})
print("Loading Controller with reward {}".format(
ctrl_state['reward']))
self.controller.load_state_dict(ctrl_state['state_dict'])
self.env = gym.make('CarRacing-v0')
self.device = device
self.time_limit = time_limit
self.mdrnn_notcell = MDRNN(LSIZE, ASIZE, RSIZE, 5)
self.mdrnn_notcell.to(device)
self.mdrnn_notcell.load_state_dict(rnn_state['state_dict'])
def loop_thru_lambda(trainset, train_loader, testset, test_loader, device,
args):
lambdas = [1, 0.5, 0.2, 0.1, 0.01, 0.001, 0.0001, 1e-5, 1e-6, 0]
for i, lam in enumerate(lambdas):
args.lam = lam
args.fig_name = str(lam) + "_lambda_model.jpg"
args.directory = "samples/" + str(lam) + "-vae.jpg"
model = VAE(args).to(device)
train_vae(model, train_loader, device, args)
test_vae(model, test_loader, device, args)
plt.figure()
for i, lam in enumerate(lambdas):
plt.plot(acc_curves[i], label="Lambda: %f" % lambdas[i])
plt.title("Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend(loc="lower right")
plt.savefig("accuracy.jpg")
plt.figure()
for i, lam in enumerate(lambdas):
plt.plot(loss_curves[i], label="Lambda: %f" % lambdas[i])
plt.title("Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend(loc="upper right")
plt.savefig("loss.jpg")
plt.show()
def main():
device = torch.device("cuda")
# Dataset
test_transform = transforms.Compose([transforms.ToTensor()])
test_dataset = PHD08(options["test_dir"], test_transform)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=options["batch_size"],
shuffle=False,
num_workers=options["workers"],
pin_memory=True,
)
checkpoint_path = "%s.pth" % options["name"]
# Model
if options["mode"] == "Normal":
model = Normal(options["hidden_dim"], options["context_dim"])
else:
model = VAE(
options["hidden_dim"], options["context_dim"], options["addition_dim"]
)
model.load_state_dict(torch.load(checkpoint_path))
model.eval()
summary(model, torch.zeros((1, 1, 28, 28)))
model = model.to(device)
# Train
test_model(options["mode"], device, model, test_loader)
def build_model(self):
'''
self.fea_extractor = torchvision.models.resnet152(pretrained=True)
for p in self.fea_extractor.parameters():
p.requires_grad = False
'''
self.GAE = GAE(n_feat=300,
n_hid=256,
n_latent=128,
adj=self.adj_norm,
dropout=0.0)
self.VAE = VAE()
#pdb.set_trace()
self.CLS = CLS(num_classes=1006)
self.CLS.apply(weights_init)
def __init__(self):
# Loading world model and vae
vae_file, rnn_file, ctrl_file = \
[join("./training", m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
logger.info("Loading {} at epoch {} "
"with test loss {}".format(
m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device).double()
self.vae.load_state_dict(vae_state['state_dict'])
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device).double()
self.mdrnn.load_state_dict(
{k.strip('_l0'): v for k, v in rnn_state['state_dict'].items()})
for p in self.vae.parameters():
p.requires_grad = False
for p in self.mdrnn.parameters():
p.requires_grad = False
self.net = Controller(LSIZE, RSIZE, ASIZE).to(device).double()
# load controller if it was previously saved
if exists(ctrl_file):
ctrl_state = torch.load(ctrl_file, map_location={'cuda:0': str(device)})
logger.info("Loading Controller with reward {}".format(
ctrl_state['reward']))
self.net.load_state_dict(ctrl_state['state_dict'])
self.training_step = 0
self.buffer = np.empty(self.buffer_capacity, dtype=transition)
self.counter = 0
self.optimizer = optim.Adam(self.net.parameters(), lr=1e-3)
class RolloutGeneratorSingle(object):
def __init__(self, mdir, device, controller_model):
""" Run one step.
Load VAE and MDRNN from files
Take the controller (exp/ctrl) an an input, so we can easily change stuff inside the other file.
"""
self.controller = controller_model.to(device)
# Load controllers
vae_file, rnn_file, ctrl_file = \
[join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)
]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
print("Loading {} at epoch {} "
"with test loss {}".format(m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device)
self.vae.load_state_dict(vae_state['state_dict'])
# MDRNNCell
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device)
self.mdrnn.load_state_dict(
{k.strip('_l0'): v
for k, v in rnn_state['state_dict'].items()})
def single_step(obs, hidden):
if params is not None:
load_parameters(params, self.controller)
obs = transform(obs).unsqueeze(0).to(self.device)
# GET ACTION
_, latent_mu, _ = self.vae(obs)
action = self.controller(latent_mu, hidden[0])
_, _, _, _, _, next_hidden = self.mdrnn(action, latent_mu, hidden)
action = action.squeeze().cpu().numpy()
return action, next_obs, next_hidden
def run(img_dims, cont_dim, cat_dims, args):
nets = {'vae': VAE(img_dims, cont_dim, cat_dims, args.temp)}
optimizers = {
'vae':
optim.Adam(nets['vae'].parameters(),
lr=args.eta,
weight_decay=args.weight_decay)
}
return nets, optimizers
def main(args):
base_path = util.make_directory(args.name, args.ow)
device = 'cuda' if torch.cuda.is_available() and len(
args.gpu_ids) > 0 else 'cpu'
start_epoch = 0
# Note: No normalization applied, since RealNVP expects inputs in (0, 1).
trainloader, testloader, in_channels, shape = get_datasets(args)
# Model
print('Building model..')
if args.model == 'ACVAE':
net = MLP_ACVAE(shape,
num_scales=args.num_scales,
in_channels=in_channels,
mid_channels=64,
num_blocks=args.num_blocks)
elif args.model == 'VAE':
net = VAE(shape, args.latent_dim, args.hidden_size)
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net, args.gpu_ids)
cudnn.benchmark = args.benchmark
pass
if args.resume:
# Load checkpoint.
ckpt_path = base_path / 'ckpts'
best_path_ckpt = ckpt_path / 'best.pth.tar'
print(f'Resuming from checkpoint at {best_path_ckpt}')
checkpoint = torch.load(best_path_ckpt)
net.load_state_dict(checkpoint['net'])
global best_loss
best_loss = checkpoint['test_loss']
start_epoch = checkpoint['epoch']
loss_fn = VAELoss if not args.use_ce_loss else VAELossCE
param_groups = util.get_param_groups(net,
args.weight_decay,
norm_suffix='weight_g')
optimizer = optim.Adam(param_groups, lr=args.lr)
for epoch in range(start_epoch, start_epoch + args.num_epochs):
if_save = (epoch == start_epoch) or (epoch == start_epoch +
args.num_epochs - 1)
train(epoch,
net,
trainloader,
device,
optimizer,
loss_fn,
args.max_grad_norm,
base_path,
save=if_save)
test(epoch, net, testloader, device, loss_fn, args.num_samples,
in_channels, base_path, args)
def __init__(self, alpha):
super().__init__()
#self.conv = Conv()
vae_file = join("./vae", 'best.tar')
assert exists(vae_file) , "vae is untrained."
vae_state = torch.load(vae_file, map_location={'cuda:0': str(device)})
logger.info("Loading VAE at epoch {} with test loss {}".format(vae_state['epoch'], vae_state['precision']))
self.vae = VAE(3, LSIZE).to(device).double()
self.vae.load_state_dict(vae_state['state_dict'])
self.W_in = nn.Linear(2*2*256+3, 512, bias=False)
self.W = nn.Linear(512, 512, bias=False)
self.W.weight.data = init_W(512, 512)
self.x_esn = None
self.alpha = alpha
def run(img_dims, cont_dim, cat_dims, args):
nets = {
'vae': VAE(img_dims, cont_dim, cat_dims, args.temp),
'disc': Discriminator(img_dims, cont_dim + sum(cat_dims))
}
optimizers = {
'enc': optim.Adam(nets['vae'].encoder.parameters(), lr=args.eta),
'dec': optim.Adam(nets['vae'].decoder.parameters(), lr=args.eta),
'disc': optim.Adam(nets['disc'].parameters(), lr=args.eta_disc)
}
return nets, optimizers
def standard_vae():
# allow for cuda
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
### 64x64 ###
encoder_params = [(3, 32, 4, 2, 1), (32, 64, 4, 2, 1), (64, 128, 4, 2, 1),
(128, 256, 4, 2, 1), 256 * 4 * 4]
decoder_params = [
256 * 4 * 4, (256, 128, 3, 1, 1), (128, 64, 3, 1, 1),
(64, 32, 3, 1, 1), (32, 3, 3, 1, 1)
]
latent_dim = 100
# set up Model
model = VAE(latent_dim, encoder_params, decoder_params)
model = model.to(device)
return model
def standard_vae32():
# allow for cuda
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
### 32x32 ###
encoder_params = [(3, 32, 5, 2, 2), (32, 64, 5, 2, 2), (64, 128, 5, 2, 2),
128 * 5 * 5]
decoder_params = [
128 * 5 * 5, (128, 64, 5, 2, 2, 1), (64, 32, 5, 2, 2, 1),
(32, 3, 5, 2, 2, 1)
]
latent_dim = 100
# set up Model
model = VAE(latent_dim, encoder_params, decoder_params)
model = model.to(device)
return model
def main(args):
torch.manual_seed(args.seed)
# device placement
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
# load vae ckpt
vae_ckpt = torch.load(args.vae_ckpt_path)
vae = VAE(vae_ckpt['in_dim'], vae_ckpt['cont_dim'],
vae_ckpt['cat_dims'], vae_ckpt['temp']).to(device)
vae.load_state_dict(vae_ckpt['state_dict'])
vae.eval()
_, datautil = _get_dataset(args.dataset)
if not args.indices:
indices = None
else:
indices = [int(e) for e in args.indices.split(',')]
assert len(indices) == args.nb_trav
vis = Visualizer(vae, device,
datautil.testdata, args.nb_trav, indices)
vis.traverse(args.save_path)
def run(args):
torch.manual_seed(args.seed)
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
# load vae ckpt
vae_ckpt = torch.load(args.vae_ckpt_path)
vae = VAE(vae_ckpt['in_dim'], vae_ckpt['cont_dim'], vae_ckpt['cat_dims'],
vae_ckpt['temp']).to(device)
vae.load_state_dict(vae_ckpt['state_dict'])
vae.eval()
attacker = Attacker(vae, args, device)
if args.vis_pair:
try:
src, tar = map(lambda x: int(x), args.vis_pair.split(','))
except:
raise ValueError(
"Argument 'vis-pair' needs to be a pair of digits")
attacker.vis_pair(src, tar)
elif args.vis_bulk:
assert args.vis_bulk > 0 and np.sqrt(args.vis_bulk).is_integer()
attacker.vis_bulk(args.vis_bulk)
else:
if not args.cls_ckpt_path:
raise ValueError("Must specify classifier path when evaluating \
attacks quantitatviely!")
# load classifier
cls = MnistClassifier().to(device)
state_dict = torch.load(args.cls_ckpt_path)
cls.load_state_dict(state_dict)
attacker.full_sweep(cls)
class Ground_truth():
""" For now it only loads the c_vae_path as ground truth"""
def __init__(self, path, color_channels, map_location=None if device is 'cuda' else 'cpu'):
# Number of iterations we trained for
self.net = VAE(device, color_channels=color_channels)
self.net.load_state_dict(torch.load(path))
self.net.to(device)
self.net.eval()
def sample_distribution(self, n, distribution, c=None):
self.net.eval()
# Improve this sampling cut to a nicer line
z = sample_z_distribution(distribution, n, self.net.lat_dim, self.net.device)
return self.sample(n, z, c)
def sample(self, n, z=None, c=None):
with torch.no_grad():
return self.net.sample(n, z, c)
def __init__(self, sizes=(128, 128), batch_size=32, save_models=True):
self.sizes = sizes
self.save_models = save_models
self.batch_size = batch_size
self.trainset = CustomClusteringDataset(
sizes=self.sizes,
path=
"/scratch/datasets/NFLsegment/experiments/vpEB_image_dataset/all_images/"
)
self.trainloader = torch.utils.data.DataLoader(dataset=self.trainset,
batch_size=batch_size,
shuffle=True,
num_workers=16)
mylogger.log(
"Loaded dataset with {} sequences of attention maps".format(
len(self.trainloader)))
self.VAE = VAE(sizes=self.sizes).cuda()
self.oracle = chamber.Oracle()
self.epochs = 2000
def main():
args = parse_args()
torch.manual_seed(args.seed)
device = torch.device("cuda" if args.cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
# load training data
train_loader = load_training_data(args, kwargs)
# load test data
test_loader = load_test_data(args, kwargs)
model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(1, args.epochs + 1):
train(args, model, optimizer, train_loader, device, epoch)
test(args, model, test_loader, device, epoch)
with torch.no_grad():
sample = torch.randn(64, 20).to(device)
sample = model.decode(sample).cpu()
save_image(sample.view(64, 1, 28, 28),
'../results/sample_' + str(epoch) + '.png')
def __init__(self, mdir, device, time_limit):
""" Build vae, rnn, controller and environment. """
# Loading world model and vae
vae_file, rnn_file, ctrl_file = \
[join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
print("Loading {} at epoch {} "
"with test loss {}".format(
m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device)
self.vae.load_state_dict(vae_state['state_dict'])
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device)
self.mdrnn.load_state_dict(
{k.strip('_l0'): v for k, v in rnn_state['state_dict'].items()})
self.controller = Controller(LSIZE, RSIZE, ASIZE).to(device)
# load controller if it was previously saved
if exists(ctrl_file):
ctrl_state = torch.load(ctrl_file, map_location={'cuda:0': str(device)})
print("Loading Controller with reward {}".format(
ctrl_state['reward']))
self.controller.load_state_dict(ctrl_state['state_dict'])
self.env = gym.make('CarRacing-v0')
self.device = device
self.time_limit = time_limit
q_net['z'].mu = Sequential([Dense(100, input_dim=500),
BatchNormalization()])
q_net['z'].var = Sequential([Dense(100, input_dim=500),
BatchNormalization(),
Activation('softplus'),])
p_net['x'].net = Sequential([Dense(500, input_dim=100),
BatchNormalization(),
Activation('relu'),
Dense(500),
BatchNormalization(),
Activation('relu'),
Dense(784),
Activation('sigmoid')])
# Set up VAE
vae = VAE(p_net=p_net, q_net=q_net, shape={'x': (784,)})
vae.compile('adam', compute_log_likelihood=True, verbose=1)
# Train VAE
dataloader = MnistVAE(100)
fh = file_handle('outputs', 'vae', 'standard')
losslog = LossLog(fh=fh)
nll = NegativeLogLikelihood(dataloader,
n_samples=50,
run_every=100,
run_training=True,
run_validation=True,
run_test=True,
display_nelbo=True,
display_epoch=True,
end_line=True,
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("shape of X:",X.shape)
padToMaxLength = 15
batch_size = 1
dim_phoneme_embeddings = 16
original_dim = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim = 500
epsilon_std = 0.01
nb_epoch =100
vae = VAE(latent_dim=latent_dim,
original_dim=original_dim,
intermediate_dim=intermediate_dim,
batch_size=batch_size,
epsilon_std=epsilon_std)
vae.load_weights("/Users/marlon/Documents/BA/Workspace/Workspace/saved_weights/ASJP_2000/pipeline_asjp_2000.h5")
print("VECTORIZE WORDS")
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("EMBED WORDS")
embeddings = vae.embed(X)
print("CHECK CORRELATION")
from scipy.stats import pearsonr,spearmanr
print("pearson",pearsonr(embeddings[:,0], embeddings[:,1]))
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("EMBED WORDS")
padToMaxLength = 15
batch_size = 271
dim_phoneme_embeddings = 16
original_dim = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim = 500
epsilon_std = 0.01
nb_epoch =100
vae = VAE(latent_dim=latent_dim,
original_dim=original_dim,
intermediate_dim=intermediate_dim,
batch_size=batch_size,
epsilon_std=epsilon_std)
vae.load_weights("../../saved_weights/ASJP_2000/pipeline_asjp_2000.h5")
embeddings = vae.embed(X)
print("plotting")
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("white")
from binary_phoneme_features import BinaryPhonemeFeatures
bpf = BinaryPhonemeFeatures()
class RolloutGenerator(object):
""" Utility to generate rollouts.
Encapsulate everything that is needed to generate rollouts in the TRUE ENV
using a controller with previously trained VAE and MDRNN.
:attr vae: VAE model loaded from mdir/vae
:attr mdrnn: MDRNN model loaded from mdir/mdrnn
:attr controller: Controller, either loaded from mdir/ctrl or randomly
initialized
:attr env: instance of the CarRacing-v0 gym environment
:attr device: device used to run VAE, MDRNN and Controller
:attr time_limit: rollouts have a maximum of time_limit timesteps
"""
def __init__(self, mdir, device, time_limit):
""" Build vae, rnn, controller and environment. """
# Loading world model and vae
vae_file, rnn_file, ctrl_file = \
[join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]
assert exists(vae_file) and exists(rnn_file),\
"Either vae or mdrnn is untrained."
vae_state, rnn_state = [
torch.load(fname, map_location={'cuda:0': str(device)})
for fname in (vae_file, rnn_file)]
for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
print("Loading {} at epoch {} "
"with test loss {}".format(
m, s['epoch'], s['precision']))
self.vae = VAE(3, LSIZE).to(device)
self.vae.load_state_dict(vae_state['state_dict'])
self.mdrnn = MDRNNCell(LSIZE, ASIZE, RSIZE, 5).to(device)
self.mdrnn.load_state_dict(
{k.strip('_l0'): v for k, v in rnn_state['state_dict'].items()})
self.controller = Controller(LSIZE, RSIZE, ASIZE).to(device)
# load controller if it was previously saved
if exists(ctrl_file):
ctrl_state = torch.load(ctrl_file, map_location={'cuda:0': str(device)})
print("Loading Controller with reward {}".format(
ctrl_state['reward']))
self.controller.load_state_dict(ctrl_state['state_dict'])
self.env = gym.make('CarRacing-v0')
self.device = device
self.time_limit = time_limit
def get_action_and_transition(self, obs, hidden):
""" Get action and transition.
Encode obs to latent using the VAE, then obtain estimation for next
latent and next hidden state using the MDRNN and compute the controller
corresponding action.
:args obs: current observation (1 x 3 x 64 x 64) torch tensor
:args hidden: current hidden state (1 x 256) torch tensor
:returns: (action, next_hidden)
- action: 1D np array
- next_hidden (1 x 256) torch tensor
"""
_, latent_mu, _ = self.vae(obs)
action = self.controller(latent_mu, hidden[0])
_, _, _, _, _, next_hidden = self.mdrnn(action, latent_mu, hidden)
return action.squeeze().cpu().numpy(), next_hidden
def rollout(self, params, render=False):
""" Execute a rollout and returns minus cumulative reward.
Load :params: into the controller and execute a single rollout. This
is the main API of this class.
:args params: parameters as a single 1D np array
:returns: minus cumulative reward
"""
# copy params into the controller
if params is not None:
load_parameters(params, self.controller)
obs = self.env.reset()
# This first render is required !
self.env.render()
hidden = [
torch.zeros(1, RSIZE).to(self.device)
for _ in range(2)]
cumulative = 0
i = 0
while True:
obs = transform(obs).unsqueeze(0).to(self.device)
action, hidden = self.get_action_and_transition(obs, hidden)
obs, reward, done, _ = self.env.step(action)
if render:
self.env.render()
cumulative += reward
if done or i > self.time_limit:
return - cumulative
i += 1
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("shape of X:",X.shape)
padToMaxLength = 15
batch_size = 1
dim_phoneme_embeddings = 16
original_dim = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim = 500
epsilon_std = 0.01
nb_epoch =100
vae = VAE(latent_dim=latent_dim,
original_dim=original_dim,
intermediate_dim=intermediate_dim,
batch_size=batch_size,
epsilon_std=epsilon_std)
vae.load_weights("/Users/marlon/Documents/BA/Workspace/Workspace/saved_weights/ASJP_2000/pipeline_asjp_2000.h5")
print("VECTORIZE WORDS")
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
embeddings = vae.embed(X)
print("CHECK CORRELATION")
from scipy.stats import pearsonr,spearmanr
print("pearson",pearsonr(embeddings[:,0], embeddings[:,1]))
print("spearman",spearmanr(embeddings[:,0], embeddings[:,1]))
from models import VAE
batch_size = 271
dim_phoneme_embeddings = 16
original_dim_phono = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim_phono = 500
epsilon_std = 0.01
nb_epoch = 100
vae = VAE(latent_dim=latent_dim,
original_dim=original_dim_phono,
intermediate_dim=intermediate_dim_phono,
batch_size=batch_size,
epsilon_std=epsilon_std)
"""
FITTING MODELS
"""
print("FITTING MODELS")
print(word_matrices.shape,
#concepts_oneHots.shape,geo_words.shape
)
vae.fit(word_matrices,
nb_epoch=nb_epoch)
embeddings = vae.embed(X=word_matrices)
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("shape of X:",X.shape)
padToMaxLength = 15
batch_size = 1
dim_phoneme_embeddings = 16
original_dim = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim = 500
epsilon_std = 0.01
nb_epoch =100
vae = VAE(latent_dim=latent_dim,
original_dim=original_dim,
intermediate_dim=intermediate_dim,
batch_size=batch_size,
epsilon_std=epsilon_std)
vae.load_weights("/Users/marlon/Documents/BA/Workspace/Workspace/saved_weights/ASJP_2000/pipeline_asjp_2000.h5")
print("VECTORIZE WORDS")
bpf = BinaryPhonemeFeatures()
X = np.array([bpf.encodeWord(word, padToMaxLength=padToMaxLength).flatten() for word in words])
print("EMBED WORDS")
embeddings = vae.embed(X)
print("CHECK CORRELATION")
from scipy.stats import pearsonr,spearmanr
print("pearson",pearsonr(embeddings[:,0], embeddings[:,1]))
from models import VAE
batch_size = 421
dim_phoneme_embeddings = 16
original_dim_phono = dim_phoneme_embeddings * padToMaxLength
latent_dim = 2
intermediate_dim = 500
epsilon_std = 0.01
nb_epoch = 1000
vae = VAE(latent_dim=latent_dim,
original_dim=X.shape[1],
intermediate_dim=intermediate_dim,
batch_size=batch_size,
epsilon_std=epsilon_std)
"""
FITTING MODELS
"""
print("FITTING MODELS")
print(word_matrices.shape,
#concepts_oneHots.shape,geo_words.shape
)
vae.fit(X,
nb_epoch=nb_epoch)
embeddings = vae.embed(X=X)
