add data_splitor, modify dataset and encoder

2024-09-04 23:38:30 +08:00
parent e0fb9a7617
commit 129bcb872e
7 changed files with 148 additions and 50 deletions
--- a/modules/pose_encoder.py
+++ b/modules/pose_encoder.py
@@ -7,12 +7,13 @@ class PoseEncoder(nn.Module):
        super(PoseEncoder, self).__init__()
        self.config = config
        pose_dim = config["pose_dim"]
+        out_dim = config["out_dim"]
        self.act = nn.ReLU(True)

        self.pose_encoder = nn.Sequential(
-            nn.Linear(pose_dim, 256),
+            nn.Linear(pose_dim, out_dim),
            self.act,
-            nn.Linear(256, 256),
+            nn.Linear(out_dim, out_dim),
            self.act,
        )