| # PyTorch - Python Package Training |
|
|
| ## Overview |
|
|
| The directory provides code to fine tune a transformer model ([BERT-base](https://huggingface.co/bert-base-cased)) from Huggingface Transformers Library for sentiment analysis task. [BERT](https://ai.googleblog.com/2018/11/open-sourcing-bert-state-of-art-pre.html) (Bidirectional Encoder Representations from Transformers) is a transformers model pre-trained on a large corpus of unlabeled text in a self-supervised fashion. In this sample, we use [IMDB sentiment classification dataset](https://huggingface.co/datasets/imdb) for the task. We show you packaging a PyTorch training model to submit it to Vertex AI using pre-built PyTorch containers and handling Python dependencies through Python build scripts (`setup.py`). |
|
|
| ## Prerequisites |
| * Setup your project by following the instructions from [documentation](https://cloud.google.com/vertex-ai/docs/start/cloud-environment) |
| * Change directories to this sample. |
|
|
| ## Directory Structure |
|
|
| * `trainer` directory: all Python modules to train the model. |
| * `scripts` directory: command-line scripts to train the model on Vertex AI. |
| * `setup.py`: `setup.py` scripts specifies Python dependencies required for the training job. Vertex Training uses pip to install the package on the training instances allocated for the job. |
|
|
| ### Trainer Modules |
| | File Name | Purpose | |
| | :-------- | :------ | |
| | [metadata.py](trainer/metadata.py) | Defines: metadata for classification task such as predefined model dataset name, target labels. | |
| | [utils.py](trainer/utils.py) | Includes: utility functions such as data input functions to read data, save model to GCS bucket. | |
| | [model.py](trainer/model.py) | Includes: function to create model with a sequence classification head from a pretrained model. | |
| | [experiment.py](trainer/experiment.py) | Runs the model training and evaluation experiment, and exports the final model. | |
| | [task.py](trainer/task.py) | Includes: 1) Initialize and parse task arguments (hyper parameters), and 2) Entry point to the trainer. | |
|
|
| ### Scripts |
|
|
| * [train-cloud.sh](scripts/train-cloud.sh) This script submits a training job to Vertex AI |
|
|
| ## How to run |
| For local testing, run: |
| ``` |
| !cd python_package && python -m trainer.task |
| ``` |
|
|
| For cloud training, once the prerequisites are satisfied, update the |
| `BUCKET_NAME` environment variable in `scripts/train-cloud.sh`. You may then |
| run the following script to submit an AI Platform Training job: |
| ``` |
| source ./python_package/scripts/train-cloud.sh |
| ``` |
|
|
| ## Run on GPU |
| The provided trainer code runs on a GPU if one is available including data loading and model creation. |
|
|
| To run the trainer code on a different GPU configuration or latest PyTorch pre-built container image, make the following changes to the trainer script. |
| * Update the PyTorch image URI to one of [PyTorch pre-built containers](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#available_container_images) |
| * Update the [`worker-pool-spec`](https://cloud.google.com/vertex-ai/docs/training/configure-compute?hl=hr) in the gcloud command that includes a GPU |
|
|
| Then, run the script to submit a Custom Job on Vertex Training job: |
| ``` |
| source ./scripts/train-cloud.sh |
| ``` |
|
|
| ### Versions |
| This script uses the pre-built PyTorch containers for PyTorch 1.7. |
| * `us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-7:latest` |
|
|
|
|
