| 351 | |
| 352 | |
| 353 | class MotionEncoderBiGRUCo(nn.Module): |
| 354 | def __init__(self, input_size, hidden_size, output_size, device): |
| 355 | super(MotionEncoderBiGRUCo, self).__init__() |
| 356 | self.device = device |
| 357 | |
| 358 | self.input_emb = nn.Linear(input_size, hidden_size) |
| 359 | self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True) |
| 360 | self.output_net = nn.Sequential( |
| 361 | nn.Linear(hidden_size*2, hidden_size), |
| 362 | nn.LayerNorm(hidden_size), |
| 363 | nn.LeakyReLU(0.2, inplace=True), |
| 364 | nn.Linear(hidden_size, output_size) |
| 365 | ) |
| 366 | |
| 367 | self.input_emb.apply(init_weight) |
| 368 | self.output_net.apply(init_weight) |
| 369 | self.hidden_size = hidden_size |
| 370 | self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True)) |
| 371 | |
| 372 | # input(batch_size, seq_len, dim) |
| 373 | def forward(self, inputs, m_lens): |
| 374 | num_samples = inputs.shape[0] |
| 375 | |
| 376 | input_embs = self.input_emb(inputs) |
| 377 | hidden = self.hidden.repeat(1, num_samples, 1) |
| 378 | |
| 379 | cap_lens = m_lens.data.tolist() |
| 380 | emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True) |
| 381 | |
| 382 | gru_seq, gru_last = self.gru(emb, hidden) |
| 383 | |
| 384 | gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1) |
| 385 | |
| 386 | return self.output_net(gru_last) |
| 387 | |
| 388 | |
| 389 | class MotionLenEstimatorBiGRU(nn.Module): |