Skip to content

table

This module provides the TableWriter class, which saves predictions in a table format, such as CSV.

TableWriter

Bases: BaseWriter

Writer for saving predictions in a table format, such as CSV.

Parameters:

Name Type Description Default
path str | Path

CSV filepath.

 required
writer str | Callable

Writer function or name of a registered writer.

 required
Source code in lighter/callbacks/writer/table.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class TableWriter(BaseWriter):
    """
    Writer for saving predictions in a table format, such as CSV.

    Args:
        path: CSV filepath.
        writer: Writer function or name of a registered writer.
    """

    def __init__(self, path: str | Path, writer: str | Callable) -> None:
        super().__init__(path, writer)
        self.csv_records = []

    @property
    def writers(self) -> dict[str, Callable]:
        return {
            "tensor": lambda tensor: tensor.item() if tensor.numel() == 1 else tensor.tolist(),
        }

    def write(self, tensor: Any, identifier: int | str) -> None:
        """
        Writes the tensor as a table record using the specified writer.

        Args:
            tensor: The tensor to record. Should not have a batch dimension.
            identifier: Identifier for the record.
        """
        self.csv_records.append({"identifier": identifier, "pred": self.writer(tensor)})

    def on_predict_epoch_end(self, trainer: Trainer, pl_module: System) -> None:
        """
        Called at the end of the prediction epoch to save predictions to a CSV file.

        Args:
            trainer: The trainer instance.
            pl_module: The System instance.
        """
        # If in distributed data parallel mode, gather records from all processes to rank 0.
        if trainer.world_size > 1:
            gather_csv_records = [None] * trainer.world_size if trainer.is_global_zero else None
            torch.distributed.gather_object(self.csv_records, gather_csv_records, dst=0)
            if trainer.is_global_zero:
                self.csv_records = list(itertools.chain(*gather_csv_records))

        # Save the records to a CSV file
        if trainer.is_global_zero:
            df = pd.DataFrame(self.csv_records)
            try:
                df = df.sort_values("identifier")
            except TypeError:
                pass
            df = df.set_index("identifier")
            df.to_csv(self.path)

        # Clear the records after saving
        self.csv_records = []

on_predict_epoch_end(trainer, pl_module)

Called at the end of the prediction epoch to save predictions to a CSV file.

Parameters:

Name Type Description Default
trainer Trainer

The trainer instance.

 required
pl_module System

The System instance.

 required
Source code in lighter/callbacks/writer/table.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def on_predict_epoch_end(self, trainer: Trainer, pl_module: System) -> None:
    """
    Called at the end of the prediction epoch to save predictions to a CSV file.

    Args:
        trainer: The trainer instance.
        pl_module: The System instance.
    """
    # If in distributed data parallel mode, gather records from all processes to rank 0.
    if trainer.world_size > 1:
        gather_csv_records = [None] * trainer.world_size if trainer.is_global_zero else None
        torch.distributed.gather_object(self.csv_records, gather_csv_records, dst=0)
        if trainer.is_global_zero:
            self.csv_records = list(itertools.chain(*gather_csv_records))

    # Save the records to a CSV file
    if trainer.is_global_zero:
        df = pd.DataFrame(self.csv_records)
        try:
            df = df.sort_values("identifier")
        except TypeError:
            pass
        df = df.set_index("identifier")
        df.to_csv(self.path)

    # Clear the records after saving
    self.csv_records = []

write(tensor, identifier)

Writes the tensor as a table record using the specified writer.

Parameters:

Name Type Description Default
tensor Any

The tensor to record. Should not have a batch dimension.

 required
identifier int | str

Identifier for the record.

 required
Source code in lighter/callbacks/writer/table.py 
37
38
39
40
41
42
43
44
45
def write(self, tensor: Any, identifier: int | str) -> None:
    """
    Writes the tensor as a table record using the specified writer.

    Args:
        tensor: The tensor to record. Should not have a batch dimension.
        identifier: Identifier for the record.
    """
    self.csv_records.append({"identifier": identifier, "pred": self.writer(tensor)})