1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
| # 进入要压力测试gpu 容器
kubect -n aijob-system exec -it dddd-0 -- sh
# 编写测试脚本
cat << EOF > demo.py
import os
import numpy as np
import torch
from torchvision.models import resnet18
import time
if __name__ == '__main__':
model = resnet18(pretrained=False)
device = torch.device('cuda')
model.eval()
model.to(device)
dump_input = torch.ones(1,3,224,224).to(device)
# Warn-up
for _ in range(500000):
start = time.time()
outputs = model(dump_input)
torch.cuda.synchronize()
end = time.time()
print('Time:{}ms'.format((end-start)*1000))
with torch.autograd.profiler.profile(enabled=True, use_cuda=True, record_shapes=False, profile_memory=False) as prof:
outputs = model(dump_input)
print(prof.table())
EOF
# 执行python 文件
python demo.py
|