6.2 Prototyping with PyTorch, scikit-learn, JAX & Model Garden in Notebooks

Google Cloud Professional Machine Learning Engineer PMLE

Track your progress as you read

Finished sections get a check mark.

1 Chapter 1: Exam Overview & the Agent Platform Map

1.1 PMLE Exam Facts, Policies & the June 2026 Blueprint 1.2 From Vertex AI to Gemini Enterprise Agent Platform: The Product Name Map

2 Chapter 2: Low-Code ML with BigQuery ML

2.1 Choosing, Training & Predicting with BigQuery ML Models 2.2 Feature Engineering & Feature Selection in BigQuery ML 2.3 Fine-Tuning & Evaluating Gemini Models from BigQuery

3 Chapter 3: AutoML & Google Cloud AI APIs

3.1 Training Models with Agent Platform AutoML 3.2 Building with Document AI, Vision, Translation & Other AI APIs

4 Chapter 4: Foundation Models in Model Garden

4.1 Evaluating & Selecting Models in Model Garden 4.2 Building & Tuning Solutions with Gemini, Imagen, Veo & MaaS 4.3 Optimizing Gemini Applications for Cost, Latency & Availability

5 Chapter 5: Exploring, Preprocessing & Protecting Data

5.1 Organizing & Exploring Tabular, Text & Image Data 5.2 Choosing a Preprocessing Tool: BigQuery, Dataflow, Spark or Python 5.3 Creating & Consolidating Features in Agent Platform Feature Store 5.4 Protecting PII & Sensitive Data in ML Workflows

6 Chapter 6: Notebooks, Prototyping & Experiment Tracking

6.1 Workbench vs. Colab Enterprise: Secure, Collaborative Notebooks 6.2 Prototyping with PyTorch, scikit-learn, JAX & Model Garden in Notebooks 6.3 Experiment Environments, Tracking & Lineage with Experiments and ML Metadata

7 Chapter 7: Evaluating Predictive & Generative AI Solutions

7.1 Evaluation Metrics for Predictive Models 7.2 Evaluating Gen AI: Computation Metrics, Rubrics & LLM-as-a-Judge

8 Chapter 8: Choosing the Modeling Approach

8.1 Choosing the Model Type: ARIMA, Trees, DNNs & LLMs 8.2 Choosing the Product & Deployment Strategy 8.3 Modeling Techniques for Interpretability Requirements

9 Chapter 9: Training Data & Training SDKs

9.1 Organizing & Ingesting Training Data from Cloud Storage & BigQuery 9.2 Custom Training on Agent Platform: Containers, Jobs & Worker Pools 9.3 Kubeflow on GKE, AutoML & Tabular Workflows as Training Choices

10 Chapter 10: Troubleshooting, Hyperparameter Tuning & Fine-Tuning

10.1 Troubleshooting ML Training Failures 10.2 Hyperparameter Tuning on Agent Platform 10.3 Fine-Tuning Foundation Models & When to Tune

11 Chapter 11: Training Hardware & Distributed Training

11.1 Choosing CPUs, GPUs & TPUs for Training 11.2 Distributed Training: Data & Model Parallelism

12 Chapter 12: Deploying Models for Inference

12.1 Batch vs. Online Inference on Agent Platform 12.2 Serving Beyond Agent Platform Endpoints: Model Garden, Cloud Run & GKE 12.3 Prebuilt & Custom Serving Containers

13 Chapter 13: Model Versions, Rollouts & Inference Logic

13.1 Organizing & Versioning Models in Model Registry 13.2 Rollout Strategies: A/B Tests, Canaries & Traffic Splits 13.3 Inference Preprocessing & Postprocessing

14 Chapter 14: Scaling Online Serving

14.1 Serving Features Online with Feature Store 14.2 Public, Dedicated & Private Endpoints 14.3 Serving Hardware & Autoscaling for Throughput 14.4 Tuning Models for Production Training & Serving

15 Chapter 15: End-to-End ML Pipelines

15.1 Validating Data & Models in Pipelines 15.2 Agent Platform Pipelines with KFP: Components, Templates & Caching 15.3 Ensuring Consistent Preprocessing Between Training & Serving

16 Chapter 16: Orchestration with Managed Airflow & Ray

16.1 Orchestrating ML Workflows with Managed Service for Apache Airflow 16.2 Ray on Agent Platform for Distributed ML Workloads

17 Chapter 17: Automating Retraining & CI/CD/CT

17.1 Choosing a Retraining Policy 17.2 CI/CD/CT for ML with Cloud Build

18 Chapter 18: Securing AI Systems, Responsible AI & Explainability

18.1 Protecting AI Systems from Data & Model Exfiltration 18.2 Defending Gen AI Apps: Safety Filters, Model Armor & Regex Guardrails 18.3 Responsible AI: Detecting & Monitoring Bias 18.4 Model Explainability on Agent Platform

19 Chapter 19: Monitoring AI Solutions

19.1 Configuring Model Monitoring on Agent Platform 19.2 Training-Serving Skew, Data Drift, Concept Drift & Attribution Drift 19.3 Monitoring, Testing & Evaluating Gen AI Solutions in Production

6.2 Prototyping with PyTorch, scikit-learn, JAX & Model Garden in Notebooks

Key Takeaways Explain This Topic AI

Workbench instances come with JupyterLab and preinstalled deep learning packages for TensorFlow and PyTorch, and you can add conda environments as extra kernels.

Model Garden model cards for most open foundation and fine-tunable models include an Open notebook option for tuning in a notebook.

The Model Garden SDK's OpenModel class can list deploy options and deploy an open model such as Gemma to an Agent Platform endpoint from a notebook.

