#!/bin/bash
export LD_LIBRARY_PATH=/usr/local/lib/:/usr/local/HiAI/runtime/lib64
export PATH=/usr/local/HiAI/runtime/ccec_compiler/bin:$PATH
export CUSTOM_OP_LIB_PATH=/usr/local/HiAI/runtime/ops/framework/built-in/tensorflow
export DDK_VERSION_PATH=/usr/local/HiAI/runtime/ddk_info
export WHICH_OP=GEOP
export NEW_GE_FE_ID=1
export GE_AICPU_FLAG=1
export OPTION_EXEC_EXTERN_PLUGIN_PATH=/usr/local/HiAI/runtime/lib64/plugin/opskernel/libfe.so:/usr/local/HiAI/runtime/lib64/plugin/opskernel/libaicpu_plugin.so:/usr/local/HiAI/runtime/lib64/plugin/opskernel/libge_local_engine.so:/usr/local/H
iAI/runtime/lib64/plugin/opskernel/librts_engine.so:/usr/local/HiAI/runtime/lib64/libhccl.so
 
export OP_PROTOLIB_PATH=/usr/local/HiAI/runtime/ops/built-in/
 
export DEVICE_ID=2
export PRINT_MODEL=1
#export DUMP_GE_GRAPH=2
 
#export DISABLE_REUSE_MEMORY=1
#export DUMP_OP=1
#export SLOG_PRINT_TO_STDOUT=1
 
export RANK_ID=0
export RANK_SIZE=1
export JOB_ID=10087
export OPTION_PROTO_LIB_PATH=/usr/local/HiAI/runtime/ops/op_proto/built-in/
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/fwkacllib/lib64/:/usr/local/Ascend/driver/lib64/common/:/usr/local/Ascend/driver/lib64/driver/:/usr/local/Ascend/add-ons/
export PYTHONPATH=$PYTHONPATH:/usr/local/Ascend/opp/op_impl/built-in/ai_core/tbe
export PATH=$PATH:/usr/local/Ascend/fwkacllib/ccec_compiler/bin
export ASCEND_HOME=/usr/local/Ascend
export ASCEND_OPP_PATH=/usr/local/Ascend/opp
export SOC_VERSION=Ascend910
 
rm -f *.pbtxt
rm -f *.txt
rm -r /var/log/npu/slog/*.log
rm -rf train_url/*
 
 
python3 mnist_train.py

import os
import numpy as np
import tensorflow as tf
import time
from tensorflow.examples.tutorials.mnist import input_data
 
import mnist_inference
 
from npu_bridge.estimator import npu_ops #导入NPU算子库
from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig #重写tensorFlow里的配置，针对NPU的配置
 
 
batch_size = 100
learning_rate = 0.1
training_step = 10000
 
model_save_path = "./model/"
model_name = "model.ckpt"
 
def train(mnist):
    x = tf.placeholder(tf.float32, [batch_size, mnist_inference.image_size, mnist_inference.image_size, mnist_inference.num_channels], name = ‘x-input‘)
    y_ = tf.placeholder(tf.float32, [batch_size, mnist_inference.num_labels], name = "y-input")
 
    regularizer = tf.contrib.layers.l2_regularizer(0.001)
    y = mnist_inference.inference(x, train = True, regularizer = regularizer) #推理过程
    global_step = tf.Variable(0, trainable=False)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits = y, labels = tf.argmax(y_, 1)) #损失函数
    cross_entropy_mean = tf.reduce_mean(cross_entropy)
    loss = cross_entropy_mean + tf.add_n(tf.get_collection("loss"))
 
    train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss, global_step = global_step) #优化器调用
 
    saver = tf.train.Saver() #启动训练
 
#以下代码是NPU所必须的代码，开始配置参数
    config = tf.ConfigProto(
        allow_soft_placement = True,
        log_device_placement = False)
    custom_op =  config.graph_options.rewrite_options.custom_optimizers.add()
    custom_op.name =  "NpuOptimizer"
    custom_op.parameter_map["use_off_line"].b = True
    #custom_op.parameter_map["profiling_mode"].b = True
    #custom_op.parameter_map["profiling_options"].s = tf.compat.as_bytes("task_trace:training_trace")
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
#配置参数结束
 
    writer = tf.summary.FileWriter("./log_dir", tf.get_default_graph())
    writer.close()
 
 
#参数初始化
    with tf.Session(config = config) as sess:
        tf.global_variables_initializer().run()
 
        start_time = time.time()
 
        for i in range(training_step):
xs, ys = mnist.train.next_batch(batch_size)
            reshaped_xs = np.reshape(xs, (batch_size, mnist_inference.image_size, mnist_inference.image_size, mnist_inference.num_channels))
            _, loss_value, step = sess.run([train_step, loss, global_step], feed_dict={x:reshaped_xs, y_:ys})
 
            #每训练10个epoch打印损失函数输出日志
            if i % 10 == 0:
                print("****************************++++++++++++++++++++++++++++++++*************************************\n" * 10)
                print("After %d training steps, loss on training batch is %g, total time in this 1000 steps is %s." % (step, loss_value, time.time() - start_time))
                #saver.save(sess, os.path.join(model_save_path, model_name), global_step = global_step)
                print("****************************++++++++++++++++++++++++++++++++*************************************\n" * 10)
                start_time = time.time()
def main():
    mnist = input_data.read_data_sets(‘MNIST_DATA/‘, one_hot= True)
    train(mnist)
 
if __name__ == "__main__":
    main()