Method training_step

train.py:286–388

(self, batch, batch_idx)

            pin_memory=True)

    def training_step(self, batch, batch_idx):
        # Fetch training data.
        rgbs = batch['rgbs']
        ts_w = batch['ts_w']
        grid = batch['grid']
        mk = batch['masks']
        flows = batch['flows']
        grid_c = batch['grid_c']
        ref_batch = batch['reference']
        self.seq_len = batch['seq_len']

        loss = 0
        rgbs_flattend = rearrange(rgbs, 'b h w c -> (b h w) c')

        # Forward the model.
        ret = self.forward(ts_w,
                            grid,
                            self.hparams.encode_w,
                            self.global_step,
                            flows=flows)

        # Mannually set a reference frame.
        if self.hparams.ref_step < 0: self.hparams.step = 1e10
        if (self.hparams.ref_idx is not None
                and self.global_step < self.hparams.ref_step):
            rgbs_c_flattend = rearrange(ref_batch[0],
                                        'b h w c -> (b h w) c')
            ret_c = self(ts_w, grid, False, self.global_step, flows=flows)

        # Loss computation.
        for i in range(self.num_models):
            results = ret.rgbs[i]
            mk_t = rearrange(mk[i], 'b h w c -> (b h w) c')
            mk_t = mk_t.sum(dim=-1) > 0.05

            if (self.hparams.ref_idx is not None
                and self.global_step < self.hparams.ref_step):
                mk_c_t = rearrange(ref_batch[1][i], 'b h w c -> (b h w) c')
                mk_c_t = mk_c_t.sum(dim=-1) > 0.05

            # Background regularization.
            if self.hparams.bg_loss:
                mk1 = torch.logical_not(mk_t)
                if self.hparams.self_bg:
                    grid_flattened = rgbs_flattend
                else:
                    grid_flattened = rearrange(grid, 'b n c -> (b n) c')
                    grid_flattened = torch.cat(
                        [grid_flattened, grid_flattened[:, :1]], -1)

            if self.hparams.bg_loss and self.hparams.mask_dir:
                loss = loss + self.hparams.bg_loss * self.color_loss(
                    results[mk1], grid_flattened[mk1])

            # MSE color loss.
            loss = loss + self.color_loss(results[mk_t],
                                            rgbs_flattend[mk_t])