mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
Upload a ResNet predictor example (#3765)
This example uses aiplatform and torch library to provide a ResNet predictor.
This commit is contained in:
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from torchvision.models import detection, resnet50, ResNet50_Weights
|
||||
from typing import Dict, List
|
||||
|
||||
class ResNetPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists("model.pth.tar"):
|
||||
self.model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
|
||||
stat_dic = torch.load("model.pth.tar")
|
||||
self.model.load_state_dict(stat_dic['state_dict'])
|
||||
else:
|
||||
weights = ResNet50_Weights.DEFAULT
|
||||
self.model = resnet50(weights=weights)
|
||||
self.model.eval()
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> torch.Tensor:
|
||||
instances = prediction_input["instances"]
|
||||
return torch.Tensor(instances)
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, instances: torch.Tensor) -> List[str]:
|
||||
return self._model(instances)
|
||||
|
||||
def postprocess(self, prediction_results: List[str]) -> Dict:
|
||||
return {"predictions": prediction_results}
|
||||
Reference in New Issue
Block a user