设置NPU上的循环次数
在确定该是否涉及该适配点前,需要您了解脚本是否采用了循环下沉的编码方式,实际上,阅读官方的脚本发现已经开放了开关供用户选择是否循环下沉。
从official/vision/image_classification/resnet/common.py可以看到官方脚本提供了两个入参:
- steps_per_loop传入training loop的大小,可以从注释中看出,在循环中间,只有训练的动作,不会执行任何callback之类的附加操作。
- use_tf_while_loop则决定是否循环下沉,默认为True,即trainning loop默认都会以While算子的形式执行。
flags.DEFINE_integer( name='steps_per_loop', default=None, help='Number of steps per training loop. Only training step happens ' 'inside the loop. Callbacks will not be called inside. Will be capped at ' 'steps per epoch.') flags.DEFINE_boolean( name='use_tf_while_loop', default=True, help='Whether to build a tf.while_loop inside the training loop on the ' 'host. Setting it to True is critical to have peak performance on ' 'TPU.')
所以,按照默认值,我们需要设置NPU_LOOP_SIZE环境变量的值与steps_per_loop一致。此环境变量的设置说明可参见启动单卡训练。
父主题: 手工迁移与训练