一、训练阶段内存估算

1. 计算模型参数量

• 公式：总参数量 = ∑(各层参数数)
• 工具：

PyTorch：

import torch
model = ...  # 加载模型
print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")

TensorFlow/Keras

model.summary()

2. 估算模型参数内存

• 参数存储格式：
  • INT8（1字节/参数）
  • FP16（2字节/参数）
  • FP32（4字节/参数）
• 公式

参数内存 = 总参数量 × 每参数字节数
例如：1亿参数的FP32模型 ≈ 400MB（1e8 × 4B）

3. 优化器状态内存

• 常见优化器内存开销：
  • SGD：无额外状态
  • Adam：每个参数需存储动量（m）和方差（v），即 2 × 参数量 × 每参数字节数
• 公式：

  优化器内存 = 参数内存 × (1 + 优化器系数)

例如：Adam优化器 ≈ 3 × 参数内存（含m、v、参数本身）

4. 梯度内存

• 梯度与参数量相同，内存开销 = 参数内存

5. 激活值内存

• 依赖因素：
  • 模型结构（如Transformer层数、CNN通道数）
  • 批量大小（batch size）
• 估算方法：
  • 使用框架内置工具（如PyTorch Profiler）：

    from torch.profiler import profile, record_function, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof: with record_function("model_inference"): output = model(input_tensor) print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))

6. 总训练内存估算

总内存 ≈ 参数内存 + 优化器内存 + 梯度内存 + 激活值内存

二、部署阶段内存估算流程

1. 模型参数内存

• 同训练阶段，但通常部署时使用INT8/FP16量化，内存可减少50%~75%。

2. 激活值内存

• 仅需前向传播的激活值，比训练时少（无反向传播的中间结果）。
• 估算方法：
  • 使用ONNX Runtime或TensorRT的内存分析工具：

    trtexec --onnx=model.onnx --printLayerInfo

3. 推理框架开销

• 框架自身占用内存（如TensorRT引擎加载、ONNX Runtime会话缓存）。