torchao.prototype.quant_logger.log_tensor

torchao.prototype.quant_logger.log_tensor = <CustomOpDef(quant_logger::log_tensor)>

Logs summary statistics for a tensor. This is the default implementation that prints to stdout; users can redefine this custom op to customize logging behavior (e.g. write to file, log custom stats).

Note: do not use this function directly. Instead, you can override it (see code sample below) to modify what gets logged and/or where it gets logged to.

Parameters:

• x – The tensor to log statistics for.

• fqn – Fully qualified name of the module parameter associated with this tensor.

• op – The operation being logged (e.g. "linear"), or "" for parameters.

• tag – A tag categorizing the log entry (e.g. "act" for activations, "param" for parameters).

• extra – Optional extra metadata string to include in the log line.

Example:

import torch
import torch.nn as nn

from torchao.prototype.quant_logger import log_parameter_info


# Override the default log_tensor to print per-channel mean and std
@torch.library.custom_op("quant_logger::log_tensor", mutates_args=("x",))
def log_tensor(
    x: torch.Tensor, fqn: str, op: str, tag: str, extra: str | None = None
) -> None:
    if x.ndim >= 2:
        channel_mean = x.mean(dim=1)
        channel_std = x.std(dim=1, correction=0)
        print(f"{fqn}: channel_mean={channel_mean}, channel_std={channel_std}")
    else:
        print(f"{fqn}: mean={x.mean().item():.4f}, std={x.std().item():.4f}")


model = nn.Sequential(
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 512),
)
log_parameter_info(model)
# 0.weight: channel_mean=tensor([...]), channel_std=tensor([...])
# 0.bias: mean=..., std=...
# 2.weight: channel_mean=tensor([...]), channel_std=tensor([...])
# 2.bias: mean=..., std=...