def train_iteration()

in src/trainer.py [0:0]

• 28 lines of code
• 1 McCabe index (conditional complexity)

    def train_iteration(self, data):
        '''
        one optimization step
        :param data: tuple of tensors containing mono, binaural, and quaternion data
        :return: dict containing values for all different losses
        '''
        # forward
        self.optimizer.zero_grad()

        mono, binaural, quats = data
        mono, binaural, quats = mono.cuda(), binaural.cuda(), quats.cuda()
        prediction = self.net.forward(mono, quats)
        l2 = self.l2_loss(prediction["output"], binaural)
        phase = self.phase_loss(prediction["output"], binaural)
        intermediate_binaural = th.cat([binaural] * len(prediction["intermediate"]), dim=1)
        intermediate_prediction = th.cat(prediction["intermediate"], dim=1)
        intermediate_l2 = self.l2_loss(intermediate_prediction, intermediate_binaural)
        intermediate_phase = self.phase_loss(intermediate_prediction, intermediate_binaural)

        loss = (l2 + intermediate_l2) * self.config["loss_weights"]["l2"] + \
               (phase + intermediate_phase) * self.config["loss_weights"]["phase"]

        # update model parameters
        loss.backward()
        self.optimizer.step()
        self.total_iters += 1

        return {
            "l2": l2,
            "phase": phase,
            "intermediate_l2": intermediate_l2,
            "intermediate_phase": intermediate_phase,
            "accumulated_loss": loss,
        }