Merge pull request openvinotoolkit#67 from bszmelcz/add_aten_adaptive…

…_maxpool2d

Add aten adaptive maxpool2d

src/frontends/pytorch/src/op/adaptive_max_pool2d.cpp

@@ -14,13 +14,8 @@ namespace op {
 OutputVector translate_adaptive_max_pool2d(NodeContext& context) {
     auto x = context.get_input(0);
     auto y = context.get_input(1);
-    auto adaptive_max_pool = context.mark_node(std::make_shared<opset8::AdaptiveMaxPool>(x, y));
-    auto return_indices = context.const_input<bool>(2);
-    OutputVector res{adaptive_max_pool->output(0)};
-    if (return_indices) {
-        res.push_back(adaptive_max_pool->output(1));
-    }
-    return res;
+    auto adaptive_max_pool = context.mark_node(std::make_shared<opset8::AdaptiveMaxPool>(x, y, ov::element::i32));
+    return {adaptive_max_pool->output(0), adaptive_max_pool->output(1)};
 };
 
 }  // namespace op

tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py

@@ -0,0 +1,51 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from pytorch_layer_test_class import PytorchLayerTest
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TestAdaptiveMaxPool2D(PytorchLayerTest):
+
+    def _prepare_input(self):
+        return (self.input_tensor,)
+
+    def create_model(self, output_size=None, return_indices=False):
+
+        class aten_adaptive_max_pool2d(torch.nn.Module):
+
+            def __init__(self, output_size=None, return_indices=False) -> None:
+                super().__init__()
+                self.output_size = output_size 
+                self.return_indices = return_indices
+
+            def forward(self, input_tensor):
+                if self.return_indices:
+                    output, indices = F.adaptive_max_pool2d(input_tensor, self.output_size, True)
+                    return output
+                return F.adaptive_max_pool2d(input_tensor, self.output_size, False)
+
+        ref_net = None
+
+        return aten_adaptive_max_pool2d(output_size, return_indices), ref_net, "aten::adaptive_max_pool2d"
+
+    @pytest.mark.parametrize('input_tensor', ([
+        np.random.randn(1, 1, 4, 4).astype(np.float32),
+        np.random.randn(1, 3, 32, 32).astype(np.float32)
+    ]))
+    @pytest.mark.parametrize('output_size', ([
+        [2, 2],
+        [4, 4],
+    ]))
+    @pytest.mark.parametrize('return_indices', ([
+        False,
+        True,
+    ]))
+    @pytest.mark.nightly
+    def test_adaptive_max_pool2d(self, ie_device, precision, ir_version, input_tensor, output_size, return_indices):
+        self.input_tensor = input_tensor
+        self._test(*self.create_model(output_size, return_indices), ie_device, precision, ir_version)