Caffe2 - hipify python API #2

.jenkins/caffe2/test.sh

@@ -109,6 +109,10 @@ if [[ $BUILD_ENVIRONMENT == *-rocm* ]]; then
   # Our cuda top_k op has some asm code, the hipified version doesn't
   # compile yet, so we don't have top_k operator for now
   rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/top_k_test.py")
+
+  # this is a multi-gpu test
+  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/data_parallel_model_test.py")
+
 fi
 
 # Python tests

caffe2/python/cnn.py

@@ -5,7 +5,7 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from caffe2.python import brew, workspace
+from caffe2.python import brew
 from caffe2.python.model_helper import ModelHelper
 from caffe2.proto import caffe2_pb2
 import logging
@@ -17,7 +17,7 @@ class CNNModelHelper(ModelHelper):
     """
 
     def __init__(self, order="NCHW", name=None,
-                 use_gpu_engine=True, gpu_engine_exhaustive_search=False,
+                 use_cudnn=True, cudnn_exhaustive_search=False,
                  ws_nbytes_limit=None, init_params=True,
                  skip_sparse_optim=False,
                  param_model=None):
@@ -31,8 +31,8 @@ def __init__(self, order="NCHW", name=None,
 
         cnn_arg_scope = {
             'order': order,
-            'use_gpu_engine': use_gpu_engine,
-            'gpu_engine_exhaustive_search': gpu_engine_exhaustive_search,
+            'use_cudnn': use_cudnn,
+            'cudnn_exhaustive_search': cudnn_exhaustive_search,
         }
         if ws_nbytes_limit:
             cnn_arg_scope['ws_nbytes_limit'] = ws_nbytes_limit
@@ -45,8 +45,8 @@ def __init__(self, order="NCHW", name=None,
         )
 
         self.order = order
-        self.use_gpu_engine = use_gpu_engine
-        self.gpu_engine_exhaustive_search = gpu_engine_exhaustive_search
+        self.use_cudnn = use_cudnn
+        self.cudnn_exhaustive_search = cudnn_exhaustive_search
         self.ws_nbytes_limit = ws_nbytes_limit
         if self.order != "NHWC" and self.order != "NCHW":
             raise ValueError(
@@ -79,9 +79,9 @@ def ConvNd(self, *args, **kwargs):
         return brew.conv_nd(
             self,
             *args,
-            use_gpu_engine=self.use_gpu_engine,
+            use_cudnn=self.use_cudnn,
             order=self.order,
-            gpu_engine_exhaustive_search=self.gpu_engine_exhaustive_search,
+            cudnn_exhaustive_search=self.cudnn_exhaustive_search,
             ws_nbytes_limit=self.ws_nbytes_limit,
             **kwargs
         )
@@ -90,9 +90,9 @@ def Conv(self, *args, **kwargs):
         return brew.conv(
             self,
             *args,
-            use_gpu_engine=self.use_gpu_engine,
+            use_cudnn=self.use_cudnn,
             order=self.order,
-            gpu_engine_exhaustive_search=self.gpu_engine_exhaustive_search,
+            cudnn_exhaustive_search=self.cudnn_exhaustive_search,
             ws_nbytes_limit=self.ws_nbytes_limit,
             **kwargs
         )
@@ -101,9 +101,9 @@ def ConvTranspose(self, *args, **kwargs):
         return brew.conv_transpose(
             self,
             *args,
-            use_gpu_engine=self.use_gpu_engine,
+            use_cudnn=self.use_cudnn,
             order=self.order,
-            gpu_engine_exhaustive_search=self.gpu_engine_exhaustive_search,
+            cudnn_exhaustive_search=self.cudnn_exhaustive_search,
             ws_nbytes_limit=self.ws_nbytes_limit,
             **kwargs
         )
@@ -112,9 +112,9 @@ def GroupConv(self, *args, **kwargs):
         return brew.group_conv(
             self,
             *args,
-            use_gpu_engine=self.use_gpu_engine,
+            use_cudnn=self.use_cudnn,
             order=self.order,
-            gpu_engine_exhaustive_search=self.gpu_engine_exhaustive_search,
+            cudnn_exhaustive_search=self.cudnn_exhaustive_search,
             ws_nbytes_limit=self.ws_nbytes_limit,
             **kwargs
         )
@@ -123,9 +123,9 @@ def GroupConv_Deprecated(self, *args, **kwargs):
         return brew.group_conv_deprecated(
             self,
             *args,
-            use_gpu_engine=self.use_gpu_engine,
+            use_cudnn=self.use_cudnn,
             order=self.order,
-            gpu_engine_exhaustive_search=self.gpu_engine_exhaustive_search,
+            cudnn_exhaustive_search=self.cudnn_exhaustive_search,
             ws_nbytes_limit=self.ws_nbytes_limit,
             **kwargs
         )
@@ -147,16 +147,16 @@ def FC_Sparse(self, *args, **kwargs):
 
     def Dropout(self, *args, **kwargs):
         return brew.dropout(
-            self, *args, order=self.order, use_gpu_engine=self.use_gpu_engine, **kwargs
+            self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs
         )
 
     def LRN(self, *args, **kwargs):
         return brew.lrn(
-            self, *args, order=self.order, use_gpu_engine=self.use_gpu_engine, **kwargs
+            self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs
         )
 
     def Softmax(self, *args, **kwargs):
-        return brew.softmax(self, *args, use_gpu_engine=self.use_gpu_engine, **kwargs)
+        return brew.softmax(self, *args, use_cudnn=self.use_cudnn, **kwargs)
 
     def SpatialBN(self, *args, **kwargs):
         return brew.spatial_bn(self, *args, order=self.order, **kwargs)
@@ -169,7 +169,7 @@ def InstanceNorm(self, *args, **kwargs):
 
     def Relu(self, *args, **kwargs):
         return brew.relu(
-            self, *args, order=self.order, use_gpu_engine=self.use_gpu_engine, **kwargs
+            self, *args, order=self.order, use_cudnn=self.use_cudnn, **kwargs
         )
 
     def PRelu(self, *args, **kwargs):
@@ -187,7 +187,7 @@ def Sum(self, *args, **kwargs):
         return brew.sum(self, *args, **kwargs)
 
     def Transpose(self, *args, **kwargs):
-        return brew.transpose(self, *args, use_gpu_engine=self.use_gpu_engine, **kwargs)
+        return brew.transpose(self, *args, use_cudnn=self.use_cudnn, **kwargs)
 
     def Iter(self, *args, **kwargs):
         return brew.iter(self, *args, **kwargs)
@@ -197,15 +197,15 @@ def Accuracy(self, *args, **kwargs):
 
     def MaxPool(self, *args, **kwargs):
         return brew.max_pool(
-            self, *args, use_gpu_engine=self.use_gpu_engine, order=self.order, **kwargs
+            self, *args, use_cudnn=self.use_cudnn, order=self.order, **kwargs
         )
 
     def MaxPoolWithIndex(self, *args, **kwargs):
         return brew.max_pool_with_index(self, *args, order=self.order, **kwargs)
 
     def AveragePool(self, *args, **kwargs):
         return brew.average_pool(
-            self, *args, use_gpu_engine=self.use_gpu_engine, order=self.order, **kwargs
+            self, *args, use_cudnn=self.use_cudnn, order=self.order, **kwargs
         )
 
     @property
@@ -235,11 +235,6 @@ def CPU(self):
     @property
     def GPU(self, gpu_id=0):
         device_option = caffe2_pb2.DeviceOption()
-        if workspace.has_hip_support:
-            device_option.device_type = caffe2_pb2.HIP
-            device_option.hip_gpu_id = gpu_id
-        else:
-            device_option.device_type = caffe2_pb2.CUDA
-            device_option.cuda_gpu_id = gpu_id
-
+        device_option.device_type = caffe2_pb2.CUDA
+        device_option.cuda_gpu_id = gpu_id
         return device_option

caffe2/python/core.py

@@ -2021,16 +2021,16 @@ def DeduplicateGradientSlices(self, g, aggregator='sum'):
             raise ValueError('{} is not supported'.format(aggregator))
         return GradientSlice(indices=unique, values=new_g)
 
-    def RunAllOnGPU(self, gpu_id=0, use_gpu_engine=False):
+    def RunAllOnGPU(self, gpu_id=0, use_cudnn=False):
         """A convenient function to run everything on the GPU."""
         device_option = caffe2_pb2.DeviceOption()
         device_option.device_type = caffe2_pb2.HIP if workspace.has_hip_support else caffe2_pb2.CUDA
         device_option.cuda_gpu_id = gpu_id
         device_option.hip_gpu_id = gpu_id
         self._net.device_option.CopyFrom(device_option)
-        if use_gpu_engine:
+        if use_cudnn:
             for op in self._net.op:
-                op.engine = "MIOPEN" if workspace.has_hip_support else 'CUDNN'
+                op.engine = 'CUDNN'
 
     def RunAllOnMKL(self):
         """A convenient function to run everything using MKLDNN."""

caffe2/python/core_gradients_test.py

@@ -9,7 +9,7 @@
 import unittest
 
 from caffe2.proto import caffe2_pb2
-from caffe2.python import core, test_util
+from caffe2.python import core, test_util, workspace
 from caffe2.python.core import CreateOperator, GradientRegistry
 from caffe2.python import workspace
 
@@ -94,7 +94,7 @@ def assertOperatorListEqual(self, operatorDefList1, operatorDefList2):
 
     @given(device_option=st.sampled_from([
         None,
-        core.DeviceOption(caffe2_pb2.CUDA, 1)]))
+        core.DeviceOption(caffe2_pb2.HIP, hip_gpu_id=1) if workspace.has_hip_support else core.DeviceOption(caffe2_pb2.CUDA, cuda_gpu_id=1)]))
     def testDirect(self, device_option):
         operators = [
             CreateOperator('Direct', 'in', 'hidden'),
@@ -279,7 +279,7 @@ def testUseInputButInputHasBeenChanged(self):
 
     @given(device_option=st.sampled_from([
         None,
-        core.DeviceOption(caffe2_pb2.CUDA, 1)]))
+        core.DeviceOption(caffe2_pb2.HIP, hip_gpu_id=1) if workspace.has_hip_support else core.DeviceOption(caffe2_pb2.CUDA, cuda_gpu_id=1)]))
     def testMultiUseInput(self, device_option):
         """Test gradient for the following case: