Enabling Iteration Offload in sess.run Mode

from __future__ import print_function
import input_data
from npu_bridge.npu_init import *
 
mnist = input_data.read_data_sets("/test/", one_hot=True)
 
import tensorflow as tf
 
# Set the model.
# Set the learning rate.
learning_rate = 0.01
# Set the number of training iterations.
training_epochs = 10
# Set the batch size.
batch_size = 100
# Set the number of iterations to display the loss.
display_step = 1
 
x = tf.placeholder(tf.float32, [None, 784])
y = tf.placeholder(tf.float32, [None, 10])
 
# Set the model parameters.
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
 
# Compile the model.
pred = tf.nn.softmax(tf.matmul(x, W) + b)
 
# Define the loss function: cross entropy.
cost = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred), reduction_indices=1))
 
# Update gradients.
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
 
# Initialize all variables.
init = tf.global_variables_initializer()
 
config = tf.ConfigProto(allow_soft_placement=True)
custom_op =  config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name =  "NpuOptimizer"
custom_op.parameter_map["mix_compile_mode"].b = False  # Disable mixed computing (disabled by default).
custom_op.parameter_map["enable_data_pre_proc"].b = True # If the GetNext operator exists on the network, offload it. The GetNext operator offload is a prerequisite for iteration offload.
custom_op.parameter_map["iterations_per_loop"].i = 10  # Used for functional validation. Must be equal to iterations_per_loop set in load_iteration_per_loop_var.
config = npu_config_proto(config_proto=config)
 
# Train the model.
with tf.Session(config=config) as sess:
    sess.run(init)
    # Set the number of iterations per loop to 10 in sess.run mode.
    iteration = util.IterationPerLoop()   # Define the IterationPerLoop class object.
    train_op = iteration.create_iteration_per_loop_var(optimizer)    # optimizer indicates the defined operation for updating gradients.
    tf.train.Supervisor(logdir="/home/xxxx",init_op=init)  # Freeze the graph.
    iteration.load_iteration_per_loop_var(sess, 10)  # Set the number of iterations per loop. sess indicates the created TensorFlow session, and 10 indicates the number of training iterations on the device.

    for epoch in range(training_epochs):
        avg_cost = 0
        total_batch = int(mnist.train.num_examples / batch_size)
 
        for i in range(total_batch):
            batch_xs, batch_ys = mnist.train.next_batch(batch_size)
            _, c = sess.run([train_op, cost], feed_dict={x: batch_xs, y: batch_ys})
 
            avg_cost += c / total_batch

from __future__ import print_function
import input_data
from npu_bridge.npu_init import *
 
mnist = input_data.read_data_sets("/test/", one_hot=True)
 
import tensorflow as tf
 
# Set the model.
# Set the learning rate.
learning_rate = 0.01
# Set the number of training iterations.
training_epochs = 10
# Set the batch size.
batch_size = 100
# Set the number of iterations to display the loss.
display_step = 1
 
x = tf.placeholder(tf.float32, [None, 784])
y = tf.placeholder(tf.float32, [None, 10])
 
# Set the model parameters.
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
 
# Compile the model.
pred = tf.nn.softmax(tf.matmul(x, W) + b)
 
# Define the loss function: cross entropy.
cost = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred), reduction_indices=1))
 
# Update gradients.
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
 
# Initialize all variables.
init = tf.global_variables_initializer()
 
config = tf.ConfigProto(allow_soft_placement=True)
custom_op =  config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name =  "NpuOptimizer"
custom_op.parameter_map["use_off_line"].b = True # Perform training on the Ascend AI Processor.
custom_op.parameter_map["mix_compile_mode"].b = False  # Disable mixed computing (disabled by default).
custom_op.parameter_map["enable_data_pre_proc"].b = True # If the GetNext operator exists on the network, offload it. The GetNext operator offload is a prerequisite for iteration offload.
custom_op.parameter_map["iterations_per_loop"].i = 10 # Determine whether the training iteration is offloaded. Must be equal to iterations_per_loop set in set_iteration_per_loop.
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF
 
# Train the model.
with tf.Session(config=config) as sess:
    sess.run(init)
    # Set the number of iterations per loop to 10 in sess.run mode.
    train_op = util.set_iteration_per_loop(sess, optimizer, 10) # sess indicates the created TensorFlow session, optimizer indicates the defined gradient update operation, and 10 indicates the number of training iterations on the device.
 
    for epoch in range(training_epochs):
        avg_cost = 0
        total_batch = int(mnist.train.num_examples / batch_size)
 
        for i in range(total_batch):
            batch_xs, batch_ys = mnist.train.next_batch(batch_size)
            _, c = sess.run([train_op, cost], feed_dict={x: batch_xs, y: batch_ys})
 
            avg_cost += c / total_batch

from __future__ import print_function
import input_data
from npu_bridge.npu_init import *
 
mnist = input_data.read_data_sets("/test/", one_hot=True)
 
import tensorflow as tf
 
# Set the model.
# Set the learning rate.
learning_rate = 0.01
# Set the number of training iterations.
training_epochs = 10
# Set the batch size.
batch_size = 100
# Set the number of iterations to display the loss.
display_step = 1
 
x = tf.placeholder(tf.float32, [None, 784])
y = tf.placeholder(tf.float32, [None, 10])
 
# Set the model parameters.
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
 
# Compile the model.
pred = tf.nn.softmax(tf.matmul(x, W) + b)
 
# Define the loss function: cross entropy.
cost = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred), reduction_indices=1))
 
# Update gradients.
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
 
# Initialize all variables.
init = tf.global_variables_initializer()
 
config = tf.ConfigProto(allow_soft_placement=True)
custom_op =  config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name =  "NpuOptimizer"
custom_op.parameter_map["use_off_line"].b = True  # Perform training on the Ascend AI Processor.
custom_op.parameter_map["mix_compile_mode"].b = False  # Disable mixed computing (disabled by default).
custom_op.parameter_map["enable_data_pre_proc"].b = True  # If the GetNext operator exists on the network, offload it. The GetNext operator offload is a prerequisite for iteration offload.
custom_op.parameter_map["iterations_per_loop"].i = 10  # Determine whether the training iteration is offloaded. Must be equal to iterations_per_loop set in load_iteration_per_loop_var.
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF 
# Train the model.
with tf.Session(config=config) as sess:
    sess.run(init)
    # Set the number of iterations per loop to 10 in sess.run mode.
    iteration = util.IterationPerLoop() 
    train_op = iteration.create_iteration_per_loop_var(optimizer) # optimizer indicates the defined operation for updating gradients.
    tf.train.Supervisor(logdir="/home/xxxx",init_op=init)  # Freeze the graph.
    iteration.load_iteration_per_loop_var(sess, 10)  # Set the number of iterations per loop. sess indicates the created TensorFlow session, and 10 indicates the number of training iterations on the device.

    for epoch in range(training_epochs):
        avg_cost = 0
        total_batch = int(mnist.train.num_examples / batch_size)
 
        for i in range(total_batch):
            batch_xs, batch_ys = mnist.train.next_batch(batch_size)
            _, c = sess.run([train_op, cost], feed_dict={x: batch_xs, y: batch_ys})
 
            avg_cost += c / total_batch