问题背景

在使用MindSpore运行一个分子动力学模拟的测试程序时：

from mindspore import context
from mindspore.nn import Adam

if __name__ == "__main__":
    import sys
    sys.path.insert(0, '../..')

    from sponge import Sponge, Molecule, ForceField, set_global_units
    from sponge.callback import WriteH5MD, RunInfo
    
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
    set_global_units('nm', 'kj/mol')

    system = Molecule(template='water.spce.yaml')

    system.reduplicate([0.3, 0, 0])
    new_sys = system.copy([0, 0, -0.3])
    system.append(new_sys)

    potential = ForceField(system, parameters='SPCE')
    opt = Adam(system.trainable_params(), 1e-3)
    mini = Sponge(system, potential, opt)

    run_info = RunInfo(5)
    cb_h5md = WriteH5MD(system, 'tutorial_c01.h5md', save_freq=10, write_velocity=True, write_force=True)
    mini.run(10, callbacks=[run_info, cb_h5md])

遇到这样的一个报错：

Traceback (most recent call last):
  File "potential_cell_doc.py", line 11, in <module>
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
  File "/home/dechin/anaconda3/envs/mindsponge/lib/python3.9/site-packages/mindspore/_checkparam.py", line 1313, in wrapper
    return func(*args, **kwargs)
  File "/home/dechin/anaconda3/envs/mindsponge/lib/python3.9/site-packages/mindspore/context.py", line 1493, in set_context
    ctx.setters[key](ctx, value)
  File "/home/dechin/anaconda3/envs/mindsponge/lib/python3.9/site-packages/mindspore/context.py", line 430, in set_device_id
    self.set_param(ms_ctx_param.device_id, device_id)
  File "/home/dechin/anaconda3/envs/mindsponge/lib/python3.9/site-packages/mindspore/context.py", line 175, in set_param
    self._context_handle.set_param(param, value)
TypeError: For 'set_context', the parameter device_id can not be set repeatedly, origin value [0] has been in effect.

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/core/utils/ms_context.cc:477 CheckReadStatus

这里提示的是重复定义device_id的错误。但是实际上我通过vscode的文件检索功能，发现在整个程序引用中，并没有哪里调用到了这个device_id设置的位置。

解决思路

既然不是程序本身设置了device_id，那就是程序内部调用了MindSpore的组件去执行GPU相关的运算，或者是在JIT编译过程会自动指定device_id。那么问题解决的思路就很简单了，直接把设置device_id的这个内容放到文件的最前头，在编译和调用MindSpore进行运算之间就指定好device_id：

from mindspore import context
from mindspore.nn import Adam

if __name__ == "__main__":
    import sys
    sys.path.insert(0, '../..')

    from sponge import Sponge, Molecule, ForceField, set_global_units
    from sponge.callback import WriteH5MD, RunInfo
    
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
    set_global_units('nm', 'kj/mol')

    system = Molecule(template='water.spce.yaml')

    system.reduplicate([0.3, 0, 0])
    new_sys = system.copy([0, 0, -0.3])
    system.append(new_sys)

    potential = ForceField(system, parameters='SPCE')
    opt = Adam(system.trainable_params(), 1e-3)
    mini = Sponge(system, potential, opt)

    run_info = RunInfo(5)
    cb_h5md = WriteH5MD(system, 'tutorial_c01.h5md', save_freq=10, write_velocity=True, write_force=True)
    mini.run(10, callbacks=[run_info, cb_h5md])

重新运行后结果如下：

[MindSPONGE] Started simulation at 2024-03-22 10:08:10
[MindSPONGE] Step: 5, E_pot: 100.10054
[MindSPONGE] Step: 10, E_pot: 89.08189
[MindSPONGE] Finished simulation at 2024-03-22 10:08:13
[MindSPONGE] Simulation time: 2.83 seconds.
--------------------------------------------------------------------------------

问题成果解决。

总结概要

本文主要介绍了一个在使用MindSpore框架进行编程的时候遇到的一个小问题--重复设定运算设备的编号。之所以会出现这个问题，是因为在调用的包里面有对MindSpore的引用，里面还包含了一些基于MindSpore的运算和即时编译，因此如果不给定设备编号的话，MindSpore内部会默认分配一个设备编号。而如果我们在自己的测试案例中又希望指定一个设备编号，那么就要把这个set_context句柄放在引用的最前面。