[add]上传训练benchmark by z00560161

2020-10-19 20:22:23 +08:00
parent 22b83024f5
commit 82522e2f61
1225 changed files with 345421 additions and 0 deletions
@@ -0,0 +1,22 @@
+import tensorflow as tf
+import os,sys
+
+
+class CreateSession():
+    def __init__(self):
+        self.estimator_config = tf.ConfigProto(
+            inter_op_parallelism_threads=10,
+            intra_op_parallelism_threads=10,
+            allow_soft_placement=True)
+
+        self.estimator_config.gpu_options.allow_growth = True
+
+        self.set_env()
+
+    def set_env(self):
+        gpu_thread_count = 2
+        os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private'
+        os.environ['TF_GPU_THREAD_COUNT'] = str(gpu_thread_count)
+        os.environ['TF_USE_CUDNN_BATCHNORM_SPATIAL_PERSISTENT'] = '1'
+        os.environ['TF_ENABLE_WINOGRAD_NONFUSED'] = '1'
+
@@ -0,0 +1,133 @@
+import numpy as np
+import preprocessing
+import tensorflow as tf
+from tensorflow.python.util import nest
+import os,sys
+import numpy as np 
+
+
+class DataLoader:
+
+    def __init__(self, args):
+        self.args = args   
+
+        filename_pattern = os.path.join(args.data_dir, '%s-*')
+        filenames_train = sorted(tf.gfile.Glob(filename_pattern % 'train'))
+        self.num_training_samples = get_num_records(filenames_train)
+        self.args.num_training_samples = self.num_training_samples
+
+        filename_pattern = os.path.join(args.data_dir, '%s-*')
+        filenames_val = sorted(tf.gfile.Glob(filename_pattern % 'validation'))
+        self.num_evaluating_samples = get_num_records(filenames_val)
+        self.args.num_evaluating_samples = self.num_evaluating_samples
+        
+        print( 'total num_training_sampels: %d' %  self.num_training_samples )
+        print( 'total num_evaluating_sampels: %d' %  self.num_evaluating_samples )
+        
+        self.training_samples_per_rank = self.num_training_samples
+        
+    def get_train_input_fn(self):
+        take_count = self.training_samples_per_rank
+
+        return make_dataset(self.args, take_count, self.args.batch_size, training=True)
+
+    def get_eval_input_fn(self):
+        take_count = self.num_evaluating_samples
+
+        return make_dataset(self.args, take_count, self.args.batch_size, training=False)
+
+
+def get_num_records(filenames):
+    def count_records(tf_record_filename):
+        count = 0
+        for _ in tf.python_io.tf_record_iterator(tf_record_filename):
+            count += 1
+        return count
+
+    nfile = len(filenames)
+    return (count_records(filenames[0]) * (nfile - 1) +
+            count_records(filenames[-1]))
+
+
+def _parse_example_proto(example_serialized):
+    feature_map = {
+        'image/encoded': tf.FixedLenFeature([], dtype=tf.string,
+                                            default_value=''),
+        'image/class/label': tf.FixedLenFeature([], dtype=tf.int64, default_value=-1),
+        'image/class/text': tf.FixedLenFeature([], dtype=tf.string,
+                                               default_value=''),
+    }
+    sparse_float32 = tf.VarLenFeature(dtype=tf.float32)
+    # Sparse features in Example proto.
+    feature_map.update(
+        {k: sparse_float32 for k in ['image/object/bbox/xmin',
+                                     'image/object/bbox/ymin',
+                                     'image/object/bbox/xmax',
+                                     'image/object/bbox/ymax']})
+
+    features = tf.parse_single_example(example_serialized, feature_map)
+    label = tf.cast(features['image/class/label'], dtype=tf.int32)
+
+    xmin = tf.expand_dims(features['image/object/bbox/xmin'].values, 0)
+    ymin = tf.expand_dims(features['image/object/bbox/ymin'].values, 0)
+    xmax = tf.expand_dims(features['image/object/bbox/xmax'].values, 0)
+    ymax = tf.expand_dims(features['image/object/bbox/ymax'].values, 0)
+
+    # Note that we impose an ordering of (y, x) just to make life difficult.
+    bbox = tf.concat([ymin, xmin, ymax, xmax], 0)
+
+    # Force the variable number of bounding boxes into the shape
+    # [1, num_boxes, coords].
+    bbox = tf.expand_dims(bbox, 0)
+    bbox = tf.transpose(bbox, [0, 2, 1])
+
+    return features['image/encoded'], label, bbox
+
+
+# since the preprocessing is done here, we add args file
+def parse_record(raw_record, is_training):
+    image_buffer, label, bbox = _parse_example_proto(raw_record)
+
+    image = preprocessing.parse_and_preprocess_image_record(image_buffer, bbox, training=is_training)
+
+    # label-1 for VGG16
+    return image, label-1
+
+
+def make_dataset(args, take_count, batch_size,
+                 training=False, shard=False):
+
+    shuffle_buffer_size = 10000
+    num_readers = 10
+
+    rank_size = int(os.getenv('RANK_SIZE'))
+    rank_id = int(os.getenv('DEVICE_INDEX'))
+
+    if training:
+        filename_pattern = os.path.join(args.data_dir, '%s-*')
+        filenames = sorted(tf.gfile.Glob(filename_pattern % 'train'))
+    else:
+        filename_pattern = os.path.join(args.data_dir, '%s-*')
+        filenames = sorted(tf.gfile.Glob(filename_pattern % 'validation'))
+
+    ds = tf.data.Dataset.from_tensor_slices(filenames)
+
+    if not training:
+        ds = ds.take(take_count)
+
+    if training:
+        ds = ds.shuffle(1000, seed=7*(1+rank_id))
+
+    ds = ds.interleave(tf.data.TFRecordDataset, cycle_length=num_readers, block_length=1)
+    counter = tf.data.Dataset.range(sys.maxsize)
+    ds = tf.data.Dataset.zip((ds, counter))
+
+    if training:
+        ds = ds.apply(tf.data.experimental.shuffle_and_repeat(shuffle_buffer_size, seed=5*(1+rank_id)))
+
+    ds = ds.map(lambda image, counter: parse_record(image, training), num_parallel_calls=14)
+
+    ds = ds.batch(batch_size, drop_remainder=True)
+    return ds
+
+
@@ -0,0 +1,158 @@
+import tensorflow as tf
+from tensorflow.contrib.layers import batch_norm, flatten
+from tensorflow.contrib.framework import arg_scope
+import numpy as np
+
+class_num = 1000
+nb_blocks = 4
+nb_blocks_layers = (6, 12, 24, 16)
+bn_size = 4
+growth_rate = 32
+init_layers = 64
+
+
+'''
+denseNet：121，169，201，264
+return _densenet('densenet121', 32, (6, 12, 24, 16), 64, pretrained, progress,
+                     **kwargs)
+return _densenet('densenet161', 48, (6, 12, 36, 24), 96, pretrained, progress,
+                     **kwargs)
+return _densenet('densenet169', 32, (6, 12, 32, 32), 64, pretrained, progress,
+                     **kwargs)
+return _densenet('densenet201', 32, (6, 12, 48, 32), 64, pretrained, progress,
+                     **kwargs)                     
+'''
+
+
+
+def conv_layer(input, filter, kernel, stride=1, layer_name="conv"):
+    with tf.name_scope(layer_name):
+        network = tf.layers.conv2d(inputs=input, filters=filter, kernel_size=kernel, strides=stride, padding='SAME', use_bias=False, kernel_initializer=tf.initializers.variance_scaling(scale=5.0, mode='fan_out')) # scale=5.0, mode='fan_out'
+        return network
+
+def Global_Average_Pooling(x, stride=1):
+    
+    width = np.shape(x)[1]
+    height = np.shape(x)[2]
+    pool_size = [width, height]
+    return tf.layers.average_pooling2d(inputs=x, pool_size=pool_size, strides=stride) # The stride value does not matter
+    #It is global average pooling without tflearn
+    
+
+    #return global_avg_pool(x, name='Global_avg_pooling')
+    # But maybe you need to install h5py and curses or not
+
+
+def Batch_Normalization(x, training, scope):
+    with arg_scope([batch_norm],
+                   scope=scope,
+                   updates_collections=None,
+                   decay=0.9,
+                   center=True,
+                   scale=True,
+                   zero_debias_moving_mean=True) :
+        training = tf.cast(training, tf.bool)
+        return tf.cond(training,
+                       lambda : batch_norm(inputs=x, is_training=training, reuse=None),
+                       lambda : batch_norm(inputs=x, is_training=training, reuse=True))
+
+def Drop_out(x, rate, training) :
+    return tf.layers.dropout(inputs=x, rate=rate, training=training)
+
+def Relu(x):
+    return tf.nn.relu(x)
+
+def Average_pooling(x, pool_size=[2,2], stride=2, padding='VALID'):
+    return tf.layers.average_pooling2d(inputs=x, pool_size=pool_size, strides=stride, padding=padding)
+
+
+def Max_Pooling(x, pool_size=[3,3], stride=2, padding='VALID'):
+    return tf.layers.max_pooling2d(inputs=x, pool_size=pool_size, strides=stride, padding=padding)
+
+def Concatenation(layers):
+    return tf.concat(layers, axis=3)
+
+def Linear(x):
+    return tf.layers.dense(inputs=x, units=class_num, name='linear')
+
+
+def bottleneck_layer(x, is_training, scope):
+    # print(x)
+    with tf.name_scope(scope):
+        x = Batch_Normalization(x, training=is_training, scope=scope+'_batch1')
+        x = Relu(x)
+        x = conv_layer(x, filter= growth_rate*bn_size, kernel=[1,1], layer_name=scope+'_conv1')
+        #x = Drop_out(x, rate=dropout_rate, training=is_training)
+        #x = Drop_out(x, rate=dropout_rate, training=is_training)
+
+        x = Batch_Normalization(x, training=is_training, scope=scope+'_batch2')
+        x = Relu(x)
+        x = conv_layer(x, filter= growth_rate, kernel=[3,3], layer_name=scope+'_conv2')
+        #x = Drop_out(x, rate=dropout_rate, training=self.training)
+
+        # print(x)
+
+        return x
+
+def transition_layer(x, is_training, scope):
+    with tf.name_scope(scope):
+        x = Batch_Normalization(x, training=is_training, scope=scope+'_batch1')
+        x = Relu(x)
+        # x = conv_layer(x, filter=self.filters, kernel=[1,1], layer_name=scope+'_conv1')
+        
+        # https://github.com/taki0112/Densenet-Tensorflow/issues/10
+        
+        in_channel = int(x.shape[-1])
+        x = conv_layer(x, filter=in_channel*0.5, kernel=[1,1], layer_name=scope+'_conv1')
+        #x = Drop_out(x, rate=dropout_rate, training=self.training)
+        x = Average_pooling(x, pool_size=[2,2], stride=2)
+
+        return x
+
+def dense_block(input_x, nb_layers, is_training, layer_name):
+    with tf.name_scope(layer_name):
+        layers_concat = list()
+        layers_concat.append(input_x)
+
+        x = bottleneck_layer(input_x, is_training, scope=layer_name + '_bottleN_' + str(0))
+
+        layers_concat.append(x)
+
+        for i in range(nb_layers - 1):
+            x = Concatenation(layers_concat)
+            x = bottleneck_layer(x, is_training, scope=layer_name + '_bottleN_' + str(i + 1))
+            layers_concat.append(x)
+
+        x = Concatenation(layers_concat)
+
+        return x
+
+def Dense_net(input_x, is_training):
+    x = conv_layer(input_x, filter=init_layers , kernel=[7,7], stride=2, layer_name='conv0')
+    x = Max_Pooling(x, pool_size=[3,3], stride=2)
+
+    for i in range(nb_blocks-1) :
+        # 6 -> 12 -> 48
+        x = dense_block(input_x=x, nb_layers=nb_blocks_layers[i], is_training=is_training, layer_name='dense_'+str(i))
+        x = transition_layer(x, is_training, scope='trans_'+str(i))
+
+    """
+    x = self.dense_block(input_x=x, nb_layers=6, layer_name='dense_1')
+    x = self.transition_layer(x, scope='trans_1')
+    x = self.dense_block(input_x=x, nb_layers=12, layer_name='dense_2')
+    x = self.transition_layer(x, scope='trans_2')
+    x = self.dense_block(input_x=x, nb_layers=48, layer_name='dense_3')
+    x = self.transition_layer(x, scope='trans_3')
+    """
+
+    x = dense_block(input_x=x, nb_layers=nb_blocks_layers[nb_blocks-1], is_training=is_training, layer_name='dense_final')
+
+    # 100 Layer
+    x = Batch_Normalization(x, training=is_training, scope='linear_batch')
+    x = Relu(x)
+    x = Global_Average_Pooling(x)
+    x = flatten(x)
+    x = Linear(x)
+
+    # x = tf.reshape(x, [-1, 10])
+    return x
@@ -0,0 +1,44 @@
+import tensorflow as tf
+import math
+import numpy as np
+
+def warmup_cosine_annealing_lr(lr, steps_per_epoch, warmup_epochs, max_epoch, T_max, eta_min=0):
+    base_lr = lr
+    warmup_init_lr = 0
+    total_steps = int(max_epoch * steps_per_epoch)
+    warmup_steps = int(warmup_epochs * steps_per_epoch)
+
+    lr_each_step = []
+    for i in range(total_steps):
+        last_epoch = i // steps_per_epoch
+        if i < warmup_steps:
+            lr = linear_warmup_lr(i + 1, warmup_steps, base_lr, warmup_init_lr)
+        else:
+            lr = eta_min + (base_lr - eta_min) * (1. + math.cos(math.pi*last_epoch / T_max)) / 2
+        lr_each_step.append(lr)
+
+    return np.array(lr_each_step).astype(np.float32)
+
+
+class HyperParams:
+    def __init__(self, args):
+        self.args=args
+        nsteps_per_epoch = self.args.num_training_samples // self.args.global_batch_size
+        self.args.nsteps_per_epoch = nsteps_per_epoch
+        if self.args.max_epochs:
+            nstep = nsteps_per_epoch * self.args.max_epochs
+        else:
+            nstep = self.args.max_train_steps
+        self.args.nstep = nstep
+
+        self.cos_lr = warmup_cosine_annealing_lr(self.args.lr, nsteps_per_epoch, 0, self.args.T_max, self.args.T_max, 0.0)
+
+    def get_learning_rate(self): 
+        global_step = tf.train.get_global_step()
+    
+        learning_rate = tf.gather(tf.convert_to_tensor(self.cos_lr), global_step)
+
+        learning_rate = tf.identity(learning_rate, 'learning_rate')
+
+        return learning_rate
+
@@ -0,0 +1,25 @@
+import tensorflow as tf
+#from tensorflow.contrib.hccl.python.ops import hccl_ops
+#from npu_bridge.hccl import hccl_ops
+from benchmark_log import hwlog
+
+class Layers: 
+    def get_accuracy(self, labels, predicted_classes, logits, args):
+        accuracy = tf.metrics.accuracy(
+            labels=labels, predictions=predicted_classes) 
+        top5acc = tf.metrics.mean(
+            tf.cast(tf.nn.in_top_k(logits, labels, 5), tf.float32))
+        if args.rank_size == 1:
+            newaccuracy = (accuracy[0], accuracy[1])
+            newtop5acc = (top5acc[0], top5acc[1])
+        else:
+            from npu_bridge.hccl import hccl_ops
+            newaccuracy = (hccl_ops.allreduce(accuracy[0],"sum")/args.rank_size, accuracy[1])
+            newtop5acc = (hccl_ops.allreduce(top5acc[0],"sum")/args.rank_size, top5acc[1])
+        metrics = {'val-top1acc': newaccuracy, 'val-top5acc': newtop5acc}
+        
+        return metrics
+
+
+
+
@@ -0,0 +1,92 @@
+from __future__ import print_function
+import tensorflow as tf
+from benchmark_log import hwlog
+import logging
+import numpy as np
+import time
+import sys,os
+
+class LogSessionRunHook(tf.train.SessionRunHook):
+    def __init__(self, args, warmup_steps=5):
+        self.global_batch_size = args.global_batch_size
+        if args.iterations_per_loop is not None:
+            self.iterations_per_loop = args.iterations_per_loop
+        else:
+            self.iterations_per_loop = args.nsteps_per_epoch
+        self.warmup_steps = warmup_steps
+        self.iter_times = []
+        self.num_records = args.num_training_samples
+        self.display_every = args.display_every
+        self.logger = get_logger(args.log_name, args.log_dir)
+        rank0log(self.logger, 'PY' + str(sys.version) + 'TF' + str(tf.__version__))
+
+
+
+    def after_create_session(self, session, coord):
+        rank0log(self.logger, 'Step   Epoch   Speed   Loss   FinLoss   LR')
+        self.elapsed_secs = 0.
+        self.count = 0
+
+    def before_run(self, run_context):
+        self.t0 = time.time()
+        return tf.train.SessionRunArgs( 
+            fetches=[tf.train.get_global_step(), 'loss:0', 'total_loss:0', 'learning_rate:0'])
+
+    def after_run(self, run_context, run_values):
+        batch_time = time.time() - self.t0
+        self.iter_times.append(batch_time)
+        self.elapsed_secs += batch_time
+        self.count += 1
+        global_step, loss, total_loss, lr = run_values.results
+        if global_step == 1 or global_step % self.display_every == 0:
+            dt = self.elapsed_secs / self.count
+            img_per_sec = self.global_batch_size * self.iterations_per_loop / dt
+            epoch = global_step * self.global_batch_size / self.num_records
+            self.logger.info('step:%6i  epoch:%5.1f  FPS:%7.1f  loss:%6.3f  total_loss:%6.3f  lr:%7.5f' %
+                             (global_step, epoch, img_per_sec, loss, total_loss, lr))
+            self.elapsed_secs = 0.
+            self.count = 0
+            
+            # add by wx983399
+            hwlog.remark_print(key=hwlog.GLOBAL_STEP, value=int(global_step))
+            hwlog.remark_print(key=hwlog.CURRENT_EPOCH, value=epoch)
+            hwlog.remark_print(key=hwlog.FPS, value=img_per_sec)
+
+    def get_average_speed(self):
+        avg_time = np.mean(self.iter_times[self.warmup_steps:])
+        speed = self.global_batch_size / avg_time
+        return speed
+
+
+
+def rank0log(logger, *args, **kwargs):
+    if logger: 
+        logger.info(''.join([str(x) for x in list(args)]))
+    else:
+        print(*args, **kwargs)
+
+
+def get_logger(log_name, log_dir):
+    logger = logging.getLogger(log_name)
+    logger.setLevel(logging.INFO)  # INFO, ERROR
+    # file handler which logs debug messages
+    if not os.path.isdir(log_dir):
+        try:
+            os.makedirs(log_dir)
+        except FileExistsError:
+            # if log_dir is common for multiple ranks like on nfs
+            pass
+    # console handler
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+    # add formatter to the handlers
+    formatter = logging.Formatter('%(message)s')
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+    fh = logging.FileHandler(os.path.join(log_dir, log_name))
+    fh.setLevel(logging.DEBUG)
+    fh.setFormatter(formatter)
+    # add handlers to logger
+    logger.addHandler(fh)
+    return logger
+
@@ -0,0 +1,72 @@
+import tensorflow as tf
+from densenet import Dense_net
+
+
+class Model(object):
+    def __init__(self, args, data, hyper_param, layers, logger):
+        self.args = args
+        self.data = data
+        self.hyper_param = hyper_param
+        self.layers = layers
+        self.logger = logger  
+
+    def get_estimator_model_func(self, features, labels, mode, params=None):
+        labels = tf.reshape(labels, (-1,))  # Squash unnecessary unary dim         #----------------not use when use onehot label
+    
+        inputs = features  # TODO: Should be using feature columns?
+        is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+
+        inputs = tf.cast(inputs, self.args.dtype)
+
+        top_layer = Dense_net(inputs, is_training)
+
+        logits = top_layer
+        predicted_classes = tf.argmax(logits, axis=1, output_type=tf.int32)
+        logits = tf.cast(logits, tf.float32)
+
+        labels_one_hot = tf.one_hot(labels, depth=1000)
+        loss = tf.losses.softmax_cross_entropy(
+            logits=logits, onehot_labels=labels_one_hot, label_smoothing=self.args.label_smoothing)
+
+
+        base_loss = tf.identity(loss, name='loss')  # For access by logger (TODO: Better way to access it?)
+
+        l2_loss = tf.add_n([tf.nn.l2_loss(tf.cast(v, tf.float32)) for v in tf.trainable_variables()])
+        l2_loss = tf.multiply(l2_loss, self.args.weight_decay)
+        total_loss = base_loss + l2_loss
+
+        total_loss = tf.identity(total_loss, name = 'total_loss')
+
+        if mode == tf.estimator.ModeKeys.EVAL:
+            with tf.device(None):
+                metrics = self.layers.get_accuracy( labels, predicted_classes, logits, self.args)
+
+            return tf.estimator.EstimatorSpec(
+                mode, loss=loss, eval_metric_ops=metrics)
+
+        assert (mode == tf.estimator.ModeKeys.TRAIN)
+
+        batch_size = tf.shape(inputs)[0]
+
+        global_step = tf.train.get_global_step()
+        learning_rate = self.hyper_param.get_learning_rate()
+
+        momentum = self.args.momentum
+
+        opt = tf.train.MomentumOptimizer(
+            learning_rate, momentum, use_nesterov=self.args.use_nesterov)
+
+        from npu_bridge.estimator.npu.npu_optimizer import NPUDistributedOptimizer
+        opt = NPUDistributedOptimizer(opt)
+
+        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) or []
+
+        with tf.control_dependencies(update_ops):
+            gate_gradients = tf.train.Optimizer.GATE_NONE
+            grads_and_vars = opt.compute_gradients(total_loss, gate_gradients=gate_gradients)
+            train_op = opt.apply_gradients(grads_and_vars, global_step=global_step)
+
+        train_op = tf.group(train_op)
+        
+        return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op)
+
@@ -0,0 +1,72 @@
+import tensorflow as tf
+from tensorflow.contrib.image.python.ops import distort_image_ops
+import math
+import random
+
+def decode_jpeg(imgdata, channels=3):
+    return tf.image.decode_jpeg(imgdata, channels=channels,
+                                fancy_upscaling=False,
+                                dct_method='INTEGER_FAST')
+
+
+def random_horizontal_flip(image, prob):
+    if prob > random.random():
+        image = tf.image.flip_left_right(image)
+    return image
+
+
+def decode_crop_and_resize(record, bbox, size, scale, ratio):
+    with tf.name_scope('decode_crop_and_resize'):
+        height = 224
+        width = 224
+        crop_ratio = 0.8
+        initial_shape = [int(round(height / crop_ratio)),
+                int(round(width / crop_ratio)), 3]
+        jpeg_shape = tf.image.extract_jpeg_shape( record )
+
+        bbox_begin, bbox_size, bbox = \
+            tf.image.sample_distorted_bounding_box(
+                tf.image.extract_jpeg_shape(record),
+                bounding_boxes=bbox,
+                min_object_covered=0.1,
+                aspect_ratio_range=ratio,
+                area_range=scale,
+                max_attempts=10,
+                use_image_if_no_bounding_boxes=True)
+
+         # Reassemble the bounding box in the format the crop op requires.
+        offset_y, offset_x, _ = tf.unstack(bbox_begin)
+        target_height, target_width, _ = tf.unstack(bbox_size)
+        crop_window = tf.stack([offset_y, offset_x, target_height, target_width])
+
+        image = tf.image.decode_and_crop_jpeg( record, crop_window, channels=3 )
+        image = tf.image.resize_images( image, [height, width] )
+
+        return image
+
+
+def parse_and_preprocess_image_record(record, bbox, training):
+    with tf.name_scope('preprocess'):
+        if training:
+            image = decode_crop_and_resize(record, bbox, 224, (0.08, 1.0), (0.75, 1.333))
+            image = random_horizontal_flip(image, 0.5)
+            image = normalize(image)
+        else:
+            image = decode_jpeg(record, channels=3)
+            image = tf.image.resize_images(image, [256, 256])
+            image = tf.image.central_crop(image, 224.0/256)
+            image = normalize(image)
+
+    return image
+
+
+def normalize(inputs):
+     imagenet_mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
+     imagenet_std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
+     imagenet_mean = tf.expand_dims(tf.expand_dims(imagenet_mean, 0), 0)
+     imagenet_std = tf.expand_dims(tf.expand_dims(imagenet_std, 0), 0)
+     inputs = inputs - imagenet_mean
+     inputs = inputs * (1.0 / imagenet_std)
+
+     return inputs
+
@@ -0,0 +1,140 @@
+import tensorflow as tf
+import numpy as np
+import os
+import sys
+import ast
+sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), '../')))
+sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), '../config')))
+sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), '../../../../utils')))
+sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), '../../../../utils/atlasboost')))
+
+import data_loader as dl
+
+import model as ml
+import hyper_param as hp
+import layers as ly
+import logger as lg
+import trainer as tr
+import create_session as cs
+import argparse
+
+from benchmark_log import hwlog
+from benchmark_log.basic_utils import get_environment_info
+from benchmark_log.basic_utils import get_model_parameter
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    
+
+    parser.add_argument('--rank_size', default=1,type=int,
+                        help="""number of NPUs  to use.""")
+
+    # mode and parameters related 
+    parser.add_argument('--mode', default='train_and_evaluate',
+                        help="""mode to run the program  e.g. train, evaluate, and
+                        train_and_evaluate""")
+    parser.add_argument('--max_train_steps', default=100,type=int,
+                        help="""train steps for one NPU""")
+    parser.add_argument('--iterations_per_loop', default=10, type=int,
+                        help="""the number of steps in devices for each iteration""")
+    parser.add_argument('--max_epochs', default=None, type=int,
+                        help="""total epochs for training""")
+    parser.add_argument('--epochs_between_evals', default=5, type=int,
+                        help="""the interval between train and evaluation , only meaningful
+                        when the mode is train_and_evaluate""")
+
+    # dataset
+    parser.add_argument('--data_dir', default='path/data',
+                        help="""directory to data.""")
+
+    # path for evaluation
+    parser.add_argument('--eval_dir', default='path/eval',
+                        help="""directory to evaluate.""")
+
+    parser.add_argument('--dtype', default=tf.float32,
+                        help="""data type of inputs.""")
+    parser.add_argument('--use_nesterov', default=True, type=ast.literal_eval,
+                        help=""" used in optimizer""")
+    parser.add_argument('--label_smoothing', default=0.1, type=float,
+                        help="""label smoothing factor""")
+    parser.add_argument('--weight_decay', default=0.0001,
+                        help="""weight decay""")
+    parser.add_argument('--batch_size', default=32, type=int,
+                        help="""batch size for one NPU""")
+
+    # learning rate and momentum
+    parser.add_argument('--lr', default=0.1, type=float,
+                        help="""learning rate""")
+    parser.add_argument('--T_max', default=150, type=int,
+                        help="""T_max for cosing_annealing learning rate""")
+    parser.add_argument('--momentum', default=0.9, type=float,
+                        help="""momentum used in optimizer.""")
+
+    # display frequency
+    parser.add_argument('--display_every', default=1, type=int,
+                        help="""the frequency to display info""")
+
+    # log file
+    parser.add_argument('--log_name', default='densenet121_training.log',
+                        help="""name of log file""")
+    parser.add_argument('--log_dir', default='./model_1p',
+                        help="""log directory""")
+ 
+    args, unknown_args = parser.parse_known_args()
+    # ['--config_file', 'densenet_config_1p_npu']
+
+    print(args, unknown_args)
+    if len(unknown_args) > 0:
+        for bad_arg in unknown_args:
+            print("ERROR: Unknown command line arg: %s" % bad_arg)
+        raise ValueError("Invalid command line arg(s)")
+
+    return args
+
+
+def main():
+
+    args = parse_args()
+    args.global_batch_size = args.batch_size * args.rank_size
+
+    session = cs.CreateSession()
+    data = dl.DataLoader(args)    
+    hyper_param = hp.HyperParams(args)
+    layers = ly.Layers() 
+    logger = lg.LogSessionRunHook(args)
+    model = ml.Model(args, data, hyper_param, layers, logger)
+   
+    trainer = tr.Trainer(session, args, data, model, logger)
+
+    if args.mode == 'train':
+        trainer.train()
+    elif args.mode == 'evaluate':
+        trainer.evaluate()
+    elif args.mode == 'train_and_evaluate':
+        trainer.train_and_evaluate()
+    else:
+        raise ValueError("Invalid mode.")
+
+
+if __name__ == '__main__':
+
+    hwlog.ROOT_DIR = os.path.split(os.path.abspath(__file__))[0]
+    cpu_info, npu_info, framework_info, os_info, benchmark_version = get_environment_info("tensorflow")
+    config_info = get_model_parameter("tensorflow_config")
+    initinal_data = {"base_lr": 0.128, "dataset": "imagenet1024", "optimizer": "SGD", "loss_scale": 512,
+                     "batchsize": 32}
+    hwlog.remark_print(key=hwlog.CPU_INFO, value=cpu_info)
+    hwlog.remark_print(key=hwlog.NPU_INFO, value=npu_info)
+    hwlog.remark_print(key=hwlog.OS_INFO, value=os_info)
+    hwlog.remark_print(key=hwlog.FRAMEWORK_INFO, value=framework_info)
+    hwlog.remark_print(key=hwlog.BENCHMARK_VERSION, value=benchmark_version)
+    hwlog.remark_print(key=hwlog.CONFIG_INFO, value=config_info)
+    hwlog.remark_print(key=hwlog.BASE_LR, value=initinal_data.get("base_lr"))
+    hwlog.remark_print(key=hwlog.DATASET, value=initinal_data.get("dataset"))
+    hwlog.remark_print(key=hwlog.OPT_NAME, value=initinal_data.get("optimizer"))
+    hwlog.remark_print(key=hwlog.LOSS_SCALE, value=initinal_data.get("loss_scale"))
+    hwlog.remark_print(key=hwlog.INPUT_BATCH_SIZE, value=initinal_data.get("batchsize"))
+    main()
+
@@ -0,0 +1,22 @@
+import tensorflow as tf
+from tensorflow.python.ops import data_flow_ops
+import re
+import os
+from operator import itemgetter
+
+
+def sort_and_load_ckpts(log_dir):
+    ckpts = []
+    for f in os.listdir(log_dir):
+        m = re.match(r'model.ckpt-([0-9]+).index', f)
+        if m is None:
+            continue
+        fullpath = os.path.join(log_dir, f)
+        ckpts.append({'step': int(m.group(1)),
+                      'path': os.path.splitext(fullpath)[0],
+                      'mtime': os.stat(fullpath).st_mtime,
+                      })
+    ckpts.sort(key=itemgetter('step'))
+    return ckpts
+
+
@@ -0,0 +1,128 @@
+import tensorflow as tf
+import math
+import time
+import os
+import train_helper
+from logger import rank0log
+from benchmark_log import hwlog
+
+class Trainer(object):
+    def __init__(self, session, args, data, model, logger):
+        self.sess = session
+        self.args = args
+        self.data = data
+        self.model = model
+        self.logger = logger
+        self.print_logger = self.logger.logger
+        self.all_preds = []
+        self.all_targets = []
+
+        self.classifier, self.training_hook = self.get_npu_classifier()
+
+    def get_npu_classifier(self):
+        from npu_bridge.estimator.npu.npu_config import NPURunConfig
+        from npu_bridge.estimator.npu.npu_estimator import NPUEstimator
+
+        run_config = NPURunConfig(
+            hcom_parallel=True,
+            precision_mode="allow_mix_precision",
+            enable_data_pre_proc=True,
+            save_checkpoints_steps=self.args.nsteps_per_epoch,
+            session_config=self.sess.estimator_config,
+            model_dir=self.args.log_dir,
+            iterations_per_loop=self.args.iterations_per_loop,
+            keep_checkpoint_max=5)
+
+        classifier =NPUEstimator(
+            model_fn= self.model.get_estimator_model_func, 
+            config= run_config
+      	  )
+      
+        training_hooks = []
+        training_hooks.append(self.logger)
+
+        return classifier, training_hooks
+
+    def train(self):
+        print ('training steps: %d' % self.args.nstep)
+        self.classifier.train( input_fn=lambda:self.data.get_train_input_fn(),
+                               max_steps = self.args.nstep,
+                               hooks = self.training_hook
+                              )
+
+    def evaluate(self):
+        rank0log(self.print_logger, "Evaluating")
+        rank0log(self.print_logger, "Validation dataset size: {}".format(self.args.num_evaluating_samples))
+        time.sleep(5)  # a little extra margin...
+        try:
+            ckpts = train_helper.sort_and_load_ckpts(self.args.eval_dir)
+            print("=========ckpt==========")
+            print(ckpts)
+            print("=========ckpt==========")
+            for i, c in enumerate(ckpts):
+                eval_result = self.classifier.evaluate(
+                    input_fn=lambda: self.data.get_eval_input_fn(),
+                    checkpoint_path=c['path'])
+                c['epoch'] = math.ceil(c['step'] / (self.args.num_training_samples/ (self.args.batch_size)))
+                c['top1'] = eval_result['val-top1acc']
+                c['top5'] = eval_result['val-top5acc']
+                c['loss'] = eval_result['loss']
+                
+            rank0log(self.print_logger, ' step  epoch  top1    top5     loss   checkpoint_time(UTC)')
+            for i, c in enumerate(ckpts):
+                if 'top1' not in c:
+                    continue
+                rank0log(self.print_logger,'{:5d}  {:5.1f}  {:5.3f}  {:6.2f}  {:6.2f}  {time}'
+                         .format(c['step'],
+                                 c['epoch'],
+                                 c['top1'] * 100,
+                                 c['top5'] * 100,
+                                 c['loss'],
+                                 time=time.strftime('%Y-%m-%d %H:%M:%S', 
+                                    time.localtime(c['mtime']))))
+            rank0log(self.print_logger, "Finished evaluation")
+        except KeyboardInterrupt:
+            self.print_logger.error("Keyboard interrupt")
+
+    def train_and_evaluate(self):
+        epochs_between_evals = self.args.epochs_between_evals
+
+        for i in range(self.args.max_epochs // epochs_between_evals):
+
+            rank0log(self.print_logger, "Starting a training cycle")
+
+            self.classifier.train(input_fn=lambda:self.data.get_train_input_fn(),
+                            steps = self.args.nsteps_per_epoch*epochs_between_evals,
+                            hooks = self.training_hook )
+
+            rank0log(self.print_logger, "Starting to evaluate")
+            rank0log(self.print_logger, "Validation dataset size: {}".format(self.args.num_evaluating_samples))
+            time.sleep(5)  # a little extra margin...
+
+            ckpts = train_helper.sort_and_load_ckpts(self.args.log_dir)
+            c = ckpts[-1]
+            eval_result = self.classifier.evaluate(
+                input_fn=lambda: self.data.get_eval_input_fn(),
+                checkpoint_path=c['path'])
+
+            # top1  top5  Log dotting
+            hwlog.remark_print(key=hwlog.EVAL_ACCURACY_TOP1, value=float(eval_result.get("val-top1acc")))
+            hwlog.remark_print(key=hwlog.EVAL_ACCURACY_TOP1, value=float(eval_result.get("val-top5acc")))
+
+            
+            c['epoch'] = math.ceil(c['step'] / (self.args.num_training_samples / (self.args.batch_size * self.args.rank_size)))
+            c['top1'] = eval_result['val-top1acc']
+            c['top5'] = eval_result['val-top5acc']
+            c['loss'] = eval_result['loss']
+
+            rank0log(self.print_logger, ' step  epoch  top1    top5     loss   checkpoint_time(UTC)')
+            
+            rank0log(self.print_logger,'{:5d}  {:5.1f}  {:5.3f}  {:6.2f}  {:6.2f}  {time}'
+                    .format(c['step'],
+                            c['epoch'],
+                            c['top1'] * 100,
+                            c['top5'] * 100,
+                            c['loss'],
+                            time=time.strftime('%Y-%m-%d %H:%M:%S',
+                                time.localtime(c['mtime']))))
+