sess.run迁移

sess.run简介

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

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

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

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

头文件增加

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

1

from npu_bridge.npu_init import *

引入上述头文件后，训练脚本默认在昇腾AI处理器执行。

数据预处理

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

1

  dataset = dataset.batch(batch_size, drop_remainder=True)

1

 assert num_written_lines == num_actual_predict_examples

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

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

• 对于原始网络中的dropout，建议替换为CANN对应的API实现，以获得更优性能，但需关注对网络精度的影响。
  • 如果存在tf.nn.dropout，建议修改为：

    1	layers = npu_ops.dropout()

  • 如果存在tf.layers.dropout/tf.layers.Dropout/tf.keras.layers.Dropout/tf.keras.layers.SpatialDropout1D/tf.keras.layers.SpatialDropout2D/tf.keras.layers.SpatialDropout3D，建议增加头文件引用：

    1	from npu_bridge.estimator.npu import npu_convert_dropout

• 对于原始网络中的gelu，建议替换为CANN对应的API实现，以获得更优性能。
  TensorFlow原始代码：

  1 2 3 4 5

  def gelu(x): cdf = 0.5 * (1.0 + tf.tanh( (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))) return x*cdf layers = gelu()

  迁移后的代码：

  1	layers = npu_unary_ops.gelu(x)

创建session并初始化资源

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

• 以下配置默认关闭，请勿开启：

  rewrite_options.disable_model_pruning

• 以下配置默认开启，请勿关闭：
  • rewrite_options.function_optimization
  • rewrite_options.constant_folding
  • rewrite_options.shape_optimization
  • rewrite_options.arithmetic_optimization
  • rewrite_options.loop_optimization
  • rewrite_options.dependency_optimization
  • rewrite_options.layout_optimizer
• 以下配置默认开启，必须显式关闭：
  • rewrite_options.remapping
  • rewrite_options.memory_optimization
• 如果原始网络中使用了tf.device相关代码，需要增加session配置“allow_soft_placement=True”，允许TensorFlow自动分配设备。

TensorFlow原始代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

#构造迭代器
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))

迁移后的代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

#构造迭代器
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配置“allow_soft_placement=True”，允许TensorFlow自动分配设备。
config = tf.ConfigProto(allow_soft_placement=True)
# 添加名字为“NpuOptimizer”的NPU优化器，网络编译时，NPU只会遍历“NpuOptimizer”下的session配置。
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
# 必须显式关闭TensorFlow的remapping、memory_optimization功能，避免与NPU中的功能冲突。
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))

执行训练

此部分代码无需改造，例如：

1
2
3
4
5
6
7
8
9

#开始循环迭代
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()。

1
2
3

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

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

1
2

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

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