21 lines
549 B
Python
21 lines
549 B
Python
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)
|