Repository: chainer/chainer

Issue:
ChainerX example support: image_captioning: predict.py



Discussion:

Resolving patch:
diff --git a/examples/image_captioning/model.py b/examples/image_captioning/model.py
--- a/examples/image_captioning/model.py
+++ b/examples/image_captioning/model.py
@@ -238,24 +238,27 @@ def predict(self, img_feats, bos, eos, max_caption_length):
         """Batch of image features to captions."""
         hx, cx, _ = self.reset(img_feats)
 
-        captions = self.xp.full((img_feats.shape[0], 1), bos, dtype=np.int32)
-        for i in range(max_caption_length):
-            # Create a list of the previous tokens to treat as inputs
-            xs = [Variable(self.xp.atleast_1d(c[-1])) for c in captions]
-
-            # Get the predictions `ys`
-            hx, cx, ys = self.step(hx, cx, xs)
-
-            # From `ys`, get the indices for the highest confidence.
-            # These indices correspond to the predicted tokens
-            #
-            # Note that this is a greedy approach and that it can by replaced
-            # by e.g. beam search
-            pred = ys.array.argmax(axis=1).astype(np.int32)
-            captions = self.xp.hstack((captions, pred[:, None]))
-
-            if (pred == eos).all():
-                break
+        with chainer.using_device(self.device):
+            xp = self.xp
+            captions = xp.full(
+                (img_feats.shape[0], 1), bos, dtype=np.int32)
+            for i in range(max_caption_length):
+                # Create a list of the previous tokens to treat as inputs
+                xs = [xp.atleast_1d(c[-1]) for c in captions]
+
+                # Get the predictions `ys`
+                hx, cx, ys = self.step(hx, cx, xs)
+
+                # From `ys`, get the indices for the highest confidence.
+                # These indices correspond to the predicted tokens
+                #
+                # Note that this is a greedy approach and that it can by
+                # replaced by e.g. beam search
+                pred = ys.array.argmax(axis=1).astype(np.int32)
+                captions = xp.hstack((captions, pred[:, None]))
+
+                if (pred == eos).all():
+                    break
         return captions
 
     def reset(self, img_feats):
diff --git a/examples/image_captioning/predict.py b/examples/image_captioning/predict.py
--- a/examples/image_captioning/predict.py
+++ b/examples/image_captioning/predict.py
@@ -2,12 +2,14 @@
 import argparse
 import glob
 import os
+import sys
 
 import numpy as np
 from PIL import Image
 
 import chainer
 from chainer import serializers
+import chainerx
 
 import datasets
 from model import ImageCaptionModel
@@ -26,12 +28,27 @@ def main():
     parser.add_argument('--rnn', type=str, default='nsteplstm',
                         choices=['nsteplstm', 'lstm'],
                         help='Language model layer type')
-    parser.add_argument('--gpu', type=int, default=0,
-                        help='GPU ID (negative value indicates CPU)')
+    parser.add_argument('--device', '-d', type=str, default='-1',
+                        help='Device specifier. Either ChainerX device '
+                        'specifier or an integer. If non-negative integer, '
+                        'CuPy arrays with specified device id are used. If '
+                        'negative integer, NumPy arrays are used')
     parser.add_argument('--max-caption-length', type=int, default=30,
                         help='Maximum caption length generated')
+    group = parser.add_argument_group('deprecated arguments')
+    group.add_argument('--gpu', '-g', dest='device',
+                       type=int, nargs='?', const=0,
+                       help='GPU ID (negative value indicates CPU)')
     args = parser.parse_args()
 
+    device = chainer.get_device(args.device)
+    if device.xp is chainerx:
+        sys.stderr.write('This example does not support ChainerX devices.\n')
+        sys.exit(1)
+
+    print('Device: {}'.format(device))
+    print()
+
     # Load the dataset to obtain the vocabulary, which is needed to convert
     # predicted tokens into actual words
     train, _ = datasets.get_mscoco(args.mscoco_root)
@@ -41,6 +58,8 @@ def main():
     model = ImageCaptionModel(len(train.vocab), rnn=args.rnn)
     serializers.load_npz(args.model, model)
 
+    model.to_device(device)
+
     if args.img_dir:  # Read all images in directory
         img_paths = [
             i for i in glob.glob(os.path.join(args.img_dir, '*')) if
@@ -58,11 +77,7 @@ def main():
         img = model.prepare(img)
         imgs.append(img)
     imgs = np.asarray(imgs)
-
-    if args.gpu >= 0:
-        chainer.backends.cuda.get_device_from_id(args.gpu).use()
-        model.to_gpu()
-        imgs = chainer.backends.cuda.to_gpu(imgs)
+    imgs = device.send(imgs)
 
     bos = vocab['<bos>']
     eos = vocab['<eos>']
@@ -70,7 +85,7 @@ def main():
             chainer.no_backprop_mode():
         captions = model.predict(
             imgs, bos=bos, eos=eos, max_caption_length=args.max_caption_length)
-    captions = chainer.backends.cuda.to_cpu(captions)
+    captions = chainer.get_device('@numpy').send(captions)
 
     # Print the predicted captions
     file_names = [os.path.basename(path) for path in img_paths]


Test patch: