MindSpore 设置昇腾Ascend 910显存的默认大小,单卡可执行多任务处理

默认显存占用

最近在使用昇腾Ascend 910和MindSpore进行训练和推理时,遇到了一个非常有趣的现象,就是无论是进行模型的训练还是推理,都会占用大约30GB的显存。即使只是进行一个简单的Tensor加减法的运算,也需要消耗30GB的显存,这看起来很难受。只做一个简单的Tensor逻辑运算,却占用了30GB的显存。最恶心的是,我花了大价钱从xxx地方购买到的国产化AI加速卡,难道只能同时执行一个任务?(小模型)

import mindspore as ms
x = ms.Tensor([10,20])
y = ms.Tensor([40,50])
b = x +y

显存占用情况:
显存占用

解决方法

参考:https://www.mindspore.cn/docs/zh-CN/r2.2/index.html
mindspore的context有个有个参数可以配置显存的默认占用情况:
MindSpore 设置昇腾Ascend 910显存的默认大小,单卡可执行多任务处理_第1张图片
max_device_memory, 用于设置显存大小,它的默认值是 30GB
由于30GB的默认值,导致了上面的问题
我们只需要配置这个参数就可以限制当前任务的显存占用:

import mindspore as ms
ms.set_context(max_device_memory="1GB")
x = ms.Tensor([10,20])
y = ms.Tensor([40,50])
b = x +y

结果:
在这里插入图片描述
通过对这个配置进行配置,可以灵活设置任务所需显存的大小,并且可以根据显存消耗情况进行多任务的并行运行。

你可能感兴趣的:(人工智能,昇腾AI,人工智能)