Google Cloud AI Platform部署Torch(BERT)模型GPU与自定义类兼容问题咨询

I’ve run into this exact conflicting limitation before—Google Cloud AI Platform’s older custom prediction class approach has a strict constraint: the mls1 machine types that support CUSTOM_CLASS don’t offer GPU options, while GPU-enabled machine types (like n1-highcpu-4 or n1-standard-8) won’t work with the custom prediction class framework.

The reliable workaround here is to use Custom Containers for AI Platform Prediction. This method lets you package your PyTorch model, custom prediction logic, and all dependencies into a Docker image that can run on any GPU-enabled machine type supported by the platform.

Step 1: Create a Custom Prediction Container #

First, build a Docker image that includes everything your model needs:

• A GPU-compatible PyTorch base image (match your Torch version; for 1.0.0, use gcr.io/deeplearning-platform-release/pytorch-gpu:1.0.0)
• Your custom prediction code (replace the CustomModelPrediction class logic with an HTTP endpoint handler)
• Any additional dependencies (like your my-torch-package-0.1 package)

Here’s a sample Dockerfile:

# Use GPU-enabled PyTorch base image matching your Torch version
FROM gcr.io/deeplearning-platform-release/pytorch-gpu:1.0.0

# Install your custom package
COPY my-torch-package-0.1.tar.gz /tmp/
RUN pip install /tmp/my-torch-package-0.1.tar.gz

# Copy your prediction handler code
COPY predictor.py /app/
WORKDIR /app

# Expose ports required by AI Platform
EXPOSE 8080

# Run the prediction server (implement required endpoints here)
CMD ["python", "predictor.py"]

Your predictor.py needs to implement two core HTTP endpoints that AI Platform expects:

• A /health endpoint for liveness checks (returns a 200 OK response)
• A /predict endpoint that accepts JSON requests, runs inference with your BERT model, and returns formatted results

Step 2: Build and Push the Image to GCR #

Once your Dockerfile and predictor code are ready, build the image and push it to Google Container Registry:

# Build the image (replace YOUR_PROJECT_ID with your GCP project ID)
docker build -t gcr.io/[YOUR_PROJECT_ID]/pytorch-bert-gpu:v1 .

# Push the image to GCR
docker push gcr.io/[YOUR_PROJECT_ID]/pytorch-bert-gpu:v1

Step 3: Deploy the Model Version with GPU #

Use the custom container to create a model version with a GPU-enabled machine type:

gcloud ai-platform versions create {VERSION} \
  --model {MODEL_NAME} \
  --container-image-uri gcr.io/[YOUR_PROJECT_ID]/pytorch-bert-gpu:v1 \
  --machine-type n1-standard-8 \
  --accelerator count=1,type=nvidia-tesla-k80 \
  --runtime-version 1.14

Key Notes #

• Ensure your base image’s CUDA version is compatible with the GPU type you’re using (Tesla K80 works with CUDA 10.0, which matches the PyTorch 1.0.0 GPU image)
• The custom container approach gives you full control over the environment, so you can replicate your local setup exactly
• You no longer need the --package-uris or --prediction-class flags since all logic is packaged directly in the container

内容的提问来源于stack exchange，提问作者user6377061