The Workbench executor runs a notebook on Agent Platform custom training, once or on a schedule, with parameters that change each run.

Prototype on a data sample in the notebook, then move the same code into a training script or container for full-scale runs.

Last updated: September 2026

The exam guide expects you to develop models in Agent Platform Workbench or Colab Enterprise notebooks using common frameworks (for example, PyTorch, sklearn, and JAX) and to use foundational and open-source models in Model Garden to create prototypes in notebook environments . The goal is fast iteration that can later become production training without a rewrite.

Choosing a Framework for the Prototype

Framework Strengths Typical PMLE scenario

scikit-learn Classical ML on tabular data, simple APIs, pipelines for preprocessing plus model Baseline gradient boosting or logistic regression on a sampled BigQuery extract

XGBoost / LightGBM Strong tabular performance Churn or fraud prototype before scaling

PyTorch Flexible deep learning, large ecosystem (vision, NLP, Hugging Face) Fine-tuning a vision or transformer model

TensorFlow / Keras Deep learning with strong TPU and serving integrations Image or text model destined for TensorFlow serving

JAX Composable transformations ( jit , grad , vmap , pmap ), high performance on TPUs Research-style models and large-scale training on TPUs

Setting Up the Environment

Workbench instances ship with JupyterLab and preinstalled TensorFlow and PyTorch packages. Add conda environments for other kernels, such as a specific framework version or JAX. Use a custom container when the whole team needs the same image.

Colab Enterprise runtimes use templates for machine type and GPUs. Install extra packages at the start of the notebook or with a post-startup script, and pin versions for reproducibility.

Pick GPUs for deep learning prototypes. Tabular scikit-learn work usually needs only CPU and memory.

Read data from BigQuery (the built-in integration or BigQuery DataFrames) and Cloud Storage , not local uploads, so the same paths work in training jobs.

Prototyping Patterns That Scale

Sample first. Develop on a stratified sample (for example, 1% of rows) to iterate in seconds. Validate the pipeline end to end before using full data.

Structure code as functions ( load_data , preprocess , train , evaluate ), not a long run of cells that depend on hidden state.

Fix random seeds and log library versions.

Establish a baseline (a simple model or a BigQuery ML model) before building complex architectures, so improvements are measurable.

Log parameters and metrics to Experiments on Agent Platform from the first run (Section 6.3).

Profile before scaling. A slow input pipeline is often the bottleneck, not the model.

From notebook to scale

Need Move to

Run the same notebook on a schedule or with bigger hardware Workbench executor or Colab Enterprise scheduled runs. The Workbench executor runs on Agent Platform custom training, with parameters per run

Full-data training, distributed jobs, hyperparameter tuning A custom training job built from the notebook's functions packaged as a Python package or container (Chapter 9)

Repeatable end-to-end workflow Agent Platform Pipelines components (Chapter 15)
Notebook anti-patterns to avoid

Cells that only work when run out of order, which leaves hidden state no one else can reproduce.

Training on the full dataset inside an interactive session on an oversized GPU VM for hours.

Copying production data to the runtime's local disk, where it's lost at deletion and bypasses governance.

Hard-coded paths, project IDs, and model versions scattered through cells instead of one configuration cell.

Using Foundation Models in a Notebook

Most gen AI prototypes start in a notebook by calling Gemini through the Google Gen AI SDK configured for Agent Platform: from google import genai

client = genai.Client(vertexai=True, project=PROJECT_ID, location="global")
resp = client.models.generate_content(
model=MODEL_ID, # keep the model version in configuration
contents=["Summarize this incident report in 3 bullets:", report_text],
)
print(resp.text)

Prototype checklist for gen AI notebooks:

Build a small evaluation set (inputs plus expected outputs or rubrics) right away.

Compare prompts and models side by side, then log results as experiment runs.

Test structured output with a response schema if downstream code parses JSON.

Track token counts and latency to estimate production cost early (Chapter 4).

Never paste production PII into prompts during exploration unless it's allowed and protected.

Using Open Models from Model Garden

Model Garden supports notebook-first workflows for open models:

Open notebook on a model card: most open foundation and fine-tunable models include a sample notebook for tuning, deployment, and inference in Colab Enterprise or Workbench.

Model Garden SDK: list deployable models, check deployment options (machine type and accelerator), and deploy from code: from vertexai import model_garden

model = model_garden.OpenModel("google/gemma3@gemma-3-12b-it")
print(model.list_deploy_options())
endpoint = model.deploy(machine_type="g2-standard-48",
accelerator_type="NVIDIA_L4",
accelerator_count=4,
accept_eula=True)

Hugging Face models can be deployed through the same flow. Google-verified deployment configurations reduce trial and error.

Tuning: use the model card's fine-tuning notebook or pipeline for LoRA or full tuning, then deploy the tuned weights.

Cost warning: a deployed endpoint bills for its accelerators while it exists. Undeploy prototype endpoints when you finish, or use MaaS APIs for prototypes that don't need your own weights.

Worked Scenario

A retailer wants to test whether an open vision-language model can tag product photos better than its current rules. In Colab Enterprise, a data scientist:

Loads 500 labeled product photos from Cloud Storage.

Calls Gemini with a tagging prompt and a response schema to get a strong managed baseline.

Deploys an open vision-language model from Model Garden to a temporary GPU endpoint with the SDK, and runs the same 500 photos.

Logs accuracy, latency, and estimated cost per 1,000 images for both models as experiment runs.

Undeploys the endpoint. The team picks the approach that meets the accuracy target at the lowest total cost, then moves the chosen path into a pipeline.