code: #!/usr/bin/env/python
  
  from typing import Tuple, List, Any, Sequence
  
  import tensorflow as tf
  import time
  import os
  import json
  import numpy as np
  import pickle
  import random
  
  from utils import MLP, ThreadedIterator, SMALL_NUMBER
  
  
  class ChemModel(object):
      @classmethod
      def default_params(cls):
          return {
              'num_epochs': 3000,
              'patience': 25,
              'learning_rate': 0.001,
              'clamp_gradient_norm': 1.0,
              'out_layer_dropout_keep_prob': 1.0,
  
              'hidden_size': 100,
              'num_timesteps': 4,
              'use_graph': True,
  
              'tie_fwd_bkwd': True,
              'task_ids': [0],
  
              'random_seed': 0,
  
              'train_file': 'molecules_train.json',
              'valid_file': 'molecules_valid.json'
          }
  
      def __init__(self, args):
          self.args = args
  
          # Collect argument things:
          data_dir = ''
          if '--data_dir' in args and args['--data_dir'] is not None:
              data_dir = args['--data_dir']
          self.data_dir = data_dir
  
          self.run_id = "_".join([time.strftime("%Y-%m-%d-%H-%M-%S"), str(os.getpid())])
          log_dir = args.get('--log_dir') or '.'
          self.log_file = os.path.join(log_dir, "%s_log.json" % self.run_id)
          self.best_model_file = os.path.join(log_dir, "%s_model_best.pickle" % self.run_id)
  
          # Collect parameters:
          params = self.default_params()
          config_file = args.get('--config-file')
          if config_file is not None:
              with open(config_file, 'r') as f:
                  params.update(json.load(f))
          config = args.get('--config')
          if config is not None:
              params.update(json.loads(config))
          self.params = params
          with open(os.path.join(log_dir, "%s_params.json" % self.run_id), "w") as f:
              json.dump(params, f)
          print("Run %s starting with following parameters:\n%s" % (self.run_id, json.dumps(self.params)))
          random.seed(params['random_seed'])
          np.random.seed(params['random_seed'])
  
  
          # Load data:
          self.max_num_vertices = 0
          self.num_edge_types = 0
          self.annotation_size = 0
          self.train_data = self.load_data(params['train_file'], is_training_data=True)
          self.valid_data = self.load_data(params['valid_file'], is_training_data=False)
  
          # Build the actual model
          config = tf.ConfigProto()
          config.gpu_options.allow_growth = True
          self.graph = tf.Graph()
          self.sess = tf.Session(graph=self.graph, config=config)
          with self.graph.as_default():
              tf.set_random_seed(params['random_seed'])
              self.placeholders = {}
              self.weights = {}
              self.ops = {}
              self.make_model()
              self.make_train_step()
  
              # Restore/initialize variables:
              restore_file = args.get('--restore')
              if restore_file is not None:
                  self.restore_model(restore_file)
              else:
                  self.initialize_model()
  
      def load_data(self, file_name, is_training_data):
          full_path = os.path.join(self.data_dir, file_name)
  
          print("Loading data from %s" % full_path)
          with open(full_path, 'r') as f:
              data = json.load(f)
  
          restrict = self.args.get("--restrict_data")
          if restrict is not None and restrict > 0:
              data = data[:restrict]
  
          # Get some common data out:
          num_fwd_edge_types = 0
          for g in data:
              self.max_num_vertices = max(self.max_num_vertices, max([v for e in g['graph'] for v in [e[0], e[2]]]))
              num_fwd_edge_types = max(num_fwd_edge_types, max([e[1] for e in g['graph']]))
          self.num_edge_types = max(self.num_edge_types, num_fwd_edge_types * (1 if self.params['tie_fwd_bkwd'] else 2))
          self.annotation_size = max(self.annotation_size, len(data[0]["node_features"][0]))
  
          return self.process_raw_graphs(data, is_training_data)
  
      @staticmethod
      def graph_string_to_array(graph_string):
          return [[int(v) for v in s.split(' ')]
                  for s in graph_string.split('\n')]
  
      def process_raw_graphs(self, raw_data, is_training_data):
          raise Exception("Models have to implement process_raw_graphs!")
  
      def make_model(self):
          self.placeholders['target_values'] = tf.placeholder(tf.float32, [len(self.params['task_ids']), None],
                                                              name='target_values')
          self.placeholders['target_mask'] = tf.placeholder(tf.float32, [len(self.params['task_ids']), None],
                                                            name='target_mask')
          self.placeholders['num_graphs'] = tf.placeholder(tf.int32, [], name='num_graphs')
          self.placeholders['out_layer_dropout_keep_prob'] = tf.placeholder(tf.float32, [], name='out_layer_dropout_keep_prob')
  
          with tf.variable_scope("graph_model"):
              self.prepare_specific_graph_model()
              # This does the actual graph work:
              if self.params['use_graph']:
                  self.ops['final_node_representations'] = self.compute_final_node_representations()
              else:
                  self.ops['final_node_representations'] = tf.zeros_like(self.placeholders['initial_node_representation'])
  
          self.ops['losses'] = []
          for (internal_id, task_id) in enumerate(self.params['task_ids']):
              with tf.variable_scope("out_layer_task%i" % task_id):
                  with tf.variable_scope("regression_gate"):
                      self.weights['regression_gate_task%i' % task_id] = MLP(2 * self.params['hidden_size'], 1, [],
                                                                             self.placeholders['out_layer_dropout_keep_prob'])
                  with tf.variable_scope("regression"):
                      self.weights['regression_transform_task%i' % task_id] = MLP(self.params['hidden_size'], 1, [],
                                                                                  self.placeholders['out_layer_dropout_keep_prob'])
                  computed_values = self.gated_regression(self.ops['final_node_representations'],
                                                          self.weights['regression_gate_task%i' % task_id],
                                                          self.weights['regression_transform_task%i' % task_id])
                  diff = computed_values - self.placeholders['target_values'][internal_id,:]
                  task_target_mask = self.placeholders['target_mask'][internal_id,:]
                  task_target_num = tf.reduce_sum(task_target_mask) + SMALL_NUMBER
                  diff = diff * task_target_mask  # Mask out unused values
                  self.ops['accuracy_task%i' % task_id] = tf.reduce_sum(tf.abs(diff)) / task_target_num
                  task_loss = tf.reduce_sum(0.5 * tf.square(diff)) / task_target_num
                  # Normalise loss to account for fewer task-specific examples in batch:
                  task_loss = task_loss * (1.0 / (self.params['task_sample_ratios'].get(task_id) or 1.0))
                  self.ops['losses'].append(task_loss)
          self.ops['loss'] = tf.reduce_sum(self.ops['losses'])
  
      def make_train_step(self):
          trainable_vars = self.sess.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
          if self.args.get('--freeze-graph-model'):
              graph_vars = set(self.sess.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="graph_model"))
              filtered_vars = []
              for var in trainable_vars:
                  if var not in graph_vars:
                      filtered_vars.append(var)
                  else:
                      print("Freezing weights of variable %s." % var.name)
              trainable_vars = filtered_vars
          optimizer = tf.train.AdamOptimizer(self.params['learning_rate'])
          grads_and_vars = optimizer.compute_gradients(self.ops['loss'], var_list=trainable_vars)
          clipped_grads = []
          for grad, var in grads_and_vars:
              if grad is not None:
                  clipped_grads.append((tf.clip_by_norm(grad, self.params['clamp_gradient_norm']), var))
              else:
                  clipped_grads.append((grad, var))
          self.ops['train_step'] = optimizer.apply_gradients(clipped_grads)
          # Initialize newly-introduced variables:
          self.sess.run(tf.local_variables_initializer())
  
      def gated_regression(self, last_h, regression_gate, regression_transform):
          raise Exception("Models have to implement gated_regression!")
  
      def prepare_specific_graph_model(self):
          raise Exception("Models have to implement prepare_specific_graph_model!")
  
      def compute_final_node_representations(self):
          raise Exception("Models have to implement compute_final_node_representations!")
  
      def make_minibatch_iterator(self, data, is_training):
          raise Exception("Models have to implement make_minibatch_iterator!")
  
      def run_epoch(self, epoch_name, data, is_training):
          chemical_accuracies = np.array([0.066513725, 0.012235489, 0.071939046, 0.033730778, 0.033486113, 0.004278493,
                                          0.001330901, 0.004165489, 0.004128926, 0.00409976, 0.004527465, 0.012292586,
                                          0.037467458])
  
          loss = 0
          accuracies = []
          accuracy_ops = [self.ops['accuracy_task%i' % task_id] for task_id in self.params['task_ids']]
          start_time = time.time()
          processed_graphs = 0
          batch_iterator = ThreadedIterator(self.make_minibatch_iterator(data, is_training), max_queue_size=5)
          for step, batch_data in enumerate(batch_iterator):
              num_graphs = batch_data[self.placeholders['num_graphs']]
              processed_graphs += num_graphs
              if is_training:
                  batch_data[self.placeholders['out_layer_dropout_keep_prob']] = self.params['out_layer_dropout_keep_prob']
                  fetch_list = [self.ops['loss'], accuracy_ops, self.ops['train_step']]
              else:
                  batch_data[self.placeholders['out_layer_dropout_keep_prob']] = 1.0
                  fetch_list = [self.ops['loss'], accuracy_ops]
              result = self.sess.run(fetch_list, feed_dict=batch_data)
              (batch_loss, batch_accuracies) = (result[0], result[1])
              loss += batch_loss * num_graphs
              accuracies.append(np.array(batch_accuracies) * num_graphs)
  
              print("Running %s, batch %i (has %i graphs). Loss so far: %.4f" % (epoch_name,
                                                                                 step,
                                                                                 num_graphs,
                                                                                 loss / processed_graphs))
  
          accuracies = np.sum(accuracies, axis=0) / processed_graphs
          loss = loss / processed_graphs
          error_ratios = accuracies / chemical_accuracies[self.params["task_ids"]]
          instance_per_sec = processed_graphs / (time.time() - start_time)
          return loss, accuracies, error_ratios, instance_per_sec
  
      def train(self):
          log_to_save = []
          total_time_start = time.time()
          with self.graph.as_default():
              if self.args.get('--restore') is not None:
                  _, valid_accs, _, _ = self.run_epoch("Resumed (validation)", self.valid_data, False)
                  best_val_acc = np.sum(valid_accs)
                  best_val_acc_epoch = 0
                  print("\r\x1b[KResumed operation, initial cum. val. acc: %.5f" % best_val_acc)
              else:
                  (best_val_acc, best_val_acc_epoch) = (float("+inf"), 0)
              for epoch in range(1, self.params['num_epochs'] + 1):
                  print("== Epoch %i" % epoch)
                  train_loss, train_accs, train_errs, train_speed = self.run_epoch("epoch %i (training)" % epoch,
                                                                                   self.train_data, True)
                  accs_str = " ".join(["%i:%.5f" % (id, acc) for (id, acc) in zip(self.params['task_ids'], train_accs)])
                  errs_str = " ".join(["%i:%.5f" % (id, err) for (id, err) in zip(self.params['task_ids'], train_errs)])
                  print("\r\x1b[K Train: loss: %.5f | acc: %s | error_ratio: %s | instances/sec: %.2f" % (train_loss,
                                                                                                          accs_str,
                                                                                                          errs_str,
                                                                                                          train_speed))
                  valid_loss, valid_accs, valid_errs, valid_speed = self.run_epoch("epoch %i (validation)" % epoch,
                                                                                   self.valid_data, False)
                  accs_str = " ".join(["%i:%.5f" % (id, acc) for (id, acc) in zip(self.params['task_ids'], valid_accs)])
                  errs_str = " ".join(["%i:%.5f" % (id, err) for (id, err) in zip(self.params['task_ids'], valid_errs)])
                  print("\r\x1b[K Valid: loss: %.5f | acc: %s | error_ratio: %s | instances/sec: %.2f" % (valid_loss,
                                                                                                          accs_str,
                                                                                                          errs_str,
                                                                                                          valid_speed))
  
                  epoch_time = time.time() - total_time_start
                  log_entry = {
                      'epoch': epoch,
                      'time': epoch_time,
                      'train_results': (train_loss, train_accs.tolist(), train_errs.tolist(), train_speed),
                      'valid_results': (valid_loss, valid_accs.tolist(), valid_errs.tolist(), valid_speed),
                  }
                  log_to_save.append(log_entry)
                  with open(self.log_file, 'w') as f:
                      json.dump(log_to_save, f, indent=4)
  
                  val_acc = np.sum(valid_accs)  # type: float
                  if val_acc < best_val_acc:
                      self.save_model(self.best_model_file)
                      print("  (Best epoch so far, cum. val. acc decreased to %.5f from %.5f. Saving to '%s')" % (val_acc, best_val_acc, self.best_model_file))
                      best_val_acc = val_acc
                      best_val_acc_epoch = epoch
                  elif epoch - best_val_acc_epoch >= self.params['patience']:
                      print("Stopping training after %i epochs without improvement on validation accuracy." % self.params['patience'])
                      break
  
      def save_model(self, path):
          weights_to_save = {}
          for variable in self.sess.graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
              assert variable.name not in weights_to_save
              weights_to_save[variable.name] = self.sess.run(variable)
  
          data_to_save = {
                           "params": self.params,
                           "weights": weights_to_save
                         }
  
          with open(path, 'wb') as out_file:
              pickle.dump(data_to_save, out_file, pickle.HIGHEST_PROTOCOL)
  
      def initialize_model(self):
          init_op = tf.group(tf.global_variables_initializer(),
                             tf.local_variables_initializer())
          self.sess.run(init_op)
  
      def restore_model(self, path):
          print("Restoring weights from file %s." % path)
          with open(path, 'rb') as in_file:
              data_to_load = pickle.load(in_file)
  
          # Assert that we got the same model configuration
          assert len(self.params) == len(data_to_load['params'])
          for (par, par_value) in self.params.items():
              # Fine to have different task_ids:
              if par not in ['task_ids', 'num_epochs']:
                  assert par_value == data_to_load['params'][par]
  
          variables_to_initialize = []
          with tf.name_scope("restore"):
              restore_ops = []
              used_vars = set()
              for variable in self.sess.graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
                  used_vars.add(variable.name)
                  if variable.name in data_to_load['weights']:
                      restore_ops.append(variable.assign(data_to_load['weights'][variable.name]))
                  else:
                      print('Freshly initializing %s since no saved value was found.' % variable.name)
                      variables_to_initialize.append(variable)
              for var_name in data_to_load['weights']:
                  if var_name not in used_vars:
                      print('Saved weights for %s not used by model.' % var_name)
              restore_ops.append(tf.variables_initializer(variables_to_initialize))
              self.sess.run(restore_ops)
  
