Method validation_step

train.py:390–421

(self, batch, batch_idx)

        return loss

    def validation_step(self, batch, batch_idx):
        rgbs = batch['rgbs']
        ts_w = batch['ts_w']
        grid = batch['grid']
        mk = batch['masks']
        grid_c = grid  # batch['grid_c']
        self.seq_len = batch['seq_len']
        ret = self(ts_w, grid, self.hparams.encode_w, self.global_step)
        ret_c = self(ts_w, grid_c, False, self.global_step)

        log = {}
        W, H = self.hparams.img_wh

        rgbs_flattend = rearrange(rgbs, 'b h w c -> (b h w) c')
        img_gt = rgbs_flattend.view(H, W, 3).permute(2, 0, 1).cpu()  # (3, H, W)
        stack_list = [img_gt]
        for i in range(self.num_models):
            results = ret.rgbs[i]
            results_c = ret_c.rgbs[i]
            mk_t = rearrange(mk[i], 'b h w c -> (b h w) c')
            if batch_idx == 0:
                results[mk_t.sum(dim=-1) <= 0.05] = 0
                img = results.view(H, W, 3).permute(2, 0, 1).cpu()  # (3, H, W)
                img_c = results_c.view(H, W, 3).permute(2, 0, 1).cpu()  # (3, H, W)
                stack_list.append(img)
                stack_list.append(img_c)

        stack = torch.stack(stack_list) # (3, 3, H, W)
        self.logger.experiment.add_images('val/GT_Reconstructed', stack,
                                          self.global_step)

        return log

    def test_step(self, batch, batch_idx):
        ts_w = batch['ts_w']