update modules and pipeline

This commit is contained in:
hofee
2024-08-21 17:57:52 +08:00
parent 913d4e521d
commit 837e1c870a
8 changed files with 129 additions and 57 deletions

View File

@@ -0,0 +1,20 @@
from torch import nn
import PytorchBoot.stereotype as stereotype
@stereotype.module("pose_encoder")
class PoseEncoder(nn.Module):
def __init__(self, config):
super(PoseEncoder, self).__init__()
self.config = config
pose_dim = config["pose_dim"]
self.act = nn.ReLU(True)
self.pose_encoder = nn.Sequential(
nn.Linear(pose_dim, 256),
self.act,
nn.Linear(256, 256),
self.act,
)
def encode_pose(self, pose):
return self.pose_encoder(pose)