code: def main():
      data_root = "/data/flower_photos"  # get data root path
  
      if not os.path.exists("./save_weights"):
          os.makedirs("./save_weights")
  
      im_height = 224
      im_width = 224
      batch_size = 16
      epochs = 30
      num_classes = 5
  
      log_dir = "./logs/" + datetime.datetime.now().strftime(
          "%Y%m%d-%H%M%S"
      )
      train_writer = tf.summary.create_file_writer(
          os.path.join(log_dir, "train")
      )
      val_writer = tf.summary.create_file_writer(
          os.path.join(log_dir, "val")
      )
  
      # data generator with data augmentation
      train_ds, val_ds = generate_ds(
          data_root, im_height, im_width, batch_size
      )
  
      # create model
      model = shufflenet_v2_x1_0(
          input_shape=(im_height, im_width, 3), num_classes=num_classes
      )
  
      # load weights
      # x1.0权重链接: https://pan.baidu.com/s/1M2mp98Si9eT9qT436DcdOw  密码: mhts
      pre_weights_path = "./shufflenetv2_x1_0.h5"
      assert os.path.exists(pre_weights_path), "cannot find {}".format(
          pre_weights_path
      )
      model.load_weights(
          pre_weights_path, by_name=True, skip_mismatch=True
      )
  
      model.summary()
  
      # custom learning rate curve
      def scheduler(now_epoch):
          initial_lr = 0.1
          end_lr_rate = 0.1  # end_lr = initial_lr * end_lr_rate
          rate = ((1 + math.cos(now_epoch * math.pi / epochs)) / 2) * (
              1 - end_lr_rate
          ) + end_lr_rate  # cosine
          new_lr = rate * initial_lr
  
          # writing lr into tensorboard
          with train_writer.as_default():
              tf.summary.scalar(
                  "learning rate", data=new_lr, step=epoch
              )
  
          return new_lr
  
      # using keras low level api for training
      loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
          from_logits=False
      )
      optimizer = tf.keras.optimizers.SGD(
          learning_rate=0.1, momentum=0.9
      )
  
      train_loss = tf.keras.metrics.Mean(name="train_loss")
      train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="train_accuracy"
      )
  
      val_loss = tf.keras.metrics.Mean(name="val_loss")
      val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="val_accuracy"
      )
  
      @tf.function
      def train_step(train_images, train_labels):
          with tf.GradientTape() as tape:
              output = model(train_images, training=True)
              loss = loss_object(train_labels, output)
          gradients = tape.gradient(loss, model.trainable_variables)
          optimizer.apply_gradients(
              zip(gradients, model.trainable_variables)
          )
  
          train_loss(loss)
          train_accuracy(train_labels, output)
  
      @tf.function
      def val_step(val_images, val_labels):
          output = model(val_images, training=False)
          loss = loss_object(val_labels, output)
  
          val_loss(loss)
          val_accuracy(val_labels, output)
  
      best_val_acc = 0.0
      for epoch in range(epochs):
          train_loss.reset_states()  # clear history info
          train_accuracy.reset_states()  # clear history info
          val_loss.reset_states()  # clear history info
          val_accuracy.reset_states()  # clear history info
  
          # train
          train_bar = tqdm(train_ds, file=sys.stdout)
          for images, labels in train_bar:
              train_step(images, labels)
  
              # print train process
              train_bar.desc = (
                  "train epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(
                      epoch + 1,
                      epochs,
                      train_loss.result(),
                      train_accuracy.result(),
                  )
              )
  
          # update learning rate
          optimizer.learning_rate = scheduler(epoch)
  
          # validate
          val_bar = tqdm(val_ds, file=sys.stdout)
          for images, labels in val_bar:
              val_step(images, labels)
  
              # print val process
              val_bar.desc = (
                  "valid epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(
                      epoch + 1,
                      epochs,
                      val_loss.result(),
                      val_accuracy.result(),
                  )
              )
          # writing training loss and acc
          with train_writer.as_default():
              tf.summary.scalar("loss", train_loss.result(), epoch)
              tf.summary.scalar(
                  "accuracy", train_accuracy.result(), epoch
              )
  
          # writing validation loss and acc
          with val_writer.as_default():
              tf.summary.scalar("loss", val_loss.result(), epoch)
              tf.summary.scalar(
                  "accuracy", val_accuracy.result(), epoch
              )
  
          # only save best weights
          if val_accuracy.result() > best_val_acc:
              best_val_acc = val_accuracy.result()
              model.save_weights(
                  "./save_weights/shufflenetv2.ckpt", save_format="tf"
              )
  
file_path: deep-learning-for-image-processing/tensorflow_classification/Test7_shuffleNet/train.py
imports: import os
  import sys
  import math
  import datetime
  import tensorflow
  from tqdm import tqdm
  from model import shufflenet_v2_x1_0
  from utils import generate_ds
label_desc: Use image node_id 1 for calling a global function with example usage: main() and returns: new_lr