Add int64 as label type and set it as default (#125)

* Add int64 as label type and set it as default Signed-off-by: ptredak <[email protected]> * Fixing multiple_reader example Signed-off-by: ptredak <[email protected]> * Remove empty cell from tf example Signed-off-by: ptredak <[email protected]>
NVIDIA · Aug 21, 2018 · 2319d24 · 2319d24
1 parent 0fbf3fc
commit 2319d24
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 14 deletions.
diff --git a/dali/tensorflow/daliop.cc b/dali/tensorflow/daliop.cc
@@ -64,8 +64,8 @@ REGISTER_OP("Dali")
   .Attr("shape: shape")
   .Attr("num_threads: int = -1")
   .Attr("device_id: int = -1")
-  .Attr("image_type: {float, int32, half} = DT_FLOAT")
-  .Attr("label_type: {float, int32, half} = DT_INT32")
+  .Attr("image_type: {float, int32, int64, half} = DT_FLOAT")
+  .Attr("label_type: {float, int32, int64} = DT_INT64")
   .Output("batch: image_type")
   .Output("label: label_type")
   .SetShapeFn([](tf::shape_inference::InferenceContext* c) {
@@ -196,6 +196,8 @@ class DaliOp : public tf::OpKernel {
   int device_id_;
 };
 
+using tf::int64;
+
 #define REGISTER_KERNEL(type_b, type_l) \
   REGISTER_KERNEL_BUILDER(              \
       Name("Dali")                      \
@@ -216,8 +218,14 @@ REGISTER_KERNEL(float, float);
 REGISTER_KERNEL(int, int);
 REGISTER_KERNEL(float, int);
 REGISTER_KERNEL(int, float);
+REGISTER_KERNEL(int64, float);
+REGISTER_KERNEL(int64, int);
+REGISTER_KERNEL(float, int64);
+REGISTER_KERNEL(int, int64);
+REGISTER_KERNEL(int64, int64);
 REGISTER_KERNEL_FP16(float);
 REGISTER_KERNEL_FP16(int);
+REGISTER_KERNEL_FP16(int64);
 
 #undef REGISTER_KERNEL
 #undef REGISTER_KERNEL_FP16
diff --git a/docs/examples/tensorflow/tensorflow-resnet50-various-readers.ipynb b/docs/examples/tensorflow/tensorflow-resnet50-various-readers.ipynb
@@ -214,6 +214,7 @@
     "import nvidia.dali.plugin.tf as dali_tf\n",
     "\n",
     "def get_batch_test_dali(batch_size, pipe_type):\n",
+    "    pipe_name, label_type, _ = pipe_type\n",
     "    pipes = [pipe_name(batch_size=batch_size, num_threads=2, device_id = device_id, num_gpus = N) for device_id in range(N)]\n",
     "\n",
     "    serialized_pipes = [pipe.serialize() for pipe in pipes]\n",
@@ -226,7 +227,7 @@
     "            image, label = daliop(serialized_pipeline = serialized_pipes[d],\n",
     "                shape = [BATCH_SIZE, 3, 227, 227],\n",
     "                image_type = tf.int32,\n",
-    "                label_type = tf.float32,\n",
+    "                label_type = label_type,\n",
     "                device_id = d)\n",
     "            images.append(image)\n",
     "            labels.append(label)\n",
@@ -250,48 +251,59 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "RUN: FileReadPipeline\n",
-      "OK : FileReadPipeline\n",
       "RUN: MXNetReaderPipeline\n",
       "OK : MXNetReaderPipeline\n",
+      "RUN: FileReadPipeline\n",
+      "OK : FileReadPipeline\n",
       "RUN: TFRecordPipeline\n",
       "OK : TFRecordPipeline\n"
      ]
     }
    ],
    "source": [
-    "pipe_types = [FileReadPipeline, MXNetReaderPipeline, TFRecordPipeline]\n",
+    "import numpy as np\n",
+    "\n",
+    "pipe_types = [[MXNetReaderPipeline, tf.float32, (0, 999)], \n",
+    "              [FileReadPipeline, tf.int32, (0, 1)], \n",
+    "              [TFRecordPipeline, tf.int64, (1, 1000)]]\n",
     "for pipe_name in pipe_types:\n",
-    "    print (\"RUN: \"  + pipe_name.__name__)\n",
+    "    print (\"RUN: \"  + pipe_name[0].__name__)\n",
     "    test_batch = get_batch_test_dali(BATCH_SIZE, pipe_name)\n",
     "    x = tf.placeholder(tf.float32, shape=[BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3], name='x')\n",
     "    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)\n",
     "    config = tf.ConfigProto(gpu_options=gpu_options)\n",
     "\n",
     "    with tf.Session(config=config) as sess:\n",
     "        for i in range(ITERATIONS):\n",
-    "            sess.run(test_batch)\n",
-    "    print(\"OK : \" + pipe_name.__name__)"
+    "            imgs, labels = sess.run(test_batch)\n",
+    "            # Testing correctness of labels\n",
+    "            for label in labels:\n",
+    "                ## labels need to be integers\n",
+    "                assert(np.equal(np.mod(label, 1), 0).all())\n",
+    "                ## labels need to be in range pipe_name[2]\n",
+    "                assert((label >= pipe_name[2][0]).all())\n",
+    "                assert((label <= pipe_name[2][1]).all())\n",
+    "    print(\"OK : \" + pipe_name[0].__name__)"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.15rc1"
+   "pygments_lexer": "ipython3",
+   "version": "3.5.2"
   }
  },
  "nbformat": 4,