PyTorch模型部署前的完整测试流程
在将PyTorch模型投入生产环境前，必须进行系统性的性能测试以确保其满足实际部署要求。本文将通过具体代码示例展示完整的测试流程。
1. 基准性能测试
首先使用torch.utils.benchmark测试模型推理速度：
import torch
import torch.nn as nn
from torch.utils.benchmark import Timer
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 64, 3)
self.fc = nn.Linear(64, 10)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1)
return self.fc(x)
model = SimpleModel()
model.eval()
timer = Timer(
stmt='model(x)',
globals={'model': model, 'x': torch.randn(1, 3, 224, 224)}
)
print(timer.timeit(100))
2. 内存占用测试
使用torch.cuda.memory_allocated()监控显存使用：
model = model.cuda()
x = torch.randn(32, 3, 224, 224).cuda()
with torch.no_grad():
torch.cuda.reset_peak_memory_stats()
_ = model(x)
peak_memory = torch.cuda.max_memory_allocated()
print(f'Peak memory: {peak_memory / 1024**2:.2f} MB')
3. 多batch size测试
对比不同batch size下的性能表现：
for batch_size in [1, 8, 16, 32]:
x = torch.randn(batch_size, 3, 224, 224)
time_taken = timer.timeit(50).mean * 1000 # ms
print(f'BS {batch_size}: {time_taken:.2f}ms')
4. 混合精度测试
使用torch.cuda.amp进行FP16推理测试：
with torch.cuda.amp.autocast():
output = model(x)
完整测试应包含上述所有环节，确保模型在各种场景下稳定运行。
讨论