empose/nn/models.py [233:262]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
                                    batch.seq_lengths, batch.marker_masks)

        if self.estimate_shape:
            shape_loss = padded_loss(batch.shapes.unsqueeze(1).repeat((1, shape_hat.shape[1], 1)),
                                     shape_hat, self.shape_loss, batch.seq_lengths)
        else:
            shape_loss = torch.zeros(1).to(device=C.DEVICE)

        if self.do_fk:
            joints_gt = batch.joints_gt.reshape(batch.batch_size, batch.seq_length, -1, 3)
            joints_hat = model_out['joints_hat'].reshape(batch.batch_size, batch.seq_length, -1, 3)
            fk_loss = reconstruction_loss(joints_gt, joints_hat, batch.seq_lengths, batch.marker_masks)
        else:
            fk_loss = torch.zeros(1).to(device=C.DEVICE)

        total_loss = pose_loss + root_pose_loss + shape_loss + self.fk_loss_weight * fk_loss

        loss_vals = {'pose': pose_loss.cpu().item(),
                     'root_pose': root_pose_loss.cpu().item(),
                     'shape': shape_loss.cpu().item(),
                     'fk': fk_loss.cpu().item(),
                     'total_loss': total_loss.cpu().item()}

        if writer is not None:
            self.log_loss_vals(loss_vals, writer, global_step)

        if self.training:
            total_loss.backward()

        return total_loss, loss_vals
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



empose/nn/models.py [337:366]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
                                    batch.seq_lengths, batch.marker_masks)

        if self.estimate_shape:
            shape_loss = padded_loss(batch.shapes.unsqueeze(1).repeat((1, shape_hat.shape[1], 1)),
                                     shape_hat, self.shape_loss, batch.seq_lengths)
        else:
            shape_loss = torch.zeros(1).to(device=C.DEVICE)

        if self.do_fk:
            joints_gt = batch.joints_gt.reshape(batch.batch_size, batch.seq_length, -1, 3)
            joints_hat = model_out['joints_hat'].reshape(batch.batch_size, batch.seq_length, -1, 3)
            fk_loss = reconstruction_loss(joints_gt, joints_hat, batch.seq_lengths, batch.marker_masks)
        else:
            fk_loss = torch.zeros(1).to(device=C.DEVICE)

        total_loss = pose_loss + root_pose_loss + shape_loss + self.fk_loss_weight * fk_loss

        loss_vals = {'pose': pose_loss.cpu().item(),
                     'root_pose': root_pose_loss.cpu().item(),
                     'shape': shape_loss.cpu().item(),
                     'fk': fk_loss.cpu().item(),
                     'total_loss': total_loss.cpu().item()}

        if writer is not None:
            self.log_loss_vals(loss_vals, writer, global_step)

        if self.training:
            total_loss.backward()

        return total_loss, loss_vals
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -