Training an ML Model
A crash course on using Beam to run an entire ML workflow, from training a model to deployment.
This example demonstrates a basic movie recommendation system. The following capabilities are demonstrated:
- Training a model using the MovieLens dataset
- Saving the trained model to a Volume
- Retrieving the trained model from a Volume during inference
- Deploying a REST API that accepts a user ID and returns customized recommendations for that user
Defining the runtime
This is the runtime our code will run in. We’ll define the compute settings and Python packages to install.
Using Volumes to save trained models
We’re going to mount a Volume, which is a writable data store.
- During training, we will save our models to this volume.
- During inference, we will retrieve our trained model.
We will access the volume at this path: ./trained_models
Training the model
We use an embedding layer for both the user and movie, to compress the respective one-hot encoded vectors into rich, compact representations that are easier to model.
These two representations are concatenated into a single vector and then passed into a simple fully connected neural network.
There’s a lot going on here, but the main thing to note is the code at the end, which saves the trained model to a Volume.
We’ve wrapped this function in a run()
, which will let us run it asynchronously on Beam.
Running the training script
Running the training script is straightforward — just enter your shell, and kick off a run
:
This command will containerize your app and run it on a remote container. Feel free to close your terminal window, if you wish. The app will continue running on Beam remotely.
You’ll see the training logs stream to your shell. This will look something like this:
Making Predictions
The whole point of a recommendation system is to make predictions dynamically, and that requires us to deploy an API.
We’ll write a function that takes in a user ID, and returns a list of movies that the user is predicted to enjoy watching.
Deployment
We’re going to deploy the app as a REST API, which will allow us to generate movie recommendations for users in real-time.
Above the run_inference
function, add a rest_api
decorator:
Now, go back to the shell and run this command to deploy the app. Make sure to customize this with the actual name of the file that has the run_inference
function you created:
You’ll see some logs appear in your shell, showing the deployment status:
Calling the API
If you navigate to the link in the last line of the shell output, you’ll be able to login to your Beam web dashboard and view the details of your deployment.
In the dashboard, you’ll be able to copy a cURL request which you can use to make predictions to the API. We’ll pass in a user ID, and ask for 3 movie recommendations in response.
The movie recommendations will be returned as JSON:
Was this page helpful?