我有一组图像需要进行推理。我想做的是生成 4 个工作线程,每个工作线程对 GPU 内存有 1/4 的访问权限。即使我将每个工作人员的内存限制为 1/8,我也会遇到 OOM。下面是一些代码:
def process_image(doc_paths):
model = models.Model(my_model)
for path in doc_paths:
...do some work...
model.do_predict()
return 0 or 1
if __name__ == "__main__":
num_processes = 4
process_pool_executer = concurrent.futures.ProcessPoolExecutor(num_processes)
# Split documents into `num_processes` equal parts, so in our case -> 4 equal parts
for i, paths in enumerate(sub_file_paths):
print(f'==== Processing path {i} ====')
f = process_pool_executer.submit(process_image, paths)
futures.append(f)
for future in concurrent.futures.as_completed(futures):
r = future.result()
results.append(r)
在
do_predict()
函数中,我导入tensorflow并尝试了这两个内存分配选项:
# OPTION 1
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
tf.config.experimental.set_virtual_device_configuration(gpus[0], [
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
except RuntimeError as e:
print(e)
# OPTION 2
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
两者都给我 OOM 错误。根据 nvidia-smi,我的 GPU 有 8081 MiB 可用空间,因此为每个进程分配 1024 MB 应该没问题。所以我一定是在以错误的方式做某事或思考某事。有什么想法吗?
虽然您可能没有使用卡提供给您的所有内存,但也许您的进程正在耗尽分配给它们的所有内存。换句话说,也许给他们更少的记忆实际上是相反的方向?