<jupyter_start><jupyter_text>Import<jupyter_code>import sonnet as snt
from tqdm import tqdm
from IPython.display import clear_output
import numpy as np
import pandas as pd
import time
import os
import tensorflow as tf
import tensorflow_hub as hub
from datetime import datetime<jupyter_output>2021-12-12 05:57:57.481105: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0<jupyter_text>Code<jupyter_code>import tfenformer #the output of this enformer has been changed
# @title `get_dataset(organism, subset, num_threads=8)`
import glob
import json
import functools


def organism_path(organism):
    return os.path.join('/work/long_lab/qli/Enformer_TFs/', organism)


def get_dataset(organism, subset, num_threads=16):
    metadata = get_metadata(organism)
    dataset = tf.data.TFRecordDataset(tfrecord_files(organism, subset),
                                    compression_type='ZLIB',
                                    num_parallel_reads=num_threads)
    dataset = dataset.map(functools.partial(deserialize, metadata=metadata),
                        num_parallel_calls=num_threads)
    return dataset


def get_metadata(organism):
    # Keys:
    # num_targets, train_seqs, valid_seqs, test_seqs, seq_length,
    # pool_width, crop_bp, target_length
    path = os.path.join(organism_path(organism), 'statistics.json')
    with tf.io.gfile.GFile(path, 'r') as f:
        return json.load(f)


def tfrecord_files(organism, subset):
  # Sort the values by int(*).
  return sorted(tf.io.gfile.glob(os.path.join(
      organism_path(organism), 'tfrecords', f'{subset}-*.tfr'
  )), key=lambda x: int(x.split('-')[-1].split('.')[0]))


def deserialize(serialized_example, metadata):
    """Deserialize bytes stored in TFRecordFile."""
    feature_map = {
      'sequence': tf.io.FixedLenFeature([], tf.string),
      'target': tf.io.FixedLenFeature([], tf.string),
    }
    example = tf.io.parse_example(serialized_example, feature_map)
    sequence = tf.io.decode_raw(example['sequence'], tf.bool)
    sequence = tf.reshape(sequence, (metadata['seq_length'], 4))
    sequence = tf.cast(sequence, tf.float32)

    target = tf.io.decode_raw(example['target'], tf.float16)
    target = tf.reshape(target,
                      (metadata['target_length'], metadata['num_targets']))
    target = tf.cast(target, tf.float32)

    return {'sequence': sequence,
          'target': target}<jupyter_output><empty_output><jupyter_text>1. Load dataset<jupyter_code>human_dataset = get_dataset('human', 'train').batch(1).repeat()
#mouse_dataset = get_dataset('mouse', 'train').batch(1).repeat()
human_dataset_pre = human_dataset.prefetch(2)<jupyter_output>2021-12-12 05:58:08.885151: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-12-12 05:58:08.886839: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1
2021-12-12 05:58:10.468096: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 0 with properties: 
pciBusID: 0000:3b:00.0 name: Tesla V100-PCIE-16GB computeCapability: 7.0
coreClock: 1.38GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2021-12-12 05:58:10.468744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 1 with properties: 
pciBusID: 0000:d8:00.0 name: Tesla V100-PCIE-16GB computeCapability: 7.0
coreClock: 1.38GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2021-12-12 05:58:10.468770: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
2021-[...]<jupyter_text>2. Model training<jupyter_code>def create_step_function(model, optimizer):

    @tf.function
    def train_step(batch, head, optimizer_clip_norm_global=0.2):
        with tf.GradientTape() as tape:
          outputs = model(batch['sequence'], is_training=True)[head]
          loss = tf.reduce_mean(
              tf.keras.losses.poisson(batch['target'], outputs))

    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply(gradients, model.trainable_variables)

    return loss
  return train_step
learning_rate = tf.Variable(0., trainable=False, name='learning_rate')
optimizer = snt.optimizers.Adam(learning_rate=learning_rate)
num_warmup_steps = 0
target_learning_rate = 0.0005

model = tfenformer.Enformer(channels=1536,  # Use 4x fewer channels to train faster.
                          num_heads=8,
                          num_transformer_layers=11,
                          pooling_type='attention')

