Train a Model to Check Your Grammar Using W&B, PyTorch Lightning ⚡, and 🤗

Based on Ayush Chaurasia's awesome W&B report and colab which performs the same task using BERT, vanilla PyTorch, and W&B.

In this notebook, we are going to train a model to detect ungrammatical sentences from the CoLA dataset. To perform the classification, we will be using Pytorch Lightning ⚡ to fine tune DistilBERT, a transformer model from huggingface 🤗.

We'll use Weights & Biases to:

• Version our model inputs and outputs using W&B Artifacts, including preprocessing steps, train/validation splits, and model checkpoints
• Log and visualize training and validation performance using W&B's Pytorch Lightning integration
• Visualize and explore the raw dataset using W&B Tables
• Orchestrate a hyperparameter search using W&B Sweeps

Be sure to follow the links that each run outputs to your W&B workspace, where you will be able to see...

Your model's performance metrics updating in real time

The raw data as a W&B Table, which you can sort, group, and filter

An awesome artifact graph showing our full pipeline

Interactive visualizations of how our hyperparameter choices effect model performance

The CoLA Dataset 🥤

We’ll fine tune the model on The Corpus of Linguistic Acceptability (CoLA) dataset for single sentence classification. It’s a set of sentences labeled as grammatically correct or incorrect. It was first published in May of 2018, and is one of the tests included in the “GLUE Benchmark” on which models like DistilBERT are competing.

We'll use a reference artifact to store a pointer to the source data. The advantages of doing this are:

• Any runs that use this artifact reference will be able to trace their lineage back to the true source
• We can use W&B to download the raw data in our code.

The cell below starts a run with job type register-data. In the context of this run, we:

1. Create an artifact called cola-raw
2. Add a reference to the CoLA dataset to our cola-raw artifact
3. Log the cola-raw artifact to Weights & Biases.

Tokenization 🪙

The cell below defines the function tokenize_data, which transforms a list of sentences and a list of labels into a tuple of torch.tensor objects which can be consumed by the transormer model we'll be using. The 3 tensors returned are the tokenized form of the sentences, the attention masks indicating which tokens in each sentence correspond to actual words, and a tensor containing the original labels.

The code below executes a run of type preprocess-data, which will

1. Download the CoLA dataset using the reference artifact we logged previously
2. Log the entire dataset to W&B as a Table
3. Use the function tokenize_data to transform each sentence into a sequence of tokens and an attention mask
4. Log the preprocessed data as an artifact to W&B.

Splitting Our Data 🪓

For our training process, we want to split the data into a train and validation set. The train set is the data we will use to update the model parameters, while the validation set will be a smaller segment of data that we use to test whether our model is generalizing to examples that it hasn't been trained on.

The cell below executes a wandb.Run with job_type="split-data". In the context of this run we will:

1. Download the preprocessed-data artifact logged by our previous run
2. Use the random_split function from torch to perform a randomn 90/10 test/valiation split on the preprocessed data
3. Store the split datasets in a new artifact called split-dataset

Defining Our Model ⚡

We define our model and the associated training + validation procedures in the LightningModule below. The model itself is a pre-trained DistilBertForSequenceClassification with two labels.

Training & Tracking Our Model 📉

In the cell below, we define a function train which sets up and performs training in the context of a W&B run. The train function takes a configuration dictionary as input then passes it to wandb.init via the config keyword argument. We use the values saved in the wandb.config object associated with the run to set the parameters of our trainer and data loaders. This is a crucial best practice to ensure that the values logged in the config object (and displayed in the run table of the W&B app) represent the actual parameters of the experiment.

Running a Hyperparameter Sweep 🧹

W&B sweeps allow you to optimize your model hyperparameters with minimal effort. In general, the workflow of sweeps is:

1. Construct a dictionary or YAML file that defines the hyperparameter space
2. Call wandb.sweep(<sweep-dict>) from the python library or wandb sweep <yaml-file> from the command line to initialize the sweep in W&B
3. Run wandb.agent(<sweep-id>) (python lib) or wandb agent <sweep-id> (cli) to start a sweep agent to continuously:

• pull hyperparameter combinations from W&B
• run training with the given hyperparameters
• log training metrics back to W&B

We implement the sweeps workflow laid out above by:

1. Creating a sweep_config dictionary describing our hyperparameter space and objective

• The hyperparameters we will sweep over are learning_rate, batch_size, and epochs
• Our objective in this sweep is to maximize the validation/epoch_acc metric logged to W&B
• We will use the random strategy, which means we will sample uniformly from the parameter space indefinitely

2. Calling wandb.sweep(sweep_config) to create the sweep in our W&B project

• wandb.sweep will return a unique id for the sweep, saved as sweep_id

3. Calling wandb.agent(sweep_id, function=train) to create an agent that will execute training with different hyperparameter combinations

• The agent will repeatedly query W&B for hyperparameter combinations
• When wandb.init is called within an agent, the config dictionary of the returned run will be populated with the next hyperparameter combination in the sweep