diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 5abddf8d8d..e6b68b3d37 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -49,6 +49,10 @@ class ORTModule(torch.nn.Module): # Log level self._loglevel = getattr(logging, 'WARNING') + # TODO: debug flags + self._save_onnx = False + self._save_onnx_prefix = '' + def forward(self, *inputs, **kwargs): '''Forward pass starts here and continues at `_ORTModuleFunction.forward` @@ -65,6 +69,13 @@ class ORTModule(torch.nn.Module): self._onnx_training = ORTModule._get_forward_graph(self._original_module, *inputs, **kwargs) self._onnx_gradient = ORTModule._build_gradient_graph(self._onnx_training, self._grad_builder_config) self._onnx_forward, self._onnx_backward = ORTModule._split_forward_and_backward(self._onnx_gradient, self._grad_builder_config.weight_names_to_train) + + if self._save_onnx: + onnx.save(self._onnx_training, self._save_onnx_prefix + '_full_training.onnx') + onnx.save(self._onnx_gradient, self._save_onnx_prefix + '_with_grad.onnx') + onnx.save(self._onnx_forward, self._save_onnx_prefix + '_forward.onnx') + onnx.save(self._onnx_backward, self._save_onnx_prefix + '_backward.onnx') + # TODO: hard-coding to CPU only self._forward_session = onnxruntime.InferenceSession(self._onnx_forward.SerializeToString(), providers=['CPUExecutionProvider']) self._backward_session = onnxruntime.InferenceSession(self._onnx_backward.SerializeToString(), providers=['CPUExecutionProvider']) @@ -319,7 +330,6 @@ class ORTModule(torch.nn.Module): def _get_forward_graph(module, *inputs, **kwargs): '''Exports PyTorch `module` to ONNX with training flag, using `*inputs` as input - TODO: Support contrib OPs support? user model has no hint TODO: How to support dynamic axes? Dimensions are determined by samples TODO: How to ingest **kwargs in proper order during export? ''' @@ -329,6 +339,10 @@ class ORTModule(torch.nn.Module): # Deepcopy inputs, since input values may change after model run. sample_inputs_copy = copy.deepcopy(inputs) + # TODO: Support contrib OPs support? user model has no hint + # from onnxruntime.training import register_custom_ops_pytorch_exporter + # register_custom_ops_pytorch_exporter.register_custom_op() + # Export torch.nn.Module to ONNX torch.onnx.export(module, tuple(sample_inputs_copy), diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py new file mode 100644 index 0000000000..2bb30604bc --- /dev/null +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py @@ -0,0 +1,364 @@ +import torch +import wget +import os +import pandas as pd +import zipfile +from transformers import BertTokenizer +from keras.preprocessing.sequence import pad_sequences +from sklearn.model_selection import train_test_split +from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler +from transformers import BertForSequenceClassification, AdamW, BertConfig +from transformers import get_linear_schedule_with_warmup +import numpy as np +import random +import time +import datetime + + +import onnxruntime +from onnxruntime.training import ORTModule + +# 1. Device setup +# TODO: Hard-coding for CPU for ORTModule + +# if torch.cuda.is_available(): +# device = torch.device("cuda") +# print('There are %d GPU(s) available.' % torch.cuda.device_count()) +# print('We will use the GPU:', torch.cuda.get_device_name(0)) +# else: +# print('No GPU available, using the CPU instead.') +# device = torch.device("cpu") +device = torch.device("cpu") + +# 2. Loading CoLA Dataset +print('Downloading dataset...') +url = 'https://nyu-mll.github.io/CoLA/cola_public_1.1.zip' + +# Download the file (if we haven't already) +if not os.path.exists('./cola_public_1.1.zip'): + wget.download(url, './cola_public_1.1.zip') +else: + print('Reusing cached dataset') + +# Unzip it +if not os.path.exists('./cola_public'): + print('Extracting cached dataset') + with zipfile.ZipFile('./cola_public_1.1.zip', 'r') as zip_ref: + zip_ref.extractall('./') +else: + print('Reusing extracted dataset') + +# Load the dataset into a pandas dataframe. +df = pd.read_csv("./cola_public/raw/in_domain_train.tsv", delimiter='\t', header=None, names=['sentence_source', 'label', 'label_notes', 'sentence']) + +# Get the lists of sentences and their labels. +sentences = df.sentence.values +labels = df.label.values + +# 3. Tokenization & Input Formatting + +# Load the BERT tokenizer. +tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True) + +# Tokenize all of the sentences and map the tokens to their word IDs. +input_ids = [] +for sent in sentences: + # `encode` will: + # (1) Tokenize the sentence. + # (2) Prepend the `[CLS]` token to the start. + # (3) Append the `[SEP]` token to the end. + # (4) Map tokens to their IDs. + encoded_sent = tokenizer.encode( + sent, # Sentence to encode. + add_special_tokens = True, # Add '[CLS]' and '[SEP]' + ) + + # Add the encoded sentence to the list. + input_ids.append(encoded_sent) + +# We'll borrow the `pad_sequences` utility function to do this. +# 64 is slightly larger than the maximum training sentence length of 47... +MAX_LEN = 64 + +# Pad our input tokens with value 0. +input_ids = pad_sequences(input_ids, maxlen=MAX_LEN, dtype="long", + value=0, truncating="post", padding="post") + +# Create attention masks +attention_masks = [] + +# For each sentence... +for sent in input_ids: + # Create the attention mask. + # - If a token ID is 0, then it's padding, set the mask to 0. + # - If a token ID is > 0, then it's a real token, set the mask to 1. + att_mask = [int(token_id > 0) for token_id in sent] + + # Store the attention mask for this sentence. + attention_masks.append(att_mask) + +# Use 90% for training and 10% for validation. +train_inputs, validation_inputs, train_labels, validation_labels = train_test_split(input_ids, labels, + random_state=2018, test_size=0.1) +# Do the same for the masks. +train_masks, validation_masks, _, _ = train_test_split(attention_masks, labels, + random_state=2018, test_size=0.1) + +# Convert all inputs and labels into torch tensors, the required datatype +# for our model. +train_inputs = torch.tensor(train_inputs) +validation_inputs = torch.tensor(validation_inputs) + +train_labels = torch.tensor(train_labels) +validation_labels = torch.tensor(validation_labels) + +train_masks = torch.tensor(train_masks) +validation_masks = torch.tensor(validation_masks) + +# The DataLoader needs to know our batch size for training, so we specify it +batch_size = 32 + +# Create the DataLoader for our training set. +train_data = TensorDataset(train_inputs, train_masks, train_labels) +train_sampler = RandomSampler(train_data) +train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size) + +# Create the DataLoader for our validation set. +validation_data = TensorDataset(validation_inputs, validation_masks, validation_labels) +validation_sampler = SequentialSampler(validation_data) +validation_dataloader = DataLoader(validation_data, sampler=validation_sampler, batch_size=batch_size) + +# 4. Train (fine-tune) + +# TODO: try to use a smaller model + +# Load BertForSequenceClassification, the pretrained BERT model with a single +# linear classification layer on top. +model = BertForSequenceClassification.from_pretrained( + "bert-base-uncased", # Use the 12-layer BERT model, with an uncased vocab. + # '/home/thiagofc/dev/github/onnxruntime/orttraining/orttraining/test/python/', + num_labels = 2, # The number of output labels--2 for binary classification. + output_attentions = False, # Whether the model returns attentions weights. + output_hidden_states = False, # Whether the model returns all hidden-states. +) +model = ORTModule(model) + +# TODO: change it to False to stop saving ONNX models +model._save_onnx = True +model._save_onnx_prefix = 'BertForSequenceClassification' + +# Tell pytorch to run this model on the GPU. +# TODO: Hard coding it to CPU for ORTModule +# model.cuda() + +# Note: AdamW is a class from the huggingface library (as opposed to pytorch) +optimizer = AdamW(model.parameters(), + lr = 2e-5, # args.learning_rate - default is 5e-5, our notebook had 2e-5 + eps = 1e-8 # args.adam_epsilon - default is 1e-8. + ) + +# Number of training epochs (authors recommend between 2 and 4) +epochs = 4 + +# Total number of training steps is number of batches * number of epochs. +total_steps = len(train_dataloader) * epochs + +# Create the learning rate scheduler. +scheduler = get_linear_schedule_with_warmup(optimizer, + num_warmup_steps = 0, # Default value in run_glue.py + num_training_steps = total_steps) + +# Function to calculate the accuracy of our predictions vs labels +def flat_accuracy(preds, labels): + pred_flat = np.argmax(preds, axis=1).flatten() + labels_flat = labels.flatten() + return np.sum(pred_flat == labels_flat) / len(labels_flat) + +def format_time(elapsed): + '''Takes a time in seconds and returns a string hh:mm:ss''' + # Round to the nearest second. + elapsed_rounded = int(round((elapsed))) + + # Format as hh:mm:ss + return str(datetime.timedelta(seconds=elapsed_rounded)) + +# This training code is based on the `run_glue.py` script here: +# https://github.com/huggingface/transformers/blob/5bfcd0485ece086ebcbed2d008813037968a9e58/examples/run_glue.py#L128 + + +# Set the seed value all over the place to make this reproducible. +seed_val = 42 + +random.seed(seed_val) +np.random.seed(seed_val) +torch.manual_seed(seed_val) +onnxruntime.set_seed(seed_val) +# TODO: We are not using CUDA for ORTModule just yet +# torch.cuda.manual_seed_all(seed_val) + +# Store the average loss after each epoch so we can plot them. +loss_values = [] + +# ======================================== +# Training +# ======================================== +for epoch_i in range(0, epochs): + # Perform one full pass over the training set. + print("") + print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs)) + + # Measure how long the training epoch takes. + t0 = time.time() + + # Reset the total loss for this epoch. + total_loss = 0 + + # Put the model into training mode. Don't be mislead--the call to + # `train` just changes the *mode*, it doesn't *perform* the training. + # `dropout` and `batchnorm` layers behave differently during training + # vs. test (source: https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch) + model.train() + + # For each batch of training data... + for step, batch in enumerate(train_dataloader): + + # Progress update every 40 batches. + if step % 40 == 0 and not step == 0: + # Calculate elapsed time in minutes. + elapsed = format_time(time.time() - t0) + + # Report progress. + print(' Batch {:>5,} of {:>5,}. Elapsed: {:}.'.format(step, len(train_dataloader), elapsed)) + + # Unpack this training batch from our dataloader. + # + # As we unpack the batch, we'll also copy each tensor to the GPU using the + # `to` method. + # + # `batch` contains three pytorch tensors: + # [0]: input ids + # [1]: attention masks + # [2]: labels + b_input_ids = batch[0].to(device) + b_input_mask = batch[1].to(device) + b_labels = batch[2].to(device) + + # Always clear any previously calculated gradients before performing a + # backward pass. PyTorch doesn't do this automatically because + # accumulating the gradients is "convenient while training RNNs". + # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch) + model.zero_grad() + + # Perform a forward pass (evaluate the model on this training batch). + # This will return the loss (rather than the model output) because we + # have provided the `labels`. + # The documentation for this `model` function is here: + # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification + outputs = model(b_input_ids, + token_type_ids=None, + attention_mask=b_input_mask, + labels=b_labels) + + # The call to `model` always returns a tuple, so we need to pull the + # loss value out of the tuple. + loss = outputs[0] + + # Accumulate the training loss over all of the batches so that we can + # calculate the average loss at the end. `loss` is a Tensor containing a + # single value; the `.item()` function just returns the Python value + # from the tensor. + total_loss += loss.item() + # total_loss += loss + + # Perform a backward pass to calculate the gradients. + print(loss) + loss.backward() + + # Clip the norm of the gradients to 1.0. + # This is to help prevent the "exploding gradients" problem. + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + + # Update parameters and take a step using the computed gradient. + # The optimizer dictates the "update rule"--how the parameters are + # modified based on their gradients, the learning rate, etc. + optimizer.step() + + # Update the learning rate. + scheduler.step() + + # Calculate the average loss over the training data. + avg_train_loss = total_loss / len(train_dataloader) + + # Store the loss value for plotting the learning curve. + loss_values.append(avg_train_loss) + + print("") + print(" Average training loss: {0:.2f}".format(avg_train_loss)) + print(" Training epoch took: {:}".format(format_time(time.time() - t0))) + + # ======================================== + # Validation + # ======================================== + # After the completion of each training epoch, measure our performance on + # our validation set. + + print("") + print("Running Validation...") + + t0 = time.time() + + # Put the model in evaluation mode--the dropout layers behave differently + # during evaluation. + model.eval() + + # Tracking variables + eval_loss, eval_accuracy = 0, 0 + nb_eval_steps, nb_eval_examples = 0, 0 + + # Evaluate data for one epoch + for batch in validation_dataloader: + + # Add batch to GPU + batch = tuple(t.to(device) for t in batch) + + # Unpack the inputs from our dataloader + b_input_ids, b_input_mask, b_labels = batch + + # Telling the model not to compute or store gradients, saving memory and + # speeding up validation + with torch.no_grad(): + + # Forward pass, calculate logit predictions. + # This will return the logits rather than the loss because we have + # not provided labels. + # token_type_ids is the same as the "segment ids", which + # differentiates sentence 1 and 2 in 2-sentence tasks. + # The documentation for this `model` function is here: + # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification + outputs = model(b_input_ids, + token_type_ids=None, + attention_mask=b_input_mask) + + # Get the "logits" output by the model. The "logits" are the output + # values prior to applying an activation function like the softmax. + logits = outputs[0] + + # Move logits and labels to CPU + logits = logits.detach().cpu().numpy() + label_ids = b_labels.to('cpu').numpy() + + # Calculate the accuracy for this batch of test sentences. + tmp_eval_accuracy = flat_accuracy(logits, label_ids) + + # Accumulate the total accuracy. + eval_accuracy += tmp_eval_accuracy + + # Track the number of batches + nb_eval_steps += 1 + + # Report the final accuracy for this validation run. + print(" Accuracy: {0:.2f}".format(eval_accuracy/nb_eval_steps)) + print(" Validation took: {:}".format(format_time(time.time() - t0))) + +print("") +print("Training complete!") \ No newline at end of file