apis: tensorflow.reduce_sum tensorflow.variables_initializer tensorflow.train.AdamOptimizer tensorflow.Graph tensorflow.ConfigProto tensorflow.clip_by_norm tensorflow.name_scope tensorflow.Session tensorflow.square tensorflow.placeholder tensorflow.global_variables_initializer tensorflow.zeros_like tensorflow.set_random_seed numpy.array numpy.sum tensorflow.local_variables_initializer numpy.random.seed tensorflow.variable_scope tensorflow.abs
api_extract: [(50, 'os.path.join', 'os.path.join', (['log_dir', "('%s_log.json' % self.run_id)"], {}), False, 'import os\n'), (51, 'os.path.join', 'os.path.join', (['log_dir', "('%s_model_best.pickle' % self.run_id)"], {}), False, 'import os\n'), (66, 'random.seed', 'random.seed', (["params['random_seed']"], {}), False, 'import random\n'), (67, 'numpy.random.seed', 'np.random.seed', (["params['random_seed']"], {}), True, 'import numpy as np\n'), (78, 'tensorflow.ConfigProto', 'tf.ConfigProto', ([], {}), True, 'import tensorflow as tf\n'), (80, 'tensorflow.Graph', 'tf.Graph', ([], {}), True, 'import tensorflow as tf\n'), (81, 'tensorflow.Session', 'tf.Session', ([], {'graph': 'self.graph', 'config': 'config'}), True, 'import tensorflow as tf\n'), (98, 'os.path.join', 'os.path.join', (['self.data_dir', 'file_name'], {}), False, 'import os\n'), (131, 'tensorflow.placeholder', 'tf.placeholder', (['tf.int32', '[]'], {'name': '"""num_graphs"""'}), True, 'import tensorflow as tf\n'), (132, 'tensorflow.placeholder', 'tf.placeholder', (['tf.float32', '[]'], {'name': '"""out_layer_dropout_keep_prob"""'}), True, 'import tensorflow as tf\n'), (163, 'tensorflow.reduce_sum', 'tf.reduce_sum', (["self.ops['losses']"], {}), True, 'import tensorflow as tf\n'), (176, 'tensorflow.train.AdamOptimizer', 'tf.train.AdamOptimizer', (["self.params['learning_rate']"], {}), True, 'import tensorflow as tf\n'), (201, 'numpy.array', 'np.array', (['[0.066513725, 0.012235489, 0.071939046, 0.033730778, 0.033486113, \n    0.004278493, 0.001330901, 0.004165489, 0.004128926, 0.00409976, \n    0.004527465, 0.012292586, 0.037467458]'], {}), True, 'import numpy as np\n'), (208, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (238, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (64, 'json.dump', 'json.dump', (['params', 'f'], {}), False, 'import json\n'), (83, 'tensorflow.set_random_seed', 'tf.set_random_seed', (["params['random_seed']"], {}), True, 'import tensorflow as tf\n'), (102, 'json.load', 'json.load', (['f'], {}), False, 'import json\n'), (134, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""graph_model"""'], {}), True, 'import tensorflow as tf\n'), (186, 'tensorflow.local_variables_initializer', 'tf.local_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (230, 'numpy.sum', 'np.sum', (['accuracies'], {'axis': '(0)'}), True, 'import numpy as np\n'), (299, 'pickle.dump', 'pickle.dump', (['data_to_save', 'out_file', 'pickle.HIGHEST_PROTOCOL'], {}), False, 'import pickle\n'), (302, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (303, 'tensorflow.local_variables_initializer', 'tf.local_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (309, 'pickle.load', 'pickle.load', (['in_file'], {}), False, 'import pickle\n'), (319, 'tensorflow.name_scope', 'tf.name_scope', (['"""restore"""'], {}), True, 'import tensorflow as tf\n'), (48, 'time.strftime', 'time.strftime', (['"""%Y-%m-%d-%H-%M-%S"""'], {}), False, 'import time\n'), (61, 'json.loads', 'json.loads', (['config'], {}), False, 'import json\n'), (63, 'os.path.join', 'os.path.join', (['log_dir', "('%s_params.json' % self.run_id)"], {}), False, 'import os\n'), (140, 'tensorflow.zeros_like', 'tf.zeros_like', (["self.placeholders['initial_node_representation']"], {}), True, 'import tensorflow as tf\n'), (144, 'tensorflow.variable_scope', 'tf.variable_scope', (["('out_layer_task%i' % task_id)"], {}), True, 'import tensorflow as tf\n'), (233, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (242, 'numpy.sum', 'np.sum', (['valid_accs'], {}), True, 'import numpy as np\n'), (277, 'numpy.sum', 'np.sum', (['valid_accs'], {}), True, 'import numpy as np\n'), (332, 'tensorflow.variables_initializer', 'tf.variables_initializer', (['variables_to_initialize'], {}), True, 'import tensorflow as tf\n'), (48, 'os.getpid', 'os.getpid', ([], {}), False, 'import os\n'), (58, 'json.load', 'json.load', (['f'], {}), False, 'import json\n'), (65, 'json.dumps', 'json.dumps', (['self.params'], {}), False, 'import json\n'), (145, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""regression_gate"""'], {}), True, 'import tensorflow as tf\n'), (146, 'utils.MLP', 'MLP', (["(2 * self.params['hidden_size'])", '(1)', '[]', "self.placeholders['out_layer_dropout_keep_prob']"], {}), False, 'from utils import MLP, ThreadedIterator, SMALL_NUMBER\n'), (148, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""regression"""'], {}), True, 'import tensorflow as tf\n'), (149, 'utils.MLP', 'MLP', (["self.params['hidden_size']", '(1)', '[]', "self.placeholders['out_layer_dropout_keep_prob']"], {}), False, 'from utils import MLP, ThreadedIterator, SMALL_NUMBER\n'), (156, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['task_target_mask'], {}), True, 'import tensorflow as tf\n'), (223, 'numpy.array', 'np.array', (['batch_accuracies'], {}), True, 'import numpy as np\n'), (266, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (275, 'json.dump', 'json.dump', (['log_to_save', 'f'], {'indent': '(4)'}), False, 'import json\n'), (158, 'tensorflow.abs', 'tf.abs', (['diff'], {}), True, 'import tensorflow as tf\n'), (181, 'tensorflow.clip_by_norm', 'tf.clip_by_norm', (['grad', "self.params['clamp_gradient_norm']"], {}), True, 'import tensorflow as tf\n'), (159, 'tensorflow.square', 'tf.square', (['diff'], {}), True, 'import tensorflow as tf\n')]