def _train_step(self, loader):
batch = self._request_data(loader)
if batch is None:
return 0, None, None
inp, target = batch
self.optimizer.zero_grad()
prediction = self.model(inp)
losses = []
loss_metrics = {}
for name, criterion in self.criteria.items():
loss = criterion(prediction, target)
losses.append(loss)
loss_metrics['loss_' + name] = get_loss_metric(loss.data[0])
total_loss = torch.sum(torch.cat(losses) * self.loss_weights)
total_loss.backward()
self.optimizer.step()
loss_metrics['loss'] = get_loss_metric(total_loss.data[0])
data = (inp, prediction, target)
return 1, loss_metrics, data
def _train_step(self, loader):
batch = self._request_data(loader)
if batch is None:
return 0, None, None
self.optimizer.zero_grad()
out_model = self.model(*self.train_model_input_fn(batch))
losses = []
loss_metrics = {}
for name, criterion in self.criteria.items():
loss = criterion(out_model, batch)
losses.append(loss)
loss_metrics['loss_' + name] = get_loss_metric(loss.data[0])
total_loss = torch.sum(torch.cat(losses) * self.loss_weights)
total_loss.backward()
self.optimizer.step()
loss_metrics['loss'] = get_loss_metric(total_loss.data[0])
data = (batch, out_model)
return 1, loss_metrics, data
def _train_step(self, loader):
batch = self._request_data(loader)
if batch is None:
return 0, None, None
inp, target = batch
# Propagate fake image through discriminator
out_gen = self.gen(inp)
out_disc_fake = self.disc(self.disc_input_fn(out_gen, inp, detach=True))
# Propagate real images through discriminator
out_disc_real = self.disc(self.disc_input_fn(target, inp, detach=True))
loss_metrics = {}
disc_losses = []
# Compute discriminator losses
for name, criterion in self.disc_adv_criteria.items():
loss = criterion(out_disc_fake, out_disc_real)
disc_losses.append(loss)
loss_metrics['disc_loss_' + name] = get_loss_metric(loss.data[0])
# Propagate again with non-detached input to allow gradients on the
# generator
out_disc_fake = self.disc(self.disc_input_fn(out_gen, inp, detach=False))
# Compute adversarial generator losses from discriminator output
# Order matters: first compute adversarial losses for generator, then
# the other generator losses. Otherwise the loss weights will not match
gen_losses = []
for name, criterion in self.gen_adv_criteria.items():
loss = criterion(out_disc_fake)
gen_losses.append(loss)
loss_metrics['gen_loss_' + name] = get_loss_metric(loss.data[0])
# Compute generator losses on prediction and target image
for name, criterion in self.gen_criteria.items():
loss = criterion(out_gen, target)
gen_losses.append(loss)
loss_metrics['gen_loss_' + name] = get_loss_metric(loss.data[0])
# Perform updates
total_disc_loss = self._update_step(self.disc_optimizer,
disc_losses,
self.disc_loss_weights)
total_gen_loss = self._update_step(self.gen_optimizer,
gen_losses,
self.gen_loss_weights)
loss_metrics['disc_loss'] = get_loss_metric(total_disc_loss.data[0])
loss_metrics['gen_loss'] = get_loss_metric(total_gen_loss.data[0])
return 1, loss_metrics, (inp, out_gen, target, out_disc_fake)
def _val_step(self, loader, compute_metrics=True):
batch = self._request_data(loader, volatile=True)
if batch is None:
return None, None
out_model = self.model(*self.test_model_input_fn(batch))
loss_metrics = {}
if compute_metrics:
for name, criterion in self.criteria.items():
loss = criterion(out_model, batch)
loss_metrics['loss_' + name] = get_loss_metric(loss.data[0])
return loss_metrics, (batch, out_model)
def _val_step(self, loader, compute_metrics=True):
batch = self._request_data(loader, volatile=True)
if batch is None:
return None, None
inp, target = batch
prediction = self.model(inp)
loss_metrics = {}
if compute_metrics:
for name, criterion in self.criteria.items():
loss = criterion(prediction, target)
loss_metrics['loss_' + name] = get_loss_metric(loss.data[0])
return loss_metrics, (inp, prediction, target)
def _val_step(self, loader, compute_metrics=True):
batch = self._request_data(loader, volatile=True)
if batch is None:
return None, None
inp, target = batch
prediction = self.gen(inp)
loss_metrics = {}
if compute_metrics:
# Only compute the standard losses here, adversarial losses don't make
# to much sense
for name, criterion in self.gen_criteria.items():
loss = criterion(prediction, target)
loss_metrics['gen_loss_' + name] = get_loss_metric(loss.data[0])
return loss_metrics, (inp, prediction, target, None)
def _val_step(self, loader, compute_metrics=True):
batch = self._request_data(loader, volatile=True)
if batch is None:
return None, None
gen_inp = self.test_model_input_fn(batch)
out_gen = self.gen(*gen_inp)
if self.disc is not None:
out_disc_fake = self.disc(
self.val_disc_input_fn(out_gen,
gen_inp[0],
out_gen,
is_real_input=False,
detach=True))
target = batch['target']
out_disc_real = self.disc(
self.val_disc_input_fn(target,
gen_inp[0],
out_gen,
is_real_input=True,
detach=True))
else:
out_disc_fake = None
out_disc_real = None
loss_metrics = {}
if compute_metrics:
# Only compute the standard losses here, adversarial losses don't make
# to much sense
for name, criterion in self.gen_criteria.items():
loss = criterion(out_gen, batch)
loss_metrics['gen_loss_' + name] = get_loss_metric(
loss.data[0])
return loss_metrics, (batch, out_gen, out_disc_fake, out_disc_real)
def _train_multiple_steps(self, loader):
"""Train generator and discriminator for multiple steps at once"""
last_batch = None
max_updates = max(self.disc_updates_per_step,
self.gen_updates_per_step)
# Deque input data upfront (this could lead to memory problems)
batches = []
for _ in range(max_updates):
batch = self._request_data(loader)
if batch is None:
break
batches.append(batch)
gen_uses_feature_matching = 'FeatureMatching' in self.gen_adv_criteria
loss_metrics = {}
# Train discriminator
for idx, batch in enumerate(batches[:self.disc_updates_per_step]):
if not self.discriminator_enabled:
continue
last_batch = batch
# Propagate fake image through discriminator
gen_inp = self.train_model_input_fn(batch)
out_gen = self.gen(*gen_inp)
out_disc_fake = self.disc(
self.disc_input_fn(out_gen,
gen_inp[0],
out_gen,
is_real_input=False,
detach=True))
# Propagate real images through discriminator
target = batch['target']
out_disc_real = self.disc(
self.disc_input_fn(target,
gen_inp[0],
out_gen,
is_real_input=True,
detach=True))
disc_losses = []
# Compute discriminator losses
for name, criterion in self.disc_adv_criteria.items():
loss = criterion(out_disc_fake, out_disc_real)
disc_losses.append(loss)
accumulate_metric(loss_metrics, 'disc_loss_' + name,
get_loss_metric(loss.data[0]))
# Perform discriminator update
total_disc_loss = self._update_step(self.disc_optimizer,
disc_losses,
self.disc_loss_weights)
accumulate_metric(loss_metrics, 'disc_loss',
get_loss_metric(total_disc_loss.data[0]))
if idx < len(batches) - 1 and idx < self.disc_updates_per_step - 1:
del out_gen
del out_disc_real
del out_disc_fake
elif self.generator_enabled:
del out_gen
del out_disc_fake
# Train generator
for idx, batch in enumerate(batches[:self.gen_updates_per_step]):
if not self.generator_enabled:
continue
last_batch = batch
gen_losses = []
gen_inp = self.train_model_input_fn(batch)
out_gen = self.gen(*gen_inp)
if self.discriminator_enabled:
# Propagate again with non-detached input to allow gradients on the
# generator
out_disc_fake = self.disc(
self.disc_input_fn(out_gen,
gen_inp[0],
out_gen,
is_real_input=False,
detach=False))
if gen_uses_feature_matching:
# Only need to compute the discriminator output for real targets if we
# use feature matching loss
target = batch['target']
out_disc_real = self.disc(
self.disc_input_fn(target,
gen_inp[0],
out_gen,
is_real_input=True,
detach=True))
else:
out_disc_real = None
# Compute adversarial generator losses from discriminator output
# Order matters: first compute adversarial losses for generator, then
# the other generator losses. Otherwise the loss weights will not match
for name, criterion in self.gen_adv_criteria.items():
loss = criterion(out_disc_fake, out_disc_real)
gen_losses.append(loss)
accumulate_metric(loss_metrics, 'gen_loss_' + name,
get_loss_metric(loss.data[0]))
# Compute generator losses on prediction and target image
for name, criterion in self.gen_criteria.items():
loss = criterion(out_gen, batch)
gen_losses.append(loss)
accumulate_metric(loss_metrics, 'gen_loss_' + name,
get_loss_metric(loss.data[0]))
# Perform generator update
total_gen_loss = self._update_step(self.gen_optimizer, gen_losses,
self.gen_loss_weights)
accumulate_metric(loss_metrics, 'gen_loss',
get_loss_metric(total_gen_loss.data[0]))
if idx < len(batch) - 1 and idx < self.gen_updates_per_step - 1:
del out_gen
if self.discriminator_enabled:
del out_disc_fake
if out_disc_real is not None:
del out_disc_real
# For simplicity, we just return the last batch of data
# This is a bit of a smell, as our metrics will only be on this last batch
# of data, whereas the loss metrics are averaged over all updates
if len(batches) > 0:
avg_loss_metrics = {
name: metric.average()
for name, metric in loss_metrics.items()
}
if not self.discriminator_enabled:
out_disc_fake = None
out_disc_real = None
data = (last_batch, out_gen, out_disc_fake, out_disc_real)
else:
avg_loss_metrics = None
data = None
return len(batches), avg_loss_metrics, data
def _train_single_step(self, loader):
batch = self._request_data(loader)
if batch is None:
return 0, None, None
loss_metrics = {}
gen_inp = self.train_model_input_fn(batch)
out_gen = self.gen(*gen_inp)
if self.discriminator_enabled:
# Propagate fake image through discriminator
out_disc_fake = self.disc(
self.disc_input_fn(out_gen,
gen_inp[0],
out_gen,
is_real_input=False,
detach=True))
# Propagate real images through discriminator
target = batch['target']
out_disc_real = self.disc(
self.disc_input_fn(target,
gen_inp[0],
out_gen,
is_real_input=True,
detach=True))
disc_losses = []
# Compute discriminator losses
for name, criterion in self.disc_adv_criteria.items():
loss = criterion(out_disc_fake, out_disc_real)
disc_losses.append(loss)
loss_metrics['disc_loss_' + name] = get_loss_metric(
loss.data[0])
if self.generator_enabled:
gen_losses = []
if self.discriminator_enabled:
# Propagate again with non-detached input to allow gradients on the
# generator
out_disc_fake = self.disc(
self.disc_input_fn(out_gen,
gen_inp[0],
out_gen,
is_real_input=False,
detach=False))
# Compute adversarial generator losses from discriminator output
# Order matters: first compute adversarial losses for generator, then
# the other generator losses. Otherwise the loss weights will not match
for name, criterion in self.gen_adv_criteria.items():
loss = criterion(out_disc_fake, out_disc_real)
gen_losses.append(loss)
loss_metrics['gen_loss_' + name] = get_loss_metric(
loss.data[0])
# Compute generator losses on prediction and target image
for name, criterion in self.gen_criteria.items():
loss = criterion(out_gen, batch)
gen_losses.append(loss)
loss_metrics['gen_loss_' + name] = get_loss_metric(
loss.data[0])
# Perform updates
if self.discriminator_enabled:
total_disc_loss = self._update_step(self.disc_optimizer,
disc_losses,
self.disc_loss_weights)
loss_metrics['disc_loss'] = get_loss_metric(
total_disc_loss.data[0])
if self.generator_enabled:
total_gen_loss = self._update_step(self.gen_optimizer, gen_losses,
self.gen_loss_weights)
loss_metrics['gen_loss'] = get_loss_metric(total_gen_loss.data[0])
if not self.discriminator_enabled:
out_disc_fake = None
out_disc_real = None
return 1, loss_metrics, (batch, out_gen, out_disc_fake, out_disc_real)