def process_batch_data(self, batch):
"""
Processes the batch returned by the dataloader iterator
:param batch: object returned by the dataloader iterator
:return: tuple of Torch Variable objects
"""
inputs, labels = batch
inputs = to_cuda_variable(inputs)
labels = to_cuda_variable(labels)
return inputs, labels
def forward_test(self, measure_score_tensor: Variable):
"""
Implements the forward pass of the VAE
:param measure_score_tensor: torch Variable,
(batch_size, num_measures, measure_seq_length)
:return: torch Variable,
(batch_size, measure_seq_length, self.num_notes)
"""
# check input
batch_size, num_measures, seq_len = measure_score_tensor.size()
assert(seq_len == self.num_ticks_per_measure)
# compute output of encoding layer
z = []
for i in range(num_measures):
z_dist = self.encoder(measure_score_tensor[:, i, :])
z.append(z_dist.rsample().unsqueeze(1))
z_tilde = torch.cat(z, 1)
# compute output of decoding layer
weights = []
samples = []
dummy_measure_tensor = to_cuda_variable(torch.zeros(batch_size, seq_len))
for i in range(num_measures):
w, s = self.decoder(
z=z_tilde[:, i, :],
score_tensor=dummy_measure_tensor,
train=False
)
samples.append(s)
weights.append(w.unsqueeze(1))
samples = torch.cat(samples, 2)
weights = torch.cat(weights, 1)
return weights, samples
def process_batch_data(self, batch, test=False):
"""
Processes the batch returned by the dataloader iterator
:param batch: object returned by the dataloader iterator
:return: tuple of Torch Variable objects
"""
inputs, labels = batch
b = inputs.size(0)
inputs = to_cuda_variable(inputs)
labels = to_cuda_variable(labels)
if test:
return inputs, labels
else:
if b % 2 != 0:
b -= 1
batch_1 = (inputs[:b // 2], labels[:b // 2])
batch_2 = (inputs[b // 2:b], labels[b // 2:b])
return (batch_1, batch_2)
def hidden_init(self, batch_size):
"""
Initializes the hidden state of the encoder GRU
:param batch_size: int
:return: torch tensor,
(self.num_encoder_layers x self.num_directions, batch_size, self.encoder_hidden_size)
"""
hidden = torch.zeros(self.num_layers * self.num_directions,
batch_size,
self.rnn_hidden_size
)
return to_cuda_variable(hidden)
def loss_and_acc_test(self, data_loader):
mean_loss = 0
mean_accuracy = 0
for sample_id, batch in tqdm(enumerate(data_loader)):
inputs, _ = self.process_batch_data(batch)
inputs = to_cuda_variable(inputs)
# compute forward pass
outputs, _, _, _, _ = self.model(inputs)
# compute loss
recons_loss = self.reconstruction_loss(inputs, outputs,
self.dec_dist)
loss = recons_loss
# compute mean loss and accuracy
mean_loss += to_numpy(loss.mean())
accuracy = self.mean_accuracy(weights=torch.sigmoid(outputs),
targets=inputs)
mean_accuracy += to_numpy(accuracy)
mean_loss /= len(data_loader)
mean_accuracy /= len(data_loader)
return (mean_loss, mean_accuracy)
def __init__(
self,
dataset,
model: FactorVAE,
lr=1e-4,
reg_type: Tuple[str] = None,
beta=0.001,
capacity=0.0,
gamma=10,
rand=0,
):
super(FactorVAETrainer, self).__init__(dataset, model, lr)
self.attr_dict = LATENT_ATTRIBUTES
self.reverse_attr_dict = {
v: k for k, v in self.attr_dict.items()
}
self.metrics = {}
self.beta = beta
self.start_beta = 0.0
self.cur_beta = self.start_beta
self.capacity = to_cuda_variable(torch.FloatTensor([capacity]))
self.gamma = gamma
self.cur_gamma = 0
self.delta = 0.0
self.cur_epoch_num = 0
self.warm_up_epochs = -1
self.num_iterations = 100000
self.exp_rate = np.log(1 + self.beta) / self.num_iterations
self.anneal_iterations = 0
self.reg_type = reg_type
self.reg_dim = ()
self.rand_seed = rand
torch.manual_seed(self.rand_seed)
np.random.seed(self.rand_seed)
self.trainer_config = f'_b_{self.beta}_c_{capacity}_g_{self.gamma}_r_{self.rand_seed}_nowarm_'
self.model.update_trainer_config(self.trainer_config)
# re-instantiate optimizers
self.optimizer = optim.Adam(model.VAE.parameters(), lr=lr)
self.D_optim = optim.Adam(model.discriminator.parameters(), lr=1e-4, betas=(0.8, 0.9))
def plot_latent_interpolations(self, attr_str, dim, num_points=10):
x1 = torch.linspace(-4, 4.0, num_points)
_, _, data_loader = self.dataset.data_loaders(batch_size=1)
for sample_id, batch in tqdm(enumerate(data_loader)):
if sample_id in [0, 1, 2]:
inputs, labels = self.process_batch_data(batch)
inputs = to_cuda_variable(inputs)
recons, _, _, z, _ = self.model(inputs)
recons = torch.sigmoid(recons)
z = z.repeat(num_points, 1)
z[:, dim] = x1.contiguous()
outputs = torch.sigmoid(self.model.decode(z))
# save interpolation
save_filepath = os.path.join(
Trainer.get_save_dir(self.model),
f'latent_interpolations_{attr_str}_{sample_id}.png')
save_image(outputs.cpu(),
save_filepath,
nrow=num_points,
pad_value=1.0)
# save original image
org_save_filepath = os.path.join(
Trainer.get_save_dir(self.model),
f'original_{sample_id}.png')
save_image(inputs.cpu(),
org_save_filepath,
nrow=1,
pad_value=1.0)
# save reconstruction
recons_save_filepath = os.path.join(
Trainer.get_save_dir(self.model),
f'recons_{sample_id}.png')
save_image(recons.cpu(),
recons_save_filepath,
nrow=1,
pad_value=1.0)
if sample_id == 5:
break