def _compute_metric(self, state: State):
criterion_kwargs = {
self.real_data_criterion_key:
state.batch_in[self.input_key],
self.fake_data_criterion_key:
state.batch_out[self.output_key],
self.critic_criterion_key:
state.model[self.critic_model_key],
self.condition_args_criterion_key:
[state.batch_in[key] for key in self.condition_keys]
}
criterion = state.get_attr("criterion", self.criterion_key)
return criterion(**criterion_kwargs)
def on_batch_end(self, state: State):
"""On batch end event"""
super().on_batch_end(state)
if not state.is_train_loader:
return
optimizer = state.get_attr(key="optimizer",
inner_key=self.optimizer_key)
need_gradient_step = \
self._accumulation_counter % self.accumulation_steps == 0
if need_gradient_step:
for group in optimizer.param_groups:
for param in group["params"]:
param.data.clamp_(min=-self.weight_clamp_value,
max=self.weight_clamp_value)