code: def main(unused_argv=None):
      tf.logging.set_verbosity(tf.logging.INFO)
      with tf.Graph().as_default():
          # Forces all input processing onto CPU in order to reserve the GPU for the
          # forward inference and back-propagation.
          device = (
              "/cpu:0" if not FLAGS.ps_tasks else "/job:worker/cpu:0"
          )
          with tf.device(
              tf.train.replica_device_setter(
                  FLAGS.ps_tasks, worker_device=device
              )
          ):
              # Loads content images.
              content_inputs_, _ = image_utils.imagenet_inputs(
                  FLAGS.batch_size, FLAGS.image_size
              )
  
              # Loads style images.
              [
                  style_inputs_,
                  _,
                  style_inputs_orig_,
              ] = image_utils.arbitrary_style_image_inputs(
                  FLAGS.style_dataset_file,
                  batch_size=FLAGS.batch_size,
                  image_size=FLAGS.image_size,
                  shuffle=True,
                  center_crop=FLAGS.center_crop,
                  augment_style_images=FLAGS.augment_style_images,
                  random_style_image_size=FLAGS.random_style_image_size,
              )
  
          with tf.device(
              tf.train.replica_device_setter(FLAGS.ps_tasks)
          ):
              # Process style and content weight flags.
              content_weights = ast.literal_eval(FLAGS.content_weights)
              style_weights = ast.literal_eval(FLAGS.style_weights)
  
              # Define the model
              (
                  stylized_images,
                  total_loss,
                  loss_dict,
                  _,
              ) = build_model.build_model(
                  content_inputs_,
                  style_inputs_,
                  trainable=True,
                  is_training=True,
                  inception_end_point="Mixed_6e",
                  style_prediction_bottleneck=100,
                  adds_losses=True,
                  content_weights=content_weights,
                  style_weights=style_weights,
                  total_variation_weight=FLAGS.total_variation_weight,
              )
  
              # Adding scalar summaries to the tensorboard.
              for key, value in loss_dict.items():
                  tf.summary.scalar(key, value)
  
              # Adding Image summaries to the tensorboard.
              tf.summary.image(
                  "image/0_content_inputs", content_inputs_, 3
              )
              tf.summary.image(
                  "image/1_style_inputs_orig", style_inputs_orig_, 3
              )
              tf.summary.image(
                  "image/2_style_inputs_aug", style_inputs_, 3
              )
              tf.summary.image(
                  "image/3_stylized_images", stylized_images, 3
              )
  
              # Set up training
              optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
              train_op = slim.learning.create_train_op(
                  total_loss,
                  optimizer,
                  clip_gradient_norm=FLAGS.clip_gradient_norm,
                  summarize_gradients=False,
              )
  
              # Function to restore VGG16 parameters.
              init_fn_vgg = slim.assign_from_checkpoint_fn(
                  vgg.checkpoint_file(), slim.get_variables("vgg_16")
              )
  
              # Function to restore Inception_v3 parameters.
              inception_variables_dict = {
                  var.op.name: var
                  for var in slim.get_model_variables("InceptionV3")
              }
              init_fn_inception = slim.assign_from_checkpoint_fn(
                  FLAGS.inception_v3_checkpoint,
                  inception_variables_dict,
              )
  
              # Function to restore VGG16 and Inception_v3 parameters.
              def init_sub_networks(session):
                  init_fn_vgg(session)
                  init_fn_inception(session)
  
              # Run training
              slim.learning.train(
                  train_op=train_op,
                  logdir=os.path.expanduser(FLAGS.train_dir),
                  master=FLAGS.master,
                  is_chief=FLAGS.task == 0,
                  number_of_steps=FLAGS.train_steps,
                  init_fn=init_sub_networks,
                  save_summaries_secs=FLAGS.save_summaries_secs,
                  save_interval_secs=FLAGS.save_interval_secs,
              )
  
file_path: magenta/magenta/models/arbitrary_image_stylization/arbitrary_image_stylization_train.py
imports: import ast
  import os
  from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_model
  from magenta.models.image_stylization import image_utils
  from magenta.models.image_stylization import vgg
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types