rnn.py

可选 recurrent 层

原文件 · 80 行 · 2929 字节

algorithms/utils/rnn.py
SHA-256: bb6ca30a91503f012df3c45d9be9011642f9895fb9804b05530bcb49db47fe50

1import torch2import torch.nn as nn34"""RNN modules."""567class RNNLayer(nn.Module):8    def __init__(self, inputs_dim, outputs_dim, recurrent_N, use_orthogonal):9        super(RNNLayer, self).__init__()10        self._recurrent_N = recurrent_N11        self._use_orthogonal = use_orthogonal1213        self.rnn = nn.GRU(inputs_dim, outputs_dim, num_layers=self._recurrent_N)14        for name, param in self.rnn.named_parameters():15            if 'bias' in name:16                nn.init.constant_(param, 0)17            elif 'weight' in name:18                if self._use_orthogonal:19                    nn.init.orthogonal_(param)20                else:21                    nn.init.xavier_uniform_(param)22        self.norm = nn.LayerNorm(outputs_dim)2324    def forward(self, x, hxs, masks):25        if x.size(0) == hxs.size(0):26            x, hxs = self.rnn(x.unsqueeze(0),27                              (hxs * masks.repeat(1, self._recurrent_N).unsqueeze(-1)).transpose(0, 1).contiguous())28            x = x.squeeze(0)29            hxs = hxs.transpose(0, 1)30        else:31            # x is a (T, N, -1) tensor that has been flatten to (T * N, -1)32            N = hxs.size(0)33            T = int(x.size(0) / N)3435            # unflatten36            x = x.view(T, N, x.size(1))3738            # Same deal with masks39            masks = masks.view(T, N)4041            # Let's figure out which steps in the sequence have a zero for any agent42            # We will always assume t=0 has a zero in it as that makes the logic cleaner43            has_zeros = ((masks[1:] == 0.0)44                         .any(dim=-1)45                         .nonzero()46                         .squeeze()47                         .cpu())4849            # +1 to correct the masks[1:]50            if has_zeros.dim() == 0:51                # Deal with scalar52                has_zeros = [has_zeros.item() + 1]53            else:54                has_zeros = (has_zeros + 1).numpy().tolist()5556            # add t=0 and t=T to the list57            has_zeros = [0] + has_zeros + [T]5859            hxs = hxs.transpose(0, 1)6061            outputs = []62            for i in range(len(has_zeros) - 1):63                # We can now process steps that don't have any zeros in masks together!64                # This is much faster65                start_idx = has_zeros[i]66                end_idx = has_zeros[i + 1]67                temp = (hxs * masks[start_idx].view(1, -1, 1).repeat(self._recurrent_N, 1, 1)).contiguous()68                rnn_scores, hxs = self.rnn(x[start_idx:end_idx], temp)69                outputs.append(rnn_scores)7071            # assert len(outputs) == T72            # x is a (T, N, -1) tensor73            x = torch.cat(outputs, dim=0)7475            # flatten76            x = x.reshape(T * N, -1)77            hxs = hxs.transpose(0, 1)7879        x = self.norm(x)80        return x, hxs