train_step = create_step_function(model, optimizer)
### Train the model (Do only if you need to start from scratch)

checkpoint_root = "TFEnformer_Transfer_Train_all_trainable_variables/checkpoints"
checkpoint_name = "example"
save_prefix = os.path.join(checkpoint_root, checkpoint_name)
checkpoint = tf.train.Checkpoint(module=model)
latest = tf.train.latest_checkpoint(checkpoint_root)
if latest is not None:
    checkpoint.restore(latest)

steps_per_epoch = 20
start_epoch = 0
num_epochs = 200

data_it = iter(human_dataset_pre)
global_step = 0
for epoch_i in range(start_epoch, num_epochs):
    loss_log=open("TFEnformer_Transfer_Train_all_trainable_variables/loss_log_human","a")
    for i in tqdm(range(steps_per_epoch)):
        global_step += 1

        if global_step > 1:
            learning_rate_frac = tf.math.minimum(1.0, global_step / tf.math.maximum(1.0, num_warmup_steps))      
            learning_rate.assign(target_learning_rate * learning_rate_frac)

        batch_human = next(data_it)

        loss_human = train_step(batch=batch_human, head='human')
        #End of the step
        print('global_step',global_step)

    # End of epoch.
    print('epoch_i',epoch_i,'loss_human', loss_human.numpy(),'learning_rate', optimizer.learning_rate.numpy())
    now = datetime.now()
    date_time = now.strftime("%m/%d/%Y, %H:%M:%S")
    loss_log.write("date and time: "+ date_time+", epoch: "+str(epoch_i)+", global_step: "+str(global_step)+", loss_human: "+str(loss_human.numpy())+', learning_rate:'+str(optimizer.learning_rate.numpy())+', learning_rate: '+str(optimizer.learning_rate.numpy())+"\n")
    loss_log.close()
    if epoch_i and not epoch_i % 2:
        checkpoint.save(save_prefix)
checkpoint.save(save_prefix)<jupyter_output><empty_output><jupyter_text>3. Model evaluating<jupyter_code># @title `PearsonR` and `R2` metrics

def _reduced_shape(shape, axis):
    if axis is None:
        return tf.TensorShape([])
    return tf.TensorShape([d for i, d in enumerate(shape) if i not in axis])


class CorrelationStats(tf.keras.metrics.Metric):
  """Contains shared code for PearsonR and R2."""

  def __init__(self, reduce_axis=None, name='pearsonr'):
    """Pearson correlation coefficient.

    Args:
      reduce_axis: Specifies over which axis to compute the correlation (say
        (0, 1). If not specified, it will compute the correlation across the
        whole tensor.
      name: Metric name.
    """
    super(CorrelationStats, self).__init__(name=name)
    self._reduce_axis = reduce_axis
    self._shape = None  # Specified in _initialize.

    def _initialize(self, input_shape):
        # Remaining dimensions after reducing over self._reduce_axis.
        self._shape = _reduced_shape(input_shape, self._reduce_axis)

        weight_kwargs = dict(shape=self._shape, initializer='zeros')
        self._count = self.add_weight(name='count', **weight_kwargs)
        self._product_sum = self.add_weight(name='product_sum', **weight_kwargs)
        self._true_sum = self.add_weight(name='true_sum', **weight_kwargs)
        self._true_squared_sum = self.add_weight(name='true_squared_sum',
                                                 **weight_kwargs)
        self._pred_sum = self.add_weight(name='pred_sum', **weight_kwargs)
        self._pred_squared_sum = self.add_weight(name='pred_squared_sum',
                                                 **weight_kwargs)

    def update_state(self, y_true, y_pred, sample_weight=None):
        """Update the metric state.

        Args:
          y_true: Multi-dimensional float tensor [batch, ...] containing the ground
            truth values.
          y_pred: float tensor with the same shape as y_true containing predicted
            values.
          sample_weight: 1D tensor aligned with y_true batch dimension specifying
            the weight of individual observations.
        """
        if self._shape is None:
          # Explicit initialization check.
          self._initialize(y_true.shape)
        y_true.shape.assert_is_compatible_with(y_pred.shape)
        y_true = tf.cast(y_true, 'float32')
        y_pred = tf.cast(y_pred, 'float32')

        self._product_sum.assign_add(
            tf.reduce_sum(y_true * y_pred, axis=self._reduce_axis))

        self._true_sum.assign_add(
            tf.reduce_sum(y_true, axis=self._reduce_axis))

        self._true_squared_sum.assign_add(
            tf.reduce_sum(tf.math.square(y_true), axis=self._reduce_axis))

        self._pred_sum.assign_add(
            tf.reduce_sum(y_pred, axis=self._reduce_axis))

        self._pred_squared_sum.assign_add(
            tf.reduce_sum(tf.math.square(y_pred), axis=self._reduce_axis))

        self._count.assign_add(
            tf.reduce_sum(tf.ones_like(y_true), axis=self._reduce_axis))

    def result(self):
        raise NotImplementedError('Must be implemented in subclasses.')

    def reset_states(self):
        if self._shape is not None:
              tf.keras.backend.batch_set_value([(v, np.zeros(self._shape))
                                            for v in self.variables])


