PyTorch框架提供与算子编译相关的二进制配置参数,可设置模型编译时是否优先在线编译,以此优化模型训练性能。参数配置代码如下,可设置True或False。当前针对Atlas 训练系列产品,若不添加以下代码则默认为True。
torch_npu.npu.set_compile_mode(jit_compile=False)
用户可根据模型是固定shape或动态shape配置该参数以提高性能。
用户在模型训练后,可参见PyTorch Analyse迁移分析工具,使用“dynamic_shape”模式,添加动态分析代码并运行脚本,根据分析结果判断为固定shape或动态shape,然后选择如下配置:
torch_npu.npu.set_compile_mode(jit_compile=False)
if __name__ == '__main__': torch_npu.npu.set_compile_mode(jit_compile=False) main()
mp.spawn(main_worker,...) ... def main_worker(): torch_npu.npu.set_compile_mode(jit_compile=False)
if __name__ == '__main__': torch_npu.npu.set_compile_mode(jit_compile=False) main()
mp.spawn(main_worker,...) ... def main_worker(): torch_npu.npu.set_compile_mode(jit_compile=False)