add test attention.py

Author: huangshiyu13

openrl/configs/config.py

@@ -618,7 +618,7 @@ def create_config_parser():
     )
     parser.add_argument(
         "--use_average_pool",
-        action="store_false",
+        type=bool,
         default=True,
         help="by default True, use average pooling for attn model.",
     )

openrl/modules/networks/utils/attention.py

@@ -234,10 +234,13 @@ def forward(self, x, self_idx=-1):
             K = self.split_shape[i][0]
             L = self.split_shape[i][1]
             for j in range(K):
-                torch.cat((x[i][:, (L * j) : (L * j + L)], self_x), dim=-1)
-                exec("x1.append(self.fc_{}(temp))".format(i))
-        x[self_idx]
-        exec("x1.append(self.fc_{}(temp))".format(N - 1))
+                # torch.cat((x[i][:, (L * j) : (L * j + L)], self_x), dim=-1)
+                # exec("x1.append(self.fc_{}(temp))".format(i))
+                temp = torch.cat((x[i][:, (L * j) : (L * j + L)], self_x), dim=-1)
+                x1.append(getattr(self, "fc_" + str(i))(temp))
+        x1.append(getattr(self, "fc_" + str(N - 1))(self_x))
+        # x[self_idx]
+        # exec("x1.append(self.fc_{}(temp))".format(N - 1))
 
         out = torch.stack(x1, 1)
 
@@ -278,8 +281,10 @@ def forward(self, x, self_idx=None):
             K = self.split_shape[i][0]
             L = self.split_shape[i][1]
             for j in range(K):
-                x[i][:, (L * j) : (L * j + L)]
-                exec("x1.append(self.fc_{}(temp))".format(i))
+                # x[i][:, (L * j) : (L * j + L)]
+                # exec("x1.append(self.fc_{}(temp))".format(i))
+                temp = x[i][:, (L * j) : (L * j + L)]
+                x1.append(getattr(self, "fc_" + str(i))(temp))
 
         out = torch.stack(x1, 1)
 

tests/test_modules/test_networks/test_attention.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 2023 The OpenRL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""""""
+
+import os
+import sys
+
+import pytest
+import torch
+
+from openrl.configs.config import create_config_parser
+from openrl.modules.networks.utils.attention import Encoder
+
+
+@pytest.fixture(
+    scope="module", params=["--use_average_pool True", "--use_average_pool False"]
+)
+def config(request):
+    cfg_parser = create_config_parser()
+    cfg = cfg_parser.parse_args(request.param.split())
+    return cfg
+
+
+@pytest.mark.unittest
+def test_attention(config):
+    for cat_self in [False, True]:
+        net = Encoder(cfg=config, split_shape=[[1, 1], [1, 1]], cat_self=cat_self)
+        net(torch.zeros((1, 1)))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", os.path.basename(__file__)]))