How to collect outputs from test_step?

implementation help
1

Hello, I have been using Lightning version 1.6 and have recently migrated to 2.0.
During the evaluation of our model, I have collected the outputs from the test_step and utilized them
for post-calculation such as calculating test accuracy, saving results in csv. Please look at the below code. How can I do something similar in 2.0? It seems that test_epoch_end is no longer available and I have to use on_test_epoch_end. However, I don’t see much documentation for my use. Can someone please guide me? thank you in advance!

    def test_step(self, batch, batch_idx):
        feature, y_true = batch["feature"], batch["y_true"]
        y_pred = self.model(feature)
        test_loss = self.loss(y_pred.squeeze(), y_true)
        return {
            "test_loss": test_loss,
            "y_pred": y_pred,
            "y_true": y_true,
        }

    def test_epoch_end(self, outputs):
        test_loss = torch.stack([x["test_loss"] for x in outputs]).mean().cpu()
        y_pred = torch.cat([torch.argmax(x["y_pred"], dim=1) for x in outputs]).cpu()
        y_true = torch.cat([x["y_true"] for x in outputs]).cpu()

        # Save the test results in the output directory
        test_label = pd.read_csv(f"{self.cfg.save_output_path}/test_label.csv")
        test_label["test_loss"] = test_loss.item()
        test_label["stop_epoch"] = self.trainer.early_stopping_callback.stopped_epoch
        test_label["y_pred"] = y_pred.numpy()
        test_label["y_true"] = y_true.numpy()
        test_label["cv_num"] = self.cfg.task.validation_cv_num
        test_label["model_name"] = self.cfg.model.model_name
        test_label["task_name"] = self.cfg.task.task_name
        test_label.to_csv(
            f"{self.cfg.save_output_path}/cv{self.cfg.task.validation_cv_num}_test_label.csv", index=False
        )
        )
2

A self reply. Let me know if there is a better implementation!

def on_test_epoch_start(self) → None:
super().on_test_epoch_start()
self.test_output_list = defaultdict(list)
return

def test_step(self, batch, batch_idx):
    feature, y_true = batch["feature"], batch["y_true"]
    y_pred = self.model(feature)
    test_loss = self.loss(y_pred.squeeze(), y_true)
    self.test_output_list["test_loss"].append(test_loss)
    self.test_output_list["y_pred"].append(y_pred)
    self.test_output_list["y_true"].append(y_true)

def on_test_epoch_end(self):
    outputs = self.test_output_list
    test_loss = torch.stack(outputs["test_loss"]).mean().cpu()
    y_pred = torch.argmax(torch.cat(outputs["y_pred"]), dim=1).cpu()
    y_true = torch.cat(outputs["y_true"]).cpu()