97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ROCm MIOpen Cache Warmer for MinerU
|
||
|
||
在 AMD RDNA 架构上，MIOpen 遇到新尺寸的卷积运算时需要搜索最优 kernel（冷启动）。
|
||
预热脚本提前跑一遍常用尺寸，将 kernel 缓存到 ~/.cache/miopen/，避免运行时等待。
|
||
|
||
缓存持久化到磁盘，重启不丢失；只有升级 ROCm 后才需要重新跑。
|
||
|
||
用法:
|
||
python cache_warmer.py --device cuda --max_side 960 --step 32
|
||
"""
|
||
|
||
import argparse
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from tqdm import tqdm
|
||
|
||
|
||
def get_args():
|
||
p = argparse.ArgumentParser(description="ROCm MIOpen Cache Warmer")
|
||
p.add_argument("--device", type=str, default="cuda")
|
||
p.add_argument("--max_side", type=int, default=960)
|
||
p.add_argument("--step", type=int, default=32)
|
||
return p.parse_args()
|
||
|
||
|
||
class MockOCRModel(nn.Module):
|
||
"""模拟 MinerU OCR 模型的卷积结构，覆盖 MIOpen 常用 kernel 尺寸。"""
|
||
|
||
def __init__(self, in_channels: int = 3):
|
||
super().__init__()
|
||
self.stem = nn.Conv2d(in_channels, 16, 3, stride=2, padding=1)
|
||
self.dw_3x3 = nn.Conv2d(16, 16, 3, stride=1, padding=1, groups=16)
|
||
self.pw_1 = nn.Conv2d(16, 64, 1)
|
||
self.dw_5x5 = nn.Conv2d(64, 64, 5, stride=2, padding=2, groups=64)
|
||
self.pw_2 = nn.Conv2d(64, 128, 1)
|
||
self.dw_3x3_s2 = nn.Conv2d(128, 128, 3, stride=2, padding=1, groups=128)
|
||
self.pw_3 = nn.Conv2d(128, 256, 1)
|
||
self.out_conv = nn.Conv2d(256, 64, 1)
|
||
self.binarize_conv = nn.Conv2d(64, 1, 3, stride=1, padding=1)
|
||
self.act = nn.ReLU()
|
||
|
||
def forward(self, x):
|
||
x = self.stem(x)
|
||
x = self.act(x)
|
||
x = self.dw_3x3(x)
|
||
x = self.pw_1(x)
|
||
x = self.dw_5x5(x)
|
||
x = self.act(x)
|
||
x = self.pw_2(x)
|
||
x = self.dw_3x3_s2(x)
|
||
x = self.pw_3(x)
|
||
x = self.out_conv(x)
|
||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||
x = self.binarize_conv(x)
|
||
return x
|
||
|
||
|
||
def main():
|
||
args = get_args()
|
||
assert torch.cuda.is_available(), "GPU not available"
|
||
device = torch.device(args.device)
|
||
|
||
print("=" * 50)
|
||
print("ROCm MIOpen Cache Warmer")
|
||
print(f" GPU : {torch.cuda.get_device_name(0)}")
|
||
print(f" ROCm : {torch.version.hip}")
|
||
print(f" Cache : ~/.cache/miopen/")
|
||
print("=" * 50)
|
||
|
||
model = MockOCRModel().to(device).eval()
|
||
sizes = list(range(64, args.max_side + 1, args.step))
|
||
combos = [(h, w) for h in sizes for w in sizes]
|
||
|
||
print(f"Warming {len(combos)} shapes ({len(sizes)}x{len(sizes)} grid, "
|
||
f"step={args.step})...")
|
||
|
||
ok = 0
|
||
with torch.no_grad():
|
||
for h, w in tqdm(combos, desc="Warming"):
|
||
try:
|
||
model(torch.zeros((1, 3, h, w), device=device, dtype=torch.float32))
|
||
ok += 1
|
||
except RuntimeError as e:
|
||
if "out of memory" in str(e):
|
||
torch.cuda.empty_cache()
|
||
# 其他错误跳过，不影响后续
|
||
|
||
print(f"\nDone! {ok}/{len(combos)} shapes cached (~3–4 min)")
|
||
print("Kernels saved to ~/.cache/miopen/")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|