Update pytorch examples to use gpu if device is set (#849)

Author: vishalbollu

examples/pytorch/image-classifier/cortex.yaml

@@ -4,6 +4,8 @@
   predictor:
     type: python
     path: predictor.py
+    config:
+      device: cuda  # use "cpu" to run on CPUs
   compute:
     cpu: 1
     gpu: 1

examples/pytorch/image-classifier/predictor.py

@@ -10,7 +10,7 @@
 
 class PythonPredictor:
     def __init__(self, config):
-        model = torchvision.models.alexnet(pretrained=True)
+        model = torchvision.models.alexnet(pretrained=True).to(config["device"])
         model.eval()
         # https://github.com/pytorch/examples/blob/447974f6337543d4de6b888e244a964d3c9b71f6/imagenet/main.py#L198-L199
         normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -22,12 +22,14 @@ def __init__(self, config):
             "https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"
         ).text.split("\n")[1:]
         self.model = model
+        self.device = config["device"]
 
     def predict(self, payload):
         image = requests.get(payload["url"]).content
         img_pil = Image.open(BytesIO(image))
         img_tensor = self.preprocess(img_pil)
         img_tensor.unsqueeze_(0)
+        img_tensor = img_tensor.to(self.device)
         with torch.no_grad():
             prediction = self.model(img_tensor)
         _, index = prediction[0].max(0)

examples/pytorch/language-identifier/cortex.yaml

@@ -6,7 +6,3 @@
     path: predictor.py
   tracker:
     model_type: classification
-  compute:
-    cpu: 1
-    gpu: 1
-    mem: 4G

examples/pytorch/object-detector/cortex.yaml

@@ -4,6 +4,8 @@
   predictor:
     type: python
     path: predictor.py
+    config:
+      device: cuda  # use "cpu" to run on CPUs
   compute:
     cpu: 1
     gpu: 1

examples/pytorch/object-detector/predictor.py

@@ -11,7 +11,8 @@
 
 class PythonPredictor:
     def __init__(self, config):
-        model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+        self.device = config["device"]
+        model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True).to(self.device)
         model.eval()
 
         self.preprocess = transforms.Compose([transforms.ToTensor()])
@@ -25,17 +26,17 @@ def predict(self, payload):
         threshold = float(payload["threshold"])
         image = requests.get(payload["url"]).content
         img_pil = Image.open(BytesIO(image))
-        img_tensor = self.preprocess(img_pil)
+        img_tensor = self.preprocess(img_pil).to(self.device)
         img_tensor.unsqueeze_(0)
 
         with torch.no_grad():
             pred = self.model(img_tensor)
 
-        predicted_class = [self.coco_labels[i] for i in list(pred[0]["labels"].numpy())]
+        predicted_class = [self.coco_labels[i] for i in list(pred[0]["labels"].cpu().numpy())]
         predicted_boxes = [
-            [(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]["boxes"].detach().numpy())
+            [(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]["boxes"].detach().cpu().numpy())
         ]
-        predicted_score = list(pred[0]["scores"].detach().numpy())
+        predicted_score = list(pred[0]["scores"].detach().cpu().numpy())
         predicted_t = [predicted_score.index(x) for x in predicted_score if x > threshold]
         if len(predicted_t) == 0:
             return [], []

examples/pytorch/reading-comprehender/cortex.yaml

@@ -4,6 +4,8 @@
   predictor:
     type: python
     path: predictor.py
+    config:
+      device: cuda  # use "cpu" to run on CPUs
   compute:
     cpu: 1
     gpu: 1

examples/pytorch/reading-comprehender/predictor.py

@@ -5,8 +5,14 @@
 
 class PythonPredictor:
     def __init__(self, config):
+
+        cuda_device = -1
+        if config["device"] == "cuda":
+            cuda_device = 0
+
         self.predictor = AllenNLPPredictor.from_path(
-            "https://storage.googleapis.com/allennlp-public-models/bidaf-elmo-model-2018.11.30-charpad.tar.gz"
+            "https://storage.googleapis.com/allennlp-public-models/bidaf-elmo-model-2018.11.30-charpad.tar.gz",
+            cuda_device=cuda_device,
         )
 
     def predict(self, payload):