sess.run迁移

sess.run简介

sess.run API属于TensorFlow的低阶API,相对于Estimator来讲,灵活性较高,但模型的实现较为复杂。

使用sess.run API进行训练脚本开发的流程为:

  1. 数据预处理。
  2. 模型搭建/计算Loss/梯度更新。
  3. 创建session并初始化资源。
  4. 执行训练。

下面介绍如何迁移sess.run训练脚本,以便在昇腾AI处理器上进行训练。

头文件增加

对于以下步骤中涉及修改的python文件,新增以下头文件引用,用于导入NPU相关库。

from npu_bridge.npu_init import *

数据预处理

一般情况下,此部分代码无需改造。如下情况需要进行适配修改:

由于当前仅支持固定shape下的训练,也就是在进行图编译时shape的值必须是已知的。当原始网络脚本中使用dataset.batch(batch_size)返回动态形状时,由于数据流中剩余的样本数可能小于batch大小,因此,在昇腾AI处理器上进行训练时,请将drop_remainder设置为True:
  dataset = dataset.batch(batch_size, drop_remainder=True)

这可能会丢弃文件中的最后几个样本,以确保每个批量都具有静态形状 (batch_size)。但需要注意的是:推理时,当最后一次迭代的推理数据量小于batch size时,需要补齐空白数据到batch size,因为有些脚本最后会加个断言,验证结果的数量要和验证数据的数量一致,此种情况会导致训练失败。

  assert num_written_lines == num_actual_predict_examples

模型搭建/计算Loss/梯度更新

一般情况下,此部分代码无需改造。如下情况需要进行适配修改:

创建session并初始化资源

昇腾AI处理器上通过sess.run模式执行训练脚本时,相关配置说明:

TensorFlow原始代码:

#构造迭代器
iterator=Iterator.from_structure(train_dataset.output_types,train_dataset.output_shapes)

#取batch数据
next_batch=iterator.get_next()

#迭代器初始化
training_init_op=iterator.make_initializer(train_dataset)
 
#变量初始化
init=tf.global_variables_initializer()
sess=tf.Session()
sess.run(init)
 
#Get the number of training/validation steps per epoch
train_batches_per_epoch=int(np.floor(train_size/batch_size))

迁移后的代码:

# 构造迭代器
iterator=Iterator.from_structure(train_dataset.output_types,train_dataset.output_shapes)

# 取batch数据
next_batch=iterator.get_next()

# 迭代器初始化
training_init_op=iterator.make_initializer(train_dataset)
 
# 变量初始化
init=tf.global_variables_initializer()

# 创建session,如果原始网络中使用了tf.device相关代码,则需要增加session配置“allow_soft_placement=True”,允许TF自动分配设备。
config = tf.ConfigProto(allow_soft_placement=True)
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF  # 必须显式关闭
config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF  # 必须显式关闭
sess = tf.Session(config=config)
sess.run(init)
 
#Get the number of training/validation steps per epoch
train_batches_per_epoch=int(np.floor(train_size/batch_size))

tf.Session原生功能在Ascend平台上全部支持。

另外,Ascend平台还支持自动混合精度等功能,如果用户需要进行相关使能,可以参考对应接口说明。

执行训练

此部分代码无需改造,例如:

#开始循环迭代
for epoch in range(num_epochs):
  ##Initialize iterator with the training dataset
  sess.run(training_init_op)
  for step in range(train_batches_per_epoch):  
    #get next batch of data
    img_batch,label_batch=sess.run(next_batch)
    #run the training op
    _,train_loss = sess.run([train_op, loss],feed_dict={x:img_batch,y_:label_batch,is_training:True})

但是,如果用户训练脚本中没有使用with创建session,比如将session对象作为自己定义的一个类成员,那么需要在迁移后的脚本中显示调用sess.close()。

sess = tf.Session(config=config)
sess.run(...)
sess.close()

这是因为,GEOP的析构函数在tf.session的close方法中会被调用到,如果是with创建的session,with会调用session的__exit__方法,里面会自动调用close:

with tf.Session(config=config) as sess:
    sess.run(...)

如果是其他情况,比如是把session对象作为自己定义的一个类成员,那么退出之前需要显示调用sess.close(),这样才可以保证退出的正常。