| 40 | |
| 41 | |
| 42 | class PoseEncoderConv(nn.Module): |
| 43 | def __init__(self, length, dim): |
| 44 | super().__init__() |
| 45 | |
| 46 | self.net = nn.Sequential( |
| 47 | ConvNormRelu(dim, 32, batchnorm=True), |
| 48 | ConvNormRelu(32, 64, batchnorm=True), |
| 49 | ConvNormRelu(64, 64, True, batchnorm=True), |
| 50 | nn.Conv1d(64, 32, 3) |
| 51 | ) |
| 52 | |
| 53 | self.out_net = nn.Sequential( |
| 54 | # nn.Linear(864, 256), # for 64 frames |
| 55 | nn.Linear(384, 256), # for 34 frames |
| 56 | nn.BatchNorm1d(256), |
| 57 | nn.LeakyReLU(True), |
| 58 | nn.Linear(256, 128), |
| 59 | nn.BatchNorm1d(128), |
| 60 | nn.LeakyReLU(True), |
| 61 | nn.Linear(128, 32), |
| 62 | ) |
| 63 | |
| 64 | self.fc_mu = nn.Linear(32, 32) |
| 65 | self.fc_logvar = nn.Linear(32, 32) |
| 66 | |
| 67 | def forward(self, poses, variational_encoding): |
| 68 | # encode |
| 69 | poses = poses.transpose(1, 2) # to (bs, dim, seq) |
| 70 | out = self.net(poses) |
| 71 | out = out.flatten(1) |
| 72 | out = self.out_net(out) |
| 73 | |
| 74 | # return out, None, None |
| 75 | mu = self.fc_mu(out) |
| 76 | logvar = self.fc_logvar(out) |
| 77 | |
| 78 | if variational_encoding: |
| 79 | z = reparameterize(mu, logvar) |
| 80 | else: |
| 81 | z = mu |
| 82 | return z, mu, logvar |
| 83 | |
| 84 | |
| 85 | class PoseDecoderFC(nn.Module): |
no outgoing calls
no test coverage detected