内存/显存管理的教训
数据加载
今天换到5090 Laptop (24G VRAM) 把GDINO开到400张图的时候即使有batch，VRAM还是炸了，结果2分钟的推理变成了15分钟，之后才发现这个代码完全没有操心内存和显存占用的问题。
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
32
33
34
35
| SAMPLE_SIZE = 9
GDINO_BATCH_SIZE = 4
#采样略
paths = [os.path.join(PROJECT_ROOT, "data", p) for p in sample_images["downloaded_path"]]
# p
aths = ['/Users/yukun/projects/wakareeru/data/img/E233系/1f5c3f_Toyota Vehicle Center.jpg', "/Users/yukun/projects/wakareeru/data/img/EF510形/ca6c48_UetsuhonsenKomachiKoshuFreight.jpg",
# '/Users/yukun/projects/wakareeru/data/img/E231系/c5ea96_Ochanomizu Crossing 2020-03-17.jpg']
images = [load_img_with_orientation(path) for path in paths]
text_labels = [["a train"]] * (len(images)//2) + [["a single locomotive"]] * (len(images) - len(images)//2) # 针对动车组和机车用不同的标签，分别侦测整列车和机车自己。
inputs = processor(
images=images,
text=text_labels,
return_tensors="pt",
padding=True, # batch 内图片尺寸不同时需要 padding
).to(device)
# (width, height) -> (height, width)，每张图片单独从归一化坐标转换回原始尺寸
target_sizes = [img.size[::-1] for img in images]
# 缓存每个 batch 的模型输出；之后只改阈值或重新打印时，不需要再跑模型。
gdino_batch_outputs = []
gdino_batch_input_ids = []
gdino_batch_target_sizes = []
with torch.no_grad():
for start in range(0, len(images), GDINO_BATCH_SIZE):
end = start + GDINO_BATCH_SIZE
inputs_batch = {k: v[start:end] for k, v in inputs.items()}
gdino_batch_outputs.append(model(**inputs_batch))
gdino_batch_input_ids.append(inputs_batch["input_ids"])
gdino_batch_target_sizes.append(target_sizes[start:end])
|
能注意到两个糟糕的点，images = [load_img_with_orientation(path) for path in paths]直接把全量图片object一口气塞进内存，结果一点开就吃走了40G RAM，幸好这台Workstation是64G内存的配置，不过如果真的到了全量阶段也只会变成一坨大的。其实如果用Hugging Face和Torch带的workflow的话，他们的内存管理会帮我解决这个问题，例如可以在Dataloader(batch=32)这里只需通过参数就解决这个问题，但是因为想要快速实验所以还是自己跳了个坑。通过yield的迭代器作为Generator可以解决这个问题，吗？
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
| def iter_batches(db_path, batch_size=32):
con = sqlite3.connect(db_path)
cur = con.cursor()
cur.arraysize = batch_size # 影响 fetchmany 的默认值
cur.execute("""
SELECT file_title, cropped_path, series
FROM images
WHERE embed_status IS NULL
AND crop_status = 'done'
ORDER BY file_title -- 保证幂等，每次顺序一致
""")
while True:
batch = cur.fetchmany(batch_size) # 每次只从 SQLite 拉 batch_size 行
if not batch:
break
yield batch # 挂起，把控制权交给调用方
con.close()
|
现在的问题变成了，如果一次只加载一个Batch，那GPU岂不是很闲，还要等CPU加载？，其实说对了，所以最后还是需要PyTorch的DataLoader来帮我们解决预加载和缓存的问题，自己的随手轮子哪有别人的优化好呢。不过需要一个class实现几个接口，可以在后面的正式pipeline中完成。
1
2
3
4
5
6
7
| loader = DataLoader(
dataset,
batch_size=32,
num_workers=4, # 4 个子进程并行加载图像
prefetch_factor=2, # 每个 worker 超前准备 2 个 batch
pin_memory=True, # 锁页内存，加速 CPU→GPU 传输
)
|
这样就只用改参数就万事大吉了。
推理结果detach
好久没用pytorch，都忘了tensor是要送进device的。下面的代码可以及时把结果推到RAM里去。
1
2
3
4
5
6
7
| with torch.no_grad():
outputs = model(pixel_values) # 在 GPU
embeddings = outputs.last_hidden_state[:, 0]
embeddings = embeddings.detach().cpu().numpy() # 立刻卸载
# outputs 对象此时没有引用，下次 GC 或 cuda 分配时会自动释放
# 但主动 del 更保险：
del outputs, pixel_values
|
之前打开结果打印出来一看结果全是tensor([...], device=cuda)，啊哈哈。
