def sample(self):
buffer_iter = iter(self.buffer_loader(self.buffer))
data, *args = to_device(self.data_key(next(buffer_iter)), self.device)
self.score.eval()
data = detach(self.integrator.integrate(self.score, data, *args))
self.score.train()
detached = to_device(data, "cpu")
args = detach(args)
update = self.decompose_batch(detached, *args)
self.buffer.update(update)
return to_device((detached, *args), self.device)
def loss(self, posterior, prior, prior_target, reconstruction, target, args): ce = self.reconstruction_loss(reconstruction, target) kld = self.divergence_loss(posterior, prior_target) kld_prior = self.divergence_loss(detach(posterior), prior) loss_val = self.reconstruction_weight * ce + self.divergence_weight * (kld - kld.detach() + kld_prior) self.current_losses["reconstruction-log-likelihood"] = float(ce) self.current_losses["kullback-leibler-divergence"] = float(kld_prior) return loss_val
def stack_loss(self, task_args, scale_args):
task_loss = 0.0
count = 0
for idx, ta in enumerate(zip(*task_args)):
level_loss = self.task_loss(*ta)
self.current_losses[f"level {idx}"] = float(level_loss)
task_loss += level_loss
count += 1
task_loss = task_loss / count
self.current_losses["task"] = float(task_loss)
prior_loss = match(MatchableList(scale_args.prior),
MatchableList(detach(scale_args.posterior)))
self.current_losses["prior"] = float(prior_loss)
policy_loss = self.policy_loss(task_args, scale_args)
self.current_losses["policy"] = float(policy_loss)
result = self.task_weight * task_loss
result = result + self.prior_weight * prior_loss
result = result + self.policy_weight * policy_loss
return result
def run_energy(self, data):
make_differentiable(data)
input_data, *data_args = self.data_key(data)
real_result = self.score(input_data, *data_args)
# sample after first pass over real data, to catch
# possible batch-norm shenanigans without blowing up.
fake = self.sample()
if self.step_id % self.report_interval == 0:
detached, *args = self.data_key(to_device(fake, "cpu"))
self.each_generate(detach(detached), *args)
make_differentiable(fake)
input_fake, *fake_args = self.data_key(fake)
self.score.eval()
fake_update_result = None
if self.sampler_likelihood:
# Sampler log likelihood:
fake_result = self.score(input_fake, *fake_args)
fake_update = self.integrator.step(self.score, input_fake,
*fake_args)
self.update_target()
fake_update_result = self.target_score(fake_update, *fake_args)
comparison = None
if self.maximum_entropy:
# Sampler entropy:
compare_update = fake_update
compare_target = to_device(
self.data_key(self.buffer.sample(self.batch_size))[0],
compare_update.device)
if hasattr(self.score, "embed"):
compare_update = self.target_score.embedding(compare_update)
compare_target = self.target_score.embedding(compare_target)
comparison = self.compare(compare_update, compare_target.detach())
self.score.train()
return real_result, fake_result, fake_update_result, comparison
def integrate(self, score, data, *args):
for idx in range(self.steps):
data = self.step(score, data, *args)
data = detach(data)
return data
def each_step(self):
super().each_step()
if self.step_id % self.report_interval == 0 and self.step_id != 0:
data, *args = detach(to_device(self.sample(), "cpu"))
self.each_generate(data, *args)