Merge pull request #602 from SrivastavaKshitij/fix_test_and_doc

Test enhancement and doc correction for QAT

Author: jaybdub

examples/contrib/quantization_aware_training/README.md

@@ -25,27 +25,9 @@ RUN add-apt-repository ppa:git-core/ppa && \
 
 RUN pip install termcolor graphviz
 
-## If you have followed instructions on main README.md file to install torch2trt using scripts/build_contrib.sh
-## You dont require rest of the steps
-
-RUN git clone https://github.com/NVIDIA/TensorRT.git /sw/TensorRT/
-
-##Make sure that patch file is under the same folder where dockerfile is being called
-
-ADD pytorch_nvidia_quantization.patch /sw/TensorRT
-
-RUN cd /sw/TensorRT/ && \
-    git sparse-checkout init --cone && \
-    git sparse-checkout set /tools/pytorch-quantization/ && \
-    git apply --reject --whitespace=fix pytorch_nvidia_quantization.patch && \
-    cd tools/pytorch-quantization/ && \
-    python setup.py install 
-
-RUN git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git /sw/TensorRT/ && \
-    cd /sw/TensorRT/ && \
-    git fetch origin pull/514/head:PR514 && \
-    git checkout PR514 && \
-    python setup.py install --plugins
+RUN git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git /sw/torch2trt/ && \
+    cd /sw/torch2trt/scripts && \
+	bash build_contrib.sh
 
 ```
 

torch2trt/converters/interpolate.py

@@ -92,34 +92,36 @@ def convert_interpolate_trt7(ctx):
 
 
 class Interpolate(torch.nn.Module):
-    def __init__(self, size, mode, align_corners):
+    def __init__(self, size=None,scale_factor=None, mode=None, align_corners=None):
         super(Interpolate, self).__init__()
+		## Use either size or scale factor. 
         self.size = size
+        self.scale_factor = scale_factor 
         self.mode = mode
         self.align_corners = align_corners
 
     def forward(self, x):
-        return F.interpolate(x, self.size, mode=self.mode, align_corners=self.align_corners)
+        return F.interpolate(x, size=self.size, scale_factor=self.scale_factor,mode=self.mode, align_corners=self.align_corners)
 
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 112, 112)], enabled=trt_version() < '7.1' and has_interpolate_plugin())
 def test_interpolate_nearest():
-    return Interpolate((224, 224), 'nearest', None)
+    return Interpolate(size=(224, 224), mode='nearest', align_corners=None)
 
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 112, 112)], enabled=trt_version() < '7.1' and has_interpolate_plugin())
 def test_interpolate_bilinear():
-    return Interpolate((224, 224), 'bilinear', False)
+    return Interpolate(size=(224, 224), mode= 'bilinear', align_corners=False)
 
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 112, 112)], enabled=trt_version() < '7.1' and has_interpolate_plugin())
 def test_interpolate_bicubic():
-    return Interpolate((224, 224), 'bicubic', False)
+    return Interpolate(size=(224, 224), mode='bicubic',align_corners= False)
 
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 112, 112)], enabled=trt_version() < '7.1' and has_interpolate_plugin())
 def test_interpolate_area():
-    return Interpolate((56, 56), 'area', None)
+    return Interpolate(size=(56, 56), mode='area',align_corners= None)
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 112, 112)], enabled=trt_version() < '7.1' and has_interpolate_plugin())
 def test_upsample_scale_factor2():
@@ -135,7 +137,11 @@ def test_bilinear_mode():
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1,3,12,12)], enabled=trt_version() >= '7.1')
 def test_align_corner():
-    return torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
+    return torch.nn.Upsample(scale_factor=2.0, mode="bilinear", align_corners=True)
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1,3,12,12)], enabled=trt_version() >= '7.1')
+def test_align_corner_functional():
+    return Interpolate(scale_factor=2.0, mode="bilinear", align_corners=True)
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1,5,13,13)], enabled=trt_version() >= '7.1')
 def test_bilinear_mode_odd_input_shape():

torch2trt/test.py

@@ -6,6 +6,22 @@
 import runpy
 import traceback
 from termcolor import colored
+import math
+import numpy as np
+
+def pSNR(model_op,trt_op):
+    #model_op = model_op.cpu().detach().numpy().flatten()
+    #trt_op = trt_op.cpu().detach().numpy().flatten()
+
+    # Calculating Mean Squared Error
+    mse = np.sum(np.square(model_op - trt_op)) / len(model_op)
+    # Calcuating peak signal to noise ratio
+    try:
+    	psnr_db = 20 * math.log10(np.max(abs(model_op))) - 10 * math.log10(mse)
+    except:
+        psnr_db = np.nan
+    return mse,psnr_db
+
 
 
 def run(self):
@@ -49,6 +65,24 @@ def run(self):
 
         if max_error_i > max_error:
             max_error = max_error_i
+
+	## calculate peak signal to noise ratio
+    assert(len(outputs) == len(outputs_trt))
+	
+    ## Check if output is boolean
+    # if yes, then dont calculate psnr
+    if outputs[0].dtype == torch.bool:
+        mse = np.nan
+        psnr_db = np.nan
+    else:
+        model_op = []
+        trt_op = []
+        for i in range(len(outputs)):
+            model_op.extend(outputs[i].detach().cpu().numpy().flatten())
+            trt_op.extend(outputs_trt[i].detach().cpu().numpy().flatten())
+        model_op = np.array(model_op)
+        trt_op = np.array(trt_op)
+        mse,psnr_db = pSNR(model_op,trt_op)
 
     # benchmark pytorch throughput
     torch.cuda.current_stream().synchronize()
@@ -90,7 +124,7 @@ def run(self):
 
     ms_trt = 1000.0 * (t1 - t0) / 50.0
 
-    return max_error, fps, fps_trt, ms, ms_trt
+    return max_error,psnr_db,mse, fps, fps_trt, ms, ms_trt
 
 
 if __name__ == '__main__':
@@ -106,7 +140,7 @@ def run(self):
     for include in args.include:
         runpy.run_module(include)
 
-    num_tests, num_success, num_tolerance, num_error = 0, 0, 0, 0
+    num_tests, num_success, num_tolerance, num_error, num_tolerance_psnr = 0, 0, 0, 0, 0
     for test in MODULE_TESTS:
 
         # filter by module name
@@ -120,14 +154,17 @@ def run(self):
             if args.use_onnx:
                 test.torch2trt_kwargs.update({'use_onnx': True})
 
-            max_error, fps, fps_trt, ms, ms_trt = run(test)
+            max_error,psnr_db,mse, fps, fps_trt, ms, ms_trt = run(test)
 
             # write entry
-            line = '| %s | %s | %s | %s | %.2E | %.3g | %.3g | %.3g | %.3g |' % (name, test.dtype.__repr__().split('.')[-1], str(test.input_shapes), str(test.torch2trt_kwargs), max_error, fps, fps_trt, ms, ms_trt)
+            line = '| %70s | %s | %25s | %s | %.2E | %.2f | %.2E | %.3g | %.3g | %.3g | %.3g |' % (name, test.dtype.__repr__().split('.')[-1], str(test.input_shapes), str(test.torch2trt_kwargs), max_error,psnr_db,mse, fps, fps_trt, ms, ms_trt)
 
             if args.tolerance >= 0 and max_error > args.tolerance:
                 print(colored(line, 'yellow'))
                 num_tolerance += 1
+            elif psnr_db < 100:
+                print(colored(line, 'magenta'))
+                num_tolerance_psnr +=1
             else:
                 print(line)
             num_success += 1
@@ -144,4 +181,5 @@ def run(self):
     print('NUM_TESTS: %d' % num_tests)
     print('NUM_SUCCESSFUL_CONVERSION: %d' % num_success)
     print('NUM_FAILED_CONVERSION: %d' % num_error)
-    print('NUM_ABOVE_TOLERANCE: %d' % num_tolerance)
+    print('NUM_ABOVE_TOLERANCE: %d' % num_tolerance)
+    print('NUM_pSNR_TOLERANCE: %d' %num_tolerance_psnr)