add_activation_loggers

torchao.prototype.quant_logger.add_activation_loggers(model: Module)

Wraps nn.Linear weights with ActivationLoggingTensor so that input activations are logged each time F.linear is called.

The logging behavior is user-configurable by redefining the quant_logger::log_tensor custom op (see enable_log_tensor_save_tensors_to_disk() and enable_log_stats_to_file() for built-in alternatives).

Parameters:

model – The model whose nn.Linear weights will be wrapped.

Example:

import torch
import torch.nn as nn

from torchao.prototype.quant_logger import add_activation_loggers

model = nn.Sequential(
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 512),
)
add_activation_loggers(model)
x = torch.randn(1, 128)
model(x)
# t=act, c=0, fqn='0.weight', op='linear', extra='MKN=1|128|256', max_abs=..., avg=..., std=...
# t=act, c=1, fqn='2.weight', op='linear', extra='MKN=1|256|512', max_abs=..., avg=..., std=...