Use GPU for test/apis/batch/image-classifier-alexnet

Author: deliahu

test/apis/batch/image-classifier-alexnet/app/main.py

@@ -73,7 +73,7 @@ def handle_batch(image_urls: List[str]):
         tensor_list.append(app.preprocess(img_pil))
 
     # classify the batch of images
-    img_tensor = torch.stack(tensor_list)
+    img_tensor = torch.stack(tensor_list).to(app.device)
     with torch.no_grad():
         prediction = app.model(img_tensor)
     _, indices = prediction.max(1)

test/apis/batch/image-classifier-alexnet/cortex_gpu.yaml

@@ -5,6 +5,8 @@
     - name: api
       image: quay.io/cortexlabs-test/batch-image-classifier-alexnet-gpu:latest
       command: ["uvicorn", "--workers", "1", "--host", "0.0.0.0", "--port", "$(CORTEX_PORT)", "main:app"]
+      env:
+        TARGET_DEVICE: "cuda"
       readiness_probe:
         http_get:
           path: "/healthz"

test/apis/batch/image-classifier-alexnet/gpu.Dockerfile

@@ -8,6 +8,8 @@ RUN apt-get update \
         build-essential \
         git \
         cmake \
+        libjpeg8-dev \
+        zlib1g-dev \
     && apt-get clean -qq && rm -rf /var/lib/apt/lists/*
 
 ENV LC_ALL=C.UTF-8 LANG=C.UTF-8