Saving weights of epoch 1 at step 92

fa13e7b over 5 years ago

10.4 kB

import jax
print(jax.local_device_count())
import jax.numpy as jnp

import flax
import flax.linen as nn
from flax.training.common_utils import get_metrics,onehot,shard,shard_prng_key
from flax.training import train_state
from flax.metrics.tensorboard import SummaryWriter
from flax.training import checkpoints


import logging
import optax
import math
from tqdm import tqdm

from pathlib import Path
from typing import Callable
from itertools import chain
from flax.metrics import tensorboard

from datasets import load_dataset,load_metric
from transformers import GPT2Config,GPT2Tokenizer

from model_file  import FlaxGPT2ForMultipleChoice

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def main():


  tokenizer=GPT2Tokenizer.from_pretrained('gpt2',pad_token='<|endoftext|>')

  dataset=load_dataset('cosmos_qa')

  def preprocess(example):
    example['context&question']=example['context']+example['question']
    example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
    example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
    return example

  train_dataset=dataset['train'].map(preprocess)
  validation_dataset=dataset['validation'].map(preprocess)
  test_dataset=dataset['test'].map(preprocess)

  #Remove after experiment
  len_train_dataset=25262
  len_validation_dataset=2985
  len_test_dataset=6963

  train_dataset=train_dataset.select(range(len_train_dataset))
  test_dataset=test_dataset.select(range(len_test_dataset))
  validation_dataset=validation_dataset.select(range(len_validation_dataset))

  #remove_cols=train_dataset.column_names

  def tokenize(examples):
    a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
    a['labels']=examples['label']
    return a

  train_dataset=train_dataset.map(tokenize)
  validation_dataset=validation_dataset.map(tokenize)
  test_dataset=test_dataset.map(tokenize)

  remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']

  train_dataset=train_dataset.remove_columns(remov_col)
  validation_dataset=validation_dataset.remove_columns(remov_col)
  test_dataset=test_dataset.remove_columns(remov_col)

  per_device_batch_size=4
  seed=0
  num_train_epochs=5
  learning_rate=2e-5
  

  total_batch_size = per_device_batch_size * jax.local_device_count()
  print('The overall batch size (both for training and eval) is', total_batch_size)
  num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs
  num_validation_steps=len(validation_dataset)//total_batch_size*num_train_epochs

  learning_rate_function = optax.linear_schedule(init_value=learning_rate, end_value=0, transition_steps=num_train_steps)

  class TrainState(train_state.TrainState):
    logits_function:Callable=flax.struct.field(pytree_node=False)
    loss_function:Callable=flax.struct.field(pytree_node=False)

  def adamw(weight_decay):
    return optax.adamw(learning_rate=learning_rate_function,b1=0.9,b2=0.99,eps=1e-6,weight_decay=weight_decay)

  decay_path=lambda p:not any(x in p for x in ['bias','LayerNorm.weight'])

  def traverse(function):
    def mask(data):
      flat=flax.traverse_util.flatten_dict(data)
      return flax.traverse_util.unflatten_dict({k:function(k,v) for k,v in flat.items()})
    return mask

  gradient_transformation=optax.chain(
      optax.masked(adamw(0.0),mask=traverse(lambda path,_:decay_path(path))),
      optax.masked(adamw(0.01),mask=traverse(lambda path,_:not decay_path(path))))

  def loss_function(logits,labels):
    logits=flax.linen.log_softmax(logits)
    xentropy=optax.softmax_cross_entropy(logits,onehot(labels,num_classes=4))
    return jnp.mean(xentropy)

  def eval_function(logits):
    return logits.argmax(-1)

  model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-medium',input_shape=(1,4,1))

  state=TrainState.create(apply_fn=model.__call__,
                          params=model.params,
                          tx=gradient_transformation,
                          logits_function=eval_function,
                          loss_function=loss_function)
                      
  def train_step(state,batch,dropout_rng):
      targets=batch.pop("label")
      dropout_rng,new_dropout_rng=jax.random.split(dropout_rng)
      def loss_function(params):
          logits=state.apply_fn(**batch,params=params,dropout_rng=dropout_rng,train=True)[0]
          loss=state.loss_function(logits,targets)
          return loss
      grad_function=jax.value_and_grad(loss_function)
      loss,grad=grad_function(state.params)
      grad=jax.lax.pmean(grad,"batch")
      new_state=state.apply_gradients(grads=grad)
        #Added.
      logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
      accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
      metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
      return new_state,metrics,new_dropout_rng

  parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))

  def eval_step(state, batch):
      targets=batch.pop('label')
      logits = state.apply_fn(**batch, params=state.params, train=False)
      loss=state.loss_function(logits,targets)
      predictions=state.logits_function(logits)
      eval_accuracy=jnp.equal(predictions,targets)
      #eval_acc=jnp.equal(predictions,targets)
      metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
      #return state.logits_function(logits)  #(8,4)
      return targets,predictions,metrics

  parallel_eval_step = jax.pmap(eval_step, axis_name="batch")

  def glue_train_data_loader(rng,dataset,batch_size):
    steps_per_epoch=len_train_dataset//batch_size
    perms=jax.random.permutation(rng,len(dataset))
    perms=perms[:steps_per_epoch*batch_size]
    perms=perms.reshape((steps_per_epoch,batch_size))
    for perm in perms:
      batch=dataset[perm]
      batch={k:jnp.array(v) for k,v in batch.items()}
      batch=shard(batch)
      yield batch

  rng=jax.random.PRNGKey(seed)
  dropout_rngs=jax.random.split(rng,jax.local_device_count())

  def glue_eval_data_loader(dataset, batch_size):
      for i in range(len_validation_dataset // batch_size):
          batch = dataset[i * batch_size : (i + 1) * batch_size]
          batch = {k: jnp.array(v) for k, v in batch.items()}
          batch = shard(batch)

          yield batch

  state = flax.jax_utils.replicate(state)
  #metrics_list = list_metrics()

  actual_task = "mnli"
  metric = load_metric('glue', "mnli")
  actual_taskmetric = load_metric('glue', actual_task)

  workdir='../results_tensorboard'
  summary_writer = tensorboard.SummaryWriter(workdir)
  #summary_writer.hparams(dict(GPT2Config()))

  logger.info(f"***** Running training *****")
  logger.info(f"  Num examples = {len_train_dataset}")
  logger.info(f"  Num Epochs = {num_train_epochs}")
  logger.info(f"  Instantaneous batch size per device = {per_device_batch_size}")
  logger.info(f"  Total train batch size  = {total_batch_size}")
  logger.info(f"  Total optimization steps = {num_train_steps}")

  for i, epoch in enumerate(tqdm(range(1, num_train_epochs+1), desc=f"Epoch ...", position=0, leave=True)):
    rng, input_rng = jax.random.split(rng)
    train_acc_metrics=[]
    train_loss_metrics=[]
    eval_acc_metrics=[]
    eval_loss_metrics=[]
    # train
    with tqdm(total=len_train_dataset // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
      for idx,batch in enumerate(glue_train_data_loader(input_rng, train_dataset, total_batch_size)):
        state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
        train_acc_metrics.append(jax.device_get(train_metric['accuracy']).mean().item())
        train_loss_metrics.append(flax.jax_utils.unreplicate(train_metric)['loss'].item())
        if idx%5==0:
          summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
          summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
        if idx%50==0:
            logger.info(f"train_step_loss{idx}:   {flax.jax_utils.unreplicate(train_metric)['loss'].item()} train_step_acc{idx}:   {jax.device_get(train_metric['accuracy']).mean().item()}")
            
        progress_bar_train.update(1)

    # evaluate
    with tqdm(total=len_validation_dataset // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
      for idx,batch in enumerate(glue_eval_data_loader(validation_dataset, total_batch_size)):
          labels,predictions,eval_metric=parallel_eval_step(state, batch)
          eval_acc_metrics.append(jax.device_get(eval_metric['accuracy']).mean().item())
          eval_loss_metrics.append(flax.jax_utils.unreplicate(eval_metric)['loss'].item())
          progress_bar_eval.update(1)
          if idx%5==0:
            logger.info(f"eval_step_loss  {idx} :   {flax.jax_utils.unreplicate(eval_metric)['loss'].item()} eval_step_acc  {idx}  :   {jax.device_get(eval_metric['accuracy']).mean().item()}")
            summary_writer.scalar('eval_loss : ', flax.jax_utils.unreplicate(eval_metric)['loss'].item(),idx)
            summary_writer.scalar('eval_accuracy : ', jax.device_get(eval_metric['accuracy']).mean().item(),idx)

    if jax.process_index() == 0:
        params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
    
        model.save_pretrained(
            '../',
            params=params,
            push_to_hub=True,
            commit_message=f"Saving weights of epoch {epoch} at step {idx}",)
          
    logger.info(f"---------------------Epoch {epoch} done-----------------")
    logger.info(f"Train loss:   {jax.device_get(jnp.array(train_loss_metrics)).mean().item()} Train accuracy:  {jax.device_get(jnp.array(train_acc_metrics)).mean().item()}")
    logger.info(f"Eval loss:    {jax.device_get(jnp.array(eval_loss_metrics)).mean().item()} Eval accuracy:    {jax.device_get(jnp.array(eval_acc_metrics)).mean().item()}")
  summary_writer.flush()

if __name__ == "__main__":
    main()

	import jax
	print(jax.local_device_count())
	import jax.numpy as jnp

	import flax
	import flax.linen as nn
	from flax.training.common_utils import get_metrics,onehot,shard,shard_prng_key
	from flax.training import train_state
	from flax.metrics.tensorboard import SummaryWriter
	from flax.training import checkpoints


	import logging
	import optax
	import math
	from tqdm import tqdm

	from pathlib import Path
	from typing import Callable
	from itertools import chain
	from flax.metrics import tensorboard

	from datasets import load_dataset,load_metric
	from transformers import GPT2Config,GPT2Tokenizer

	from model_file import FlaxGPT2ForMultipleChoice

	logger = logging.getLogger()
	logger.setLevel(logging.INFO)

	def main():


	tokenizer=GPT2Tokenizer.from_pretrained('gpt2',pad_token='<\|endoftext\|>')

	dataset=load_dataset('cosmos_qa')

	def preprocess(example):
	example['context&question']=example['context']+example['question']
	example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
	example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
	return example

	train_dataset=dataset['train'].map(preprocess)
	validation_dataset=dataset['validation'].map(preprocess)
	test_dataset=dataset['test'].map(preprocess)

	#Remove after experiment
	len_train_dataset=25262
	len_validation_dataset=2985
	len_test_dataset=6963

	train_dataset=train_dataset.select(range(len_train_dataset))
	test_dataset=test_dataset.select(range(len_test_dataset))
	validation_dataset=validation_dataset.select(range(len_validation_dataset))

	#remove_cols=train_dataset.column_names

	def tokenize(examples):
	a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
	a['labels']=examples['label']
	return a

	train_dataset=train_dataset.map(tokenize)
	validation_dataset=validation_dataset.map(tokenize)
	test_dataset=test_dataset.map(tokenize)

	remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']

	train_dataset=train_dataset.remove_columns(remov_col)
	validation_dataset=validation_dataset.remove_columns(remov_col)
	test_dataset=test_dataset.remove_columns(remov_col)

	per_device_batch_size=4
	seed=0
	num_train_epochs=5
	learning_rate=2e-5


	total_batch_size = per_device_batch_size * jax.local_device_count()
	print('The overall batch size (both for training and eval) is', total_batch_size)
	num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs
	num_validation_steps=len(validation_dataset)//total_batch_size*num_train_epochs

	learning_rate_function = optax.linear_schedule(init_value=learning_rate, end_value=0, transition_steps=num_train_steps)

	class TrainState(train_state.TrainState):
	logits_function:Callable=flax.struct.field(pytree_node=False)
	loss_function:Callable=flax.struct.field(pytree_node=False)

	def adamw(weight_decay):
	return optax.adamw(learning_rate=learning_rate_function,b1=0.9,b2=0.99,eps=1e-6,weight_decay=weight_decay)

	decay_path=lambda p:not any(x in p for x in ['bias','LayerNorm.weight'])

	def traverse(function):
	def mask(data):
	flat=flax.traverse_util.flatten_dict(data)
	return flax.traverse_util.unflatten_dict({k:function(k,v) for k,v in flat.items()})
	return mask

	gradient_transformation=optax.chain(
	optax.masked(adamw(0.0),mask=traverse(lambda path,_:decay_path(path))),
	optax.masked(adamw(0.01),mask=traverse(lambda path,_:not decay_path(path))))

	def loss_function(logits,labels):
	logits=flax.linen.log_softmax(logits)
	xentropy=optax.softmax_cross_entropy(logits,onehot(labels,num_classes=4))
	return jnp.mean(xentropy)

	def eval_function(logits):
	return logits.argmax(-1)

	model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-medium',input_shape=(1,4,1))

	state=TrainState.create(apply_fn=model.__call__,
	params=model.params,
	tx=gradient_transformation,
	logits_function=eval_function,
	loss_function=loss_function)

	def train_step(state,batch,dropout_rng):
	targets=batch.pop("label")
	dropout_rng,new_dropout_rng=jax.random.split(dropout_rng)
	def loss_function(params):
	logits=state.apply_fn(**batch,params=params,dropout_rng=dropout_rng,train=True)[0]
	loss=state.loss_function(logits,targets)
	return loss
	grad_function=jax.value_and_grad(loss_function)
	loss,grad=grad_function(state.params)
	grad=jax.lax.pmean(grad,"batch")
	new_state=state.apply_gradients(grads=grad)
	#Added.
	logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
	accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
	metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
	return new_state,metrics,new_dropout_rng

	parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))

	def eval_step(state, batch):
	targets=batch.pop('label')
	logits = state.apply_fn(**batch, params=state.params, train=False)
	loss=state.loss_function(logits,targets)
	predictions=state.logits_function(logits)
	eval_accuracy=jnp.equal(predictions,targets)
	#eval_acc=jnp.equal(predictions,targets)
	metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
	#return state.logits_function(logits) #(8,4)
	return targets,predictions,metrics

	parallel_eval_step = jax.pmap(eval_step, axis_name="batch")

	def glue_train_data_loader(rng,dataset,batch_size):
	steps_per_epoch=len_train_dataset//batch_size
	perms=jax.random.permutation(rng,len(dataset))
	perms=perms[:steps_per_epoch*batch_size]
	perms=perms.reshape((steps_per_epoch,batch_size))
	for perm in perms:
	batch=dataset[perm]
	batch={k:jnp.array(v) for k,v in batch.items()}
	batch=shard(batch)
	yield batch

	rng=jax.random.PRNGKey(seed)
	dropout_rngs=jax.random.split(rng,jax.local_device_count())

	def glue_eval_data_loader(dataset, batch_size):
	for i in range(len_validation_dataset // batch_size):
	batch = dataset[i * batch_size : (i + 1) * batch_size]
	batch = {k: jnp.array(v) for k, v in batch.items()}
	batch = shard(batch)

	yield batch

	state = flax.jax_utils.replicate(state)
	#metrics_list = list_metrics()

	actual_task = "mnli"
	metric = load_metric('glue', "mnli")
	actual_taskmetric = load_metric('glue', actual_task)

	workdir='../results_tensorboard'
	summary_writer = tensorboard.SummaryWriter(workdir)
	#summary_writer.hparams(dict(GPT2Config()))

	logger.info(f"*** Running training ***")
	logger.info(f" Num examples = {len_train_dataset}")
	logger.info(f" Num Epochs = {num_train_epochs}")
	logger.info(f" Instantaneous batch size per device = {per_device_batch_size}")
	logger.info(f" Total train batch size = {total_batch_size}")
	logger.info(f" Total optimization steps = {num_train_steps}")

	for i, epoch in enumerate(tqdm(range(1, num_train_epochs+1), desc=f"Epoch ...", position=0, leave=True)):
	rng, input_rng = jax.random.split(rng)
	train_acc_metrics=[]
	train_loss_metrics=[]
	eval_acc_metrics=[]
	eval_loss_metrics=[]
	# train
	with tqdm(total=len_train_dataset // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
	for idx,batch in enumerate(glue_train_data_loader(input_rng, train_dataset, total_batch_size)):
	state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
	train_acc_metrics.append(jax.device_get(train_metric['accuracy']).mean().item())
	train_loss_metrics.append(flax.jax_utils.unreplicate(train_metric)['loss'].item())
	if idx%5==0:
	summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
	summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
	if idx%50==0:
	logger.info(f"train_step_loss{idx}: {flax.jax_utils.unreplicate(train_metric)['loss'].item()} train_step_acc{idx}: {jax.device_get(train_metric['accuracy']).mean().item()}")

	progress_bar_train.update(1)

	# evaluate
	with tqdm(total=len_validation_dataset // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
	for idx,batch in enumerate(glue_eval_data_loader(validation_dataset, total_batch_size)):
	labels,predictions,eval_metric=parallel_eval_step(state, batch)
	eval_acc_metrics.append(jax.device_get(eval_metric['accuracy']).mean().item())
	eval_loss_metrics.append(flax.jax_utils.unreplicate(eval_metric)['loss'].item())
	progress_bar_eval.update(1)
	if idx%5==0:
	logger.info(f"eval_step_loss {idx} : {flax.jax_utils.unreplicate(eval_metric)['loss'].item()} eval_step_acc {idx} : {jax.device_get(eval_metric['accuracy']).mean().item()}")
	summary_writer.scalar('eval_loss : ', flax.jax_utils.unreplicate(eval_metric)['loss'].item(),idx)
	summary_writer.scalar('eval_accuracy : ', jax.device_get(eval_metric['accuracy']).mean().item(),idx)

	if jax.process_index() == 0:
	params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))

	model.save_pretrained(
	'../',
	params=params,
	push_to_hub=True,
	commit_message=f"Saving weights of epoch {epoch} at step {idx}",)

	logger.info(f"---------------------Epoch {epoch} done-----------------")
	logger.info(f"Train loss: {jax.device_get(jnp.array(train_loss_metrics)).mean().item()} Train accuracy: {jax.device_get(jnp.array(train_acc_metrics)).mean().item()}")
	logger.info(f"Eval loss: {jax.device_get(jnp.array(eval_loss_metrics)).mean().item()} Eval accuracy: {jax.device_get(jnp.array(eval_acc_metrics)).mean().item()}")
	summary_writer.flush()

	if __name__ == "__main__":
	main()