class PearsonR(CorrelationStats):
  """Pearson correlation coefficient.

  Computed as:
  ((x - x_avg) * (y - y_avg) / sqrt(Var[x] * Var[y])
  """

    def __init__(self, reduce_axis=(0,), name='pearsonr'):
        """Pearson correlation coefficient.

        Args:
          reduce_axis: Specifies over which axis to compute the correlation.
          name: Metric name.
        """
        super(PearsonR, self).__init__(reduce_axis=reduce_axis,
                                       name=name)

    def result(self):
        true_mean = self._true_sum / self._count
        pred_mean = self._pred_sum / self._count

        covariance = (self._product_sum
                      - true_mean * self._pred_sum
                      - pred_mean * self._true_sum
                      + self._count * true_mean * pred_mean)

        true_var = self._true_squared_sum - self._count * tf.math.square(true_mean)
        pred_var = self._pred_squared_sum - self._count * tf.math.square(pred_mean)
        tp_var = tf.math.sqrt(true_var) * tf.math.sqrt(pred_var)
        correlation = covariance / tp_var

        return correlation


class R2(CorrelationStats):
  """R-squared  (fraction of explained variance)."""

  def __init__(self, reduce_axis=None, name='R2'):
    """R-squared metric.

    Args:
      reduce_axis: Specifies over which axis to compute the correlation.
      name: Metric name.
    """
    super(R2, self).__init__(reduce_axis=reduce_axis,
                             name=name)

    def result(self):
        true_mean = self._true_sum / self._count
        total = self._true_squared_sum - self._count * tf.math.square(true_mean)
        residuals = (self._pred_squared_sum - 2 * self._product_sum
                     + self._true_squared_sum)

    return tf.ones_like(residuals) - residuals / total


class MetricDict:
    def __init__(self, metrics):
        self._metrics = metrics

    def update_state(self, y_true, y_pred):
        for k, metric in self._metrics.items():
            metric.update_state(y_true, y_pred)
    
    def result(self):
        return {k: metric.result() for k, metric in self._metrics.items()}
def evaluate_model(model, dataset, head, max_steps=None):
    metric = MetricDict({'PearsonR': PearsonR(reduce_axis=(0,1))})
    @tf.function
    def predict(x):
        return model(x, is_training=False)[head]

    for i, batch in tqdm(enumerate(dataset)):
        if max_steps is not None and i > max_steps:
              break
        metric.update_state(batch['target'], predict(batch['sequence']))

    return metric.result()
metrics_human = evaluate_model(model,
                               dataset=get_dataset('human', 'valid').batch(1).prefetch(2),
                               head='human',
                               max_steps=100)
print('')
print({k: v.numpy().mean() for k, v in metrics_human.items()})<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11200568316362776, 'fraction_numerical': 0.021864393401215567, 'mean_word_length': 3.213501829065514, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '23234', 'license': 'MIT', 'n_tokens_mistral': 4232, 'n_tokens_neox': 3959, 'n_words': 814, 'path': '0_TFEnformer_Transfer_Train_all_trainable_variables.ipynb', 'repo_name': 'liqingbioinfo/TFEnformer', 'size': '12669'}