Fix test cases

Signed-off-by: AhnDW <[email protected]>

Author: zsef123

core/conversion/converters/impl/batch_norm.cpp

@@ -101,7 +101,6 @@ auto batch_norm_registrations TRTORCH_UNUSED =
               // track_running_stats=True
               LOG_DEBUG("Args[3] running_mean : " << args[3].isIValue() << " / " << args[3].IValue()->isNone());
               LOG_DEBUG("Args[4] running_var : " << args[4].isIValue() << " / " << args[4].IValue()->isNone());
-              
               LOG_DEBUG("use_input_stats, momemtum, cudnn_enabled disregarded");
               LOG_DEBUG("ctx->input_is_dynamic : " << ctx->input_is_dynamic);
 
@@ -112,15 +111,15 @@ auto batch_norm_registrations TRTORCH_UNUSED =
               }
 
               auto eps = static_cast<float>(args[7].unwrapToDouble(1e-5f));
-
+              
               auto scales = args[1].unwrapToTensor(at::ones(shape[1], options)).cpu().contiguous();
               auto bias = args[2].unwrapToTensor(at::zeros(shape[1], options)).cpu().contiguous();
 
               // track_running_stats=True
               if (!args[3].IValue()->isNone() || !args[4].IValue()->isNone()) {
-                auto running_mean = args[3].unwrapToTensor().cpu().contiguous();
-                auto running_var = args[4].unwrapToTensor().cpu().contiguous();
-                _batch_norm(ctx, n, input, orig_shape, scales, bias, running_mean, running_var, eps);
+                auto running_mean = args[3].unwrapToTensor();
+                auto running_var = args[4].unwrapToTensor();
+                _batch_norm(ctx, n, input, orig_shape, scales.to(running_mean.options()), bias.to(running_mean.options()), running_mean, running_var, eps);
                 return true;
               }
 

tests/core/conversion/converters/test_instance_norm.cpp

@@ -11,33 +11,41 @@
 //     const c10::optional<Tensor>& running_mean_opt /* optional */,
 //     const c10::optional<Tensor>& running_var_opt /* optional */,
 //     bool use_input_stats, double momentum, double eps, bool cudnn_enabled) 
-inline constexpr auto graph = R"IR(
-      graph(%0 : Tensor,
-            %1 : Tensor?,
-            %2 : Tensor?,
-            %3 : Tensor?,
-            %4 : Tensor?,
-            %5 : bool):
-        %9 : bool = prim::Constant[value=0]()
-        %6 : float = prim::Constant[value=0.10000000000000001]()
-        %7 : float = prim::Constant[value=1.0000000000000001e-05]()
-        %8 : Tensor = aten::instance_norm(%0, %1, %2, %3, %4, %5, %6, %7, %9)
-        return (%8)
+constexpr auto graph = R"IR(
+      graph(%input.1 : Tensor,
+            %weight.1 : Tensor?,
+            %bias.1 : Tensor?,
+            %running_mean.1 : Tensor?,
+            %running_var.1 : Tensor?,
+            %use_input_stats.1 : bool):
+        %cudnn_enabled.1 : bool = prim::Constant[value=1]()
+        %momentum.1 : float = prim::Constant[value=0.10000000000000001]()
+        %eps.1 : float = prim::Constant[value=1.0000000000000001e-05]()
+        %4 : Tensor = aten::instance_norm(%input.1,
+          %weight.1, %bias.1,
+          %running_mean.1, %running_var.1,
+          %use_input_stats.1, %momentum.1, %eps.1, %cudnn_enabled.1)
+        return (%4)
 )IR";
 
+
 TEST(Converters, ATenInstanceNormConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph, g.get());
 
   auto in = at::randint(1, 10, {1, 5, 5, 5}, {at::kCUDA});
   torch::jit::IValue weight, bias, mean, var; // NoneType
-  bool use_input_stats = true;
+  // https://github.com/pytorch/pytorch/blob/79693bb86a3f601a5c0d3da52d99acec95bb48c1/torch/nn/modules/instancenorm.py#L59
+  const bool use_input_stats = true; 
+  
+  auto trt_in = at::clone(in);
+  torch::jit::IValue trt_weight, trt_bias, trt_mean, trt_var;
 
   auto params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
   auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
 
-  params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
-  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {trt_weight, trt_bias, trt_mean, trt_var, use_input_stats});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }
@@ -50,37 +58,44 @@ TEST(Converters, ATenInstanceNormAffineConvertsCorrectly) {
 
   auto weight = at::randn({in.size(1)}).to(at::kCUDA);
   auto bias = at::randn({in.size(1)}).to(at::kCUDA);
+    
+  torch::jit::IValue mean, var; // NoneType  
+  const bool use_input_stats = true;
 
-  torch::jit::IValue mean, var; // NoneType
-  bool use_input_stats = true;
+  auto trt_in = at::clone(in);
+  auto trt_weight = at::clone(weight);
+  auto trt_bias = at::clone(bias);
+  torch::jit::IValue trt_mean, trt_var;
 
   auto params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
   auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
 
-  params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
-  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {trt_weight, trt_bias, trt_mean, trt_var, use_input_stats});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }
 
-
 TEST(Converters, ATenInstanceNormRunningStatsConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph, g.get());
 
-  auto in = at::randint(1, 10, {1, 5, 5, 5}, {at::kCUDA});
-
-  torch::jit::IValue weight, bias; // NoneType
-
-  auto mean = at::randn({in.size(1)}).to(at::kCUDA);
-  auto var = at::randn({in.size(1)}).to(at::kCUDA);
-  bool use_input_stats = false;
+  auto in = at::randn({1, 5, 5, 5}, {at::kCUDA});
+  
+  torch::jit::IValue weight, bias;
+  auto mean = at::zeros({in.size(1)}, {at::kCUDA});
+  auto var = at::ones({in.size(1)}, {at::kCUDA});
+  const bool use_input_stats = false;
+  
+  auto trt_in = at::clone(in);
+  torch::jit::IValue trt_weight, trt_bias;
+  auto trt_mean = at::clone(mean);
+  auto trt_var = at::clone(var);
 
   auto params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
   auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
 
-  params = trtorch::core::conversion::get_named_params(g->inputs(), {weight, bias, mean, var, use_input_stats});
-  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
-
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {trt_weight, trt_bias, trt_mean, trt_var, use_input_stats});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }