Native Keras API Support

On the Ascend platform, you can directly use the native Keras API for training. However, the number of iterations per training loop on the Ascend AI Processor is fixed at 1 in each sess.run() call. To migrate Keras-based network scripts to the Ascend platform for training, pay attention to the following points:

• use_off_line needs to be enabled for training on the Ascend AI Processor. Therefore, you need to create a TensorFlow session and register Keras first. The session should be closed when the training ends.

  import tensorflow as tf
  import tensorflow.python.keras as keras
  from tensorflow.python.keras import backend as K
  from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig
  from npu_bridge.estimator import npu_ops
  
  sess_config = tf.ConfigProto()
  custom_op = sess_config.graph_options.rewrite_options.custom_optimizers.add()
  custom_op.name = "NpuOptimizer"
  custom_op.parameter_map["use_off_line"].b = True
  sess_config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
  sess_config.graph_options.rewrite_options.optimizers.extend(["GradFusionOptimizer"]) # This line is required in the distributed training scenario.
  sess = tf.Session(config=sess_config)
  K.set_session(sess)
  
  # Preprocess the data...
  # Construct a model...
  # Build the model...
  # Train the model...
  
  sess.close()

• If tf.device is used on the original network, delete the related code.

In addition, Ascend supports functions such as automatic mixed precision. For details about how to enable the functions, see the API description.