Training vs Inference Optimization Trade-offs Questions

import torch
from torch.utils.data import DataLoader

# 1. Load model and weights
model = MyLSTMModel(...)  # contains nn.LSTM, nn.Linear
model.eval()
baseline = model  # keep baseline copy

# 2. Prepare calibration / eval dataset (representative)
calib_loader = DataLoader(calib_dataset, batch_size=32, shuffle=False)
eval_loader = DataLoader(eval_dataset, batch_size=32, shuffle=False)

# 3. Apply dynamic quantization
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.LSTM, torch.nn.Linear}, dtype=torch.qint8
)

# 4. Evaluation function
def evaluate(m, loader):
    m.eval()
    total, correct = 0, 0
    import time
    t0 = time.time()
    with torch.no_grad():
        for x, y in loader:
            out = m(x)
            pred = out.argmax(dim=-1)
            correct += (pred == y).sum().item()
            total += y.numel()
    latency = (time.time() - t0) / len(loader)
    return correct/total, latency

# 5. Run evals
acc_baseline, lat_base = evaluate(baseline, eval_loader)
acc_q, lat_q = evaluate(quantized_model, eval_loader)

# 6. Size
import io
def model_size(m):
    f = io.BytesIO(); torch.save(m.state_dict(), f); return f.getbuffer().nbytes
size_base, size_q = model_size(baseline), model_size(quantized_model)

print(acc_baseline, acc_q, lat_base, lat_q, size_base, size_q)