OpenMotionLab/MotionGPT

1import torch2import torch.nn as nn3from torch.nn.utils.rnn import pack_padded_sequence4 5 6class MovementConvEncoder(nn.Module):7    def __init__(self, input_size, hidden_size, output_size):8        super(MovementConvEncoder, self).__init__()9        self.main = nn.Sequential(10            nn.Conv1d(input_size, hidden_size, 4, 2, 1),11            nn.Dropout(0.2, inplace=True),12            nn.LeakyReLU(0.2, inplace=True),13            nn.Conv1d(hidden_size, output_size, 4, 2, 1),14            nn.Dropout(0.2, inplace=True),15            nn.LeakyReLU(0.2, inplace=True),16        )17        self.out_net = nn.Linear(output_size, output_size)18        # self.main.apply(init_weight)19        # self.out_net.apply(init_weight)20 21    def forward(self, inputs):22        inputs = inputs.permute(0, 2, 1)23        outputs = self.main(inputs).permute(0, 2, 1)24        # print(outputs.shape)25        return self.out_net(outputs)26 27 28class MotionEncoderBiGRUCo(nn.Module):29    def __init__(self, input_size, hidden_size, output_size):30        super(MotionEncoderBiGRUCo, self).__init__()31 32        self.input_emb = nn.Linear(input_size, hidden_size)33        self.gru = nn.GRU(34            hidden_size, hidden_size, batch_first=True, bidirectional=True35        )36        self.output_net = nn.Sequential(37            nn.Linear(hidden_size * 2, hidden_size),38            nn.LayerNorm(hidden_size),39            nn.LeakyReLU(0.2, inplace=True),40            nn.Linear(hidden_size, output_size),41        )42 43        # self.input_emb.apply(init_weight)44        # self.output_net.apply(init_weight)45        self.hidden_size = hidden_size46        self.hidden = nn.Parameter(47            torch.randn((2, 1, self.hidden_size), requires_grad=True)48        )49 50    # input(batch_size, seq_len, dim)51    def forward(self, inputs, m_lens):52        num_samples = inputs.shape[0]53 54        input_embs = self.input_emb(inputs)55        hidden = self.hidden.repeat(1, num_samples, 1)56 57        cap_lens = m_lens.data.tolist()58        59        # emb = pack_padded_sequence(input=input_embs, lengths=cap_lens, batch_first=True)60        emb = input_embs61 62        gru_seq, gru_last = self.gru(emb, hidden)63 64        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)65 66        return self.output_net(gru_last)67 68 69class TextEncoderBiGRUCo(nn.Module):70    def __init__(self, word_size, pos_size, hidden_size, output_size):71        super(TextEncoderBiGRUCo, self).__init__()72 73        self.pos_emb = nn.Linear(pos_size, word_size)74        self.input_emb = nn.Linear(word_size, hidden_size)75        self.gru = nn.GRU(76            hidden_size, hidden_size, batch_first=True, bidirectional=True77        )78        self.output_net = nn.Sequential(79            nn.Linear(hidden_size * 2, hidden_size),80            nn.LayerNorm(hidden_size),81            nn.LeakyReLU(0.2, inplace=True),82            nn.Linear(hidden_size, output_size),83        )84 85        # self.input_emb.apply(init_weight)86        # self.pos_emb.apply(init_weight)87        # self.output_net.apply(init_weight)88        # self.linear2.apply(init_weight)89        # self.batch_size = batch_size90        self.hidden_size = hidden_size91        self.hidden = nn.Parameter(92            torch.randn((2, 1, self.hidden_size), requires_grad=True)93        )94 95    # input(batch_size, seq_len, dim)96    def forward(self, word_embs, pos_onehot, cap_lens):97        num_samples = word_embs.shape[0]98 99        pos_embs = self.pos_emb(pos_onehot)100        inputs = word_embs + pos_embs101        input_embs = self.input_emb(inputs)102        hidden = self.hidden.repeat(1, num_samples, 1)103 104        cap_lens = cap_lens.data.tolist()105        emb = pack_padded_sequence(input=input_embs, lengths=cap_lens, batch_first=True)106 107        gru_seq, gru_last = self.gru(emb, hidden)108 109        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)110 111        return self.output_net(gru_last)112