Python torch.nn.GRU Examples

def __init__(self, batch=32, bidirectional=True):

        super(BLSTM, self).__init__()
        
        self.bidirectional = bidirectional
        self.hidden = self.init_hidden(batch)
        self.lstm = nn.LSTM(257, 50, num_layers=2, bidirectional=True)  
        self.fc = nn.Linear(50*2,1)

        ## Weights initialization
        def _weights_init(m):
            if isinstance(m, nn.Conv2d or nn.Linear or nn.GRU or nn.LSTM):
                init.xavier_normal_(m.weight)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d or nn.BatchNorm1d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
        
        self.apply(_weights_init) 

def whatCellType(input_size, hidden_size, cell_type, dropout_rate):
    if cell_type == 'rnn':
        cell = nn.RNN(input_size, hidden_size, dropout=dropout_rate, batch_first=False)
        init_gru(cell)
        return cell
    elif cell_type == 'gru':
        cell = nn.GRU(input_size, hidden_size, dropout=dropout_rate, batch_first=False)
        init_gru(cell)
        return cell
    elif cell_type == 'lstm':
        cell = nn.LSTM(input_size, hidden_size, dropout=dropout_rate, batch_first=False)
        init_lstm(cell)
        return cell
    elif cell_type == 'bigru':
        cell = nn.GRU(input_size, hidden_size, bidirectional=True, dropout=dropout_rate, batch_first=False)
        init_gru(cell)
        return cell
    elif cell_type == 'bilstm':
        cell = nn.LSTM(input_size, hidden_size, bidirectional=True, dropout=dropout_rate, batch_first=False)
        init_lstm(cell)
        return cell 

def __init__(self, num_nodes, num_rels, hidden_channels, seq_len,
                 num_layers=1, dropout=0., bias=True):
        super(RENet, self).__init__()

        self.num_nodes = num_nodes
        self.hidden_channels = hidden_channels
        self.num_rels = num_rels
        self.seq_len = seq_len
        self.dropout = dropout

        self.ent = Parameter(torch.Tensor(num_nodes, hidden_channels))
        self.rel = Parameter(torch.Tensor(num_rels, hidden_channels))

        self.sub_gru = GRU(3 * hidden_channels, hidden_channels, num_layers,
                           batch_first=True, bias=bias)
        self.obj_gru = GRU(3 * hidden_channels, hidden_channels, num_layers,
                           batch_first=True, bias=bias)

        self.sub_lin = Linear(3 * hidden_channels, num_nodes, bias=bias)
        self.obj_lin = Linear(3 * hidden_channels, num_nodes, bias=bias)

        self.reset_parameters() 

def __init__(self):
        super().__init__()
        self.prenet = PreNet(in_features=hp.E)  # [N, T, E//2]

        self.conv1d_bank = Conv1dBank(K=hp.K, in_channels=hp.E // 2, out_channels=hp.E // 2)  # [N, T, E//2 * K]

        self.conv1d_1 = Conv1d(in_channels=hp.K * hp.E // 2, out_channels=hp.E // 2, kernel_size=3)  # [N, T, E//2]
        self.conv1d_2 = Conv1d(in_channels=hp.E // 2, out_channels=hp.E // 2, kernel_size=3)  # [N, T, E//2]
        self.bn1 = BatchNorm1d(num_features=hp.E // 2)
        self.bn2 = BatchNorm1d(num_features=hp.E // 2)

        self.highways = nn.ModuleList()
        for i in range(hp.num_highways):
            self.highways.append(Highway(in_features=hp.E // 2, out_features=hp.E // 2))

        self.gru = nn.GRU(input_size=hp.E // 2, hidden_size=hp.E // 2, num_layers=2, bidirectional=True, batch_first=True) 

def __init__(self, vocab_size, word_dim, embed_size, num_layers, pooling='last',
                 use_abs=False, bid=False, glove_path='data/glove.pkl'):
        super(EncoderTextGRU, self).__init__()
        self.use_abs = use_abs
        self.embed_size = embed_size
        self.combiner = Combiner(pooling, embed_size)

        # word embedding
        self.word_dim = word_dim
        if word_dim > 300:
            self.embed = nn.Embedding(vocab_size, word_dim-300)
        _, embed_weight = pickle.load(open(glove_path, 'rb'))
        self.glove = Variable(torch.cuda.FloatTensor(embed_weight), requires_grad=False)

        # caption embedding
        self.rnn = nn.GRU(word_dim, embed_size//(2 if bid else 1), num_layers, batch_first=True, bidirectional=bid)

        self.init_weights() 

def make_model(d_vocab, N, d_model, d_ff=1024, h=4, dropout=0.1):
    """Helper: Construct a model from hyperparameters."""
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.GRU(d_model, d_model, 1),
        nn.Sequential(Embeddings(d_model, d_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, d_vocab), c(position)),
        Generator(d_model, d_vocab),
        d_model
    )
    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model 

def __init__(self, encoder, config, device=torch.device('cpu'), wandb=None):
        super().__init__(encoder, wandb, device)
        self.config = config
        for k, v in config.items():
            setattr(self, k, v)

        self.device = device
        self.steps_gen = lambda: range(self.steps_start, self.steps_end, self.steps_step)
        self.discriminators = {i: nn.Linear(self.gru_size, self.encoder.hidden_size).to(device) for i in self.steps_gen()}
        self.gru = nn.GRU(input_size=self.encoder.hidden_size, hidden_size=self.gru_size, num_layers=self.gru_layers, batch_first=True).to(device)
        self.labels = {i: torch.arange(self.batch_size * (self.sequence_length - i - 1)).to(device) for i in self.steps_gen()}
        params = list(self.encoder.parameters()) + list(self.gru.parameters())
        for disc in self.discriminators.values():
          params += disc.parameters()
        self.optimizer = torch.optim.Adam(params, lr=config['lr'])
        self.early_stopper = EarlyStopping(patience=self.patience, verbose=False, wandb=self.wandb, name="encoder") 

def make_model(d_vocab, N, d_model, d_ff=1024, h=4, dropout=0.1):
    """Helper: Construct a model from hyperparameters."""
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.GRU(d_model, d_model, 1),
        nn.Sequential(Embeddings(d_model, d_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, d_vocab), c(position)),
        Generator(d_model, d_vocab),
        d_model
    )
    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model 

def init_hidden(config, batch_size, cell):
    layers = 1
    if isinstance(cell, (nn.LSTM, nn.GRU)):
        layers = cell.num_layers
        if cell.bidirectional:
            layers = layers * 2

    if isinstance(cell, (nn.LSTM, nn.LSTMCell)):
        hidden  = Variable(torch.zeros(layers, batch_size, cell.hidden_size))
        context = Variable(torch.zeros(layers, batch_size, cell.hidden_size))
    
        if config.CONFIG.cuda:
            hidden  = hidden.cuda()
            context = context.cuda()
        return hidden, context

    if isinstance(cell, (nn.GRU, nn.GRUCell)):
        hidden  = Variable(torch.zeros(layers, batch_size, cell.hidden_size))
        if config.CONFIG.cuda:
            hidden  = hidden.cuda()
        return hidden 

def __init__(self):

        super().__init__()
        K = len(hp.ref_enc_filters)
        filters = [1] + hp.ref_enc_filters
        convs = [nn.Conv2d(in_channels=filters[i],
                           out_channels=filters[i + 1],
                           kernel_size=(3, 3),
                           stride=(2, 2),
                           padding=(1, 1)) for i in range(K)]
        self.convs = nn.ModuleList(convs)
        self.bns = nn.ModuleList([nn.BatchNorm2d(num_features=hp.ref_enc_filters[i]) for i in range(K)])

        out_channels = self.calculate_channels(hp.n_mels, 3, 2, 1, K)
        self.gru = nn.GRU(input_size=hp.ref_enc_filters[-1] * out_channels,
                          hidden_size=hp.E // 2,
                          batch_first=True) 

def __init__(self, vocab_size, embed_size, hidden_size, dropout=None, \
        bidirectional=False, shared_embed=None, init_word_embed=None, rnn_type='lstm', use_cuda=True):
        super(EncoderRNN, self).__init__()
        if not rnn_type in ('lstm', 'gru'):
            raise RuntimeError('rnn_type is expected to be lstm or gru, got {}'.format(rnn_type))
        if bidirectional:
            print('[ Using bidirectional {} encoder ]'.format(rnn_type))
        else:
            print('[ Using {} encoder ]'.format(rnn_type))
        if bidirectional and hidden_size % 2 != 0:
            raise RuntimeError('hidden_size is expected to be even in the bidirectional mode!')
        self.dropout = dropout
        self.rnn_type = rnn_type
        self.use_cuda = use_cuda
        self.hidden_size = hidden_size // 2 if bidirectional else hidden_size
        self.num_directions = 2 if bidirectional else 1
        self.embed = shared_embed if shared_embed is not None else nn.Embedding(vocab_size, embed_size, padding_idx=0)
        model = nn.LSTM if rnn_type == 'lstm' else nn.GRU
        self.model = model(embed_size, self.hidden_size, 1, batch_first=True, bidirectional=bidirectional)
        if shared_embed is None:
            self.init_weights(init_word_embed) 

def __init__(self, hidden_size, embed_size, output_size,
                 n_layers=1, dropout=0.1):
        """

        Args:
            hidden_size: GRU hidden_size
            embed_size:  embedding size
            output_size:  outputs vocab size
            n_layers:  GRU layers
            dropout:  dropout ratio,
        """
        super(AttnDecoder, self).__init__()
        # Define parameters
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = dropout
        # Define layers
        self.embedding = nn.Embedding(output_size, embed_size)
        self.dropout_layer = nn.Dropout(dropout)
        self.attn = Attn('concat', hidden_size)
        self.gru = nn.GRU(hidden_size + embed_size, hidden_size, n_layers, dropout=dropout)
        self.out = nn.Linear(hidden_size, output_size) 

def __init__(self, input_size, embed_size, hidden_size,
                 n_layers=1, dropout=0.3):
        """initialize encoder

        Args:
            input_size: <int>, encoder vocab size
            embed_size: <int>, encoder embed size
            hidden_size: <int>, GRU hidden state size
            n_layers: <int>, GRU layers
            dropout: <float>, dropout rate

        Notes:
            default batch_first, bidirectional=True
        """
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.n_layers = n_layers
        self.dropout = self.dropout
        self.embedding = nn.Embedding(input_size, embed_size)
        self.gru = nn.GRU(embed_size, hidden_size, n_layers, bidirectional=True, dropout=dropout) 

def __init__(self, embed_size, hidden_size, vocab_size, dropout_rate):
        super().__init__()
        self.attn_u = Attn(hidden_size)
        self.attn_z = Attn(hidden_size)
        self.gru = nn.GRU(embed_size + hidden_size, hidden_size, dropout=dropout_rate)
        self.ln1 = LayerNormalization(hidden_size)

        self.w1 = nn.Linear(hidden_size, vocab_size)
        self.proj_copy1 = nn.Linear(hidden_size * 2, hidden_size)
        self.v1 = nn.Linear(hidden_size, 1)

        self.proj_copy2 = nn.Linear(hidden_size * 2, hidden_size)
        self.v2 = nn.Linear(hidden_size, 1)
        self.mu = nn.Linear(vocab_size, embed_size)
        self.dropout_rate = dropout_rate

        self.gru = orth_gru(self.gru)

        self.copy_weight = 1 

def __init__(self, input_size, embed_size, hidden_size,
                 n_layers=1, dropout=0.5, batch_first=False, bidirectional=False):
        """

        :param input_size: 词典大小
        :param embed_size:  word2vec嵌入维度
        :param hidden_size:  encoder rnn 隐藏态维度
        :param n_layers:   rnn层数
        :param dropout:    dropout rate
        """
        super(Encoder, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.embed = nn.Embedding(input_size, embed_size)
        self.gru = nn.GRU(embed_size, hidden_size, n_layers,
                          dropout=dropout, bidirectional=bidirectional, batch_first=batch_first) 

def __init__(self, lang, shared_emb, vocab_size, hidden_size, dropout, slots, nb_gate):
        super(Generator, self).__init__()
        self.vocab_size = vocab_size
        self.lang = lang
        self.embedding = shared_emb
        self.dropout_layer = nn.Dropout(dropout)
        self.gru = nn.GRU(hidden_size, hidden_size, dropout=dropout)
        self.nb_gate = nb_gate
        self.hidden_size = hidden_size
        self.W_ratio = nn.Linear(3 * hidden_size, 1)
        self.softmax = nn.Softmax(dim=1)
        self.sigmoid = nn.Sigmoid()
        self.slots = slots

        self.W_gate = nn.Linear(hidden_size, nb_gate)

        # Create independent slot embeddings
        self.slot_w2i = {}
        for slot in self.slots:
            if slot.split("-")[0] not in self.slot_w2i.keys():
                self.slot_w2i[slot.split("-")[0]] = len(self.slot_w2i)
            if slot.split("-")[1] not in self.slot_w2i.keys():
                self.slot_w2i[slot.split("-")[1]] = len(self.slot_w2i)
        self.Slot_emb = nn.Embedding(len(self.slot_w2i), hidden_size)
        self.Slot_emb.weight.data.normal_(0, 0.1) 

def __init__(self, vocab_size, hidden_size, dropout, n_layers=1, data_dir=''):
        super(EncoderRNN, self).__init__()
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.dropout = dropout
        self.dropout_layer = nn.Dropout(dropout)
        self.embedding = nn.Embedding(
            vocab_size, hidden_size, padding_idx=PAD_token)
        self.embedding.weight.data.normal_(0, 0.1)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers,
                          dropout=dropout, bidirectional=True)
        # self.domain_W = nn.Linear(hidden_size, nb_domain)

        if args["load_embedding"]:
            with open(os.path.join(data_dir, 'emb{}.json'.format(vocab_size))) as f:
                E = json.load(f)
            new = self.embedding.weight.data.new
            self.embedding.weight.data.copy_(new(E))
            self.embedding.weight.requires_grad = True
            print("Encoder embedding requires_grad",
                  self.embedding.weight.requires_grad)

        if args["fix_embedding"]:
            self.embedding.weight.requires_grad = False 

def __init__(self):

        super(ConvGRU, self).__init__()

        self.conv1  = nn.Conv2d(1, 16, kernel_size=(5,5), padding=(2,2))
        nn.init.xavier_uniform_(self.conv1.weight)
        self.batch1 = nn.BatchNorm2d(16)
        self.relu1  = nn.ReLU()
        self.conv2  = nn.Conv2d(16, 32, kernel_size=(5,5), padding=(2,2))
        nn.init.xavier_uniform_(self.conv2.weight)
        self.batch2 = nn.BatchNorm2d(32)
        self.relu2  = nn.ReLU()
        self.conv3  = nn.Conv2d(32, 64, kernel_size=(5,5), padding=(2,2))
        self.batch3 = nn.BatchNorm2d(64)
        self.relu3  = nn.ReLU()
        self.gru    = nn.GRU(25700, 1, bidirectional=True)
        self.fc1    = nn.Linear(512, 256)
        nn.init.xavier_uniform_(self.fc1.weight)
        self.drop1  = nn.Dropout(0.5)
        self.relu4  = nn.ReLU()
        self.fc2    = nn.Linear(256, 128)
        nn.init.xavier_uniform_(self.fc2.weight)
        self.drop2  = nn.Dropout(0.5)
        self.relu5  = nn.ReLU()
        self.fc3    = nn.Linear(128, 2) 

def __init__(self, vocab_size, embed_size, hidden_size, input_dropout_p=0, dropout_p=0, n_layers=1, 
                 rnn_cell='GRU', variable_lengths=False, embedding=None, update_embedding=True):
        super(Encoder, self).__init__()
        
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.input_dropout = nn.Dropout(p=input_dropout_p)
        if rnn_cell == 'LSTM':
            self.rnn_cell = nn.LSTM
        elif rnn_cell == 'GRU':
            self.rnn_cell = nn.GRU
        else:
            raise ValueError("Unsupported RNN Cell: {0}".format(rnn_cell))
        
        self.variable_lengths = variable_lengths
        self.embedding = nn.Embedding(vocab_size, embed_size)
        if embedding is not None:
            self.embedding.weight = nn.Parameter(embedding)
        self.embedding.weight.requires_grad = update_embedding
        self.rnn = self.rnn_cell(embed_size, hidden_size, n_layers, batch_first=True, dropout=dropout_p) 

def __init__(self, vocab, embedding_dim, hop, dropout, unk_mask):
        super(DecoderMemNN, self).__init__()
        self.num_vocab = vocab
        self.max_hops = hop
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        self.unk_mask = unk_mask
        for hop in range(self.max_hops+1):
            C = nn.Embedding(self.num_vocab, embedding_dim, padding_idx=PAD_token)
            C.weight.data.normal_(0, 0.1)
            self.add_module("C_{}".format(hop), C)
        self.C = AttrProxy(self, "C_")
        self.softmax = nn.Softmax(dim=1)
        self.W = nn.Linear(embedding_dim,1)
        self.W1 = nn.Linear(2*embedding_dim,self.num_vocab)
        self.gru = nn.GRU(embedding_dim, embedding_dim, dropout=dropout) 

def forward(self, input_seq, input_lengths, hidden=None):
        # Convert word indexes to embeddings
        embedded = self.embedding(input_seq)
        # Pack padded batch of sequences for RNN module
        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
        # Forward pass through GRU
        outputs, hidden = self.gru(packed, hidden)
        # Unpack padding
        outputs, _ = torch.nn.utils.rnn.pad_packed_sequence(outputs)
        # Sum bidirectional GRU outputs
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]
        # Return output and final hidden state
        return outputs, hidden


# Luong attention layer 

def __init__(self, attn_model, hidden_size, output_size, n_layers=1, dropout=0.1):
        super(LuongAttnDecoderRNN, self).__init__()

        # Keep for reference
        self.attn_model = attn_model
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = dropout

        # Define layers
        self.embedding = nn.Embedding(output_size, hidden_size)
        self.embedding_dropout = nn.Dropout(dropout)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout))
        self.concat = nn.Linear(hidden_size * 2, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)

        self.attn = Attn(attn_model, hidden_size) 

def forward(self, input_step, last_hidden, encoder_outputs):
        # Note: we run this one step (word) at a time
        # Get embedding of current input word
        embedded = self.embedding(input_step)
        embedded = self.embedding_dropout(embedded)
        # Forward through unidirectional GRU
        rnn_output, hidden = self.gru(embedded, last_hidden)
        # Calculate attention weights from the current GRU output
        attn_weights = self.attn(rnn_output, encoder_outputs)
        # Multiply attention weights to encoder outputs to get new "weighted sum" context vector
        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))
        # Concatenate weighted context vector and GRU output using Luong eq. 5
        rnn_output = rnn_output.squeeze(0)
        context = context.squeeze(1)
        concat_input = torch.cat((rnn_output, context), 1)
        concat_output = torch.tanh(self.concat(concat_input))
        # Predict next word using Luong eq. 6
        output = self.out(concat_output)
        output = F.softmax(output, dim=1)
        # Return output and final hidden state
        return output, hidden 

def forward(self, inputs, hidden):
        def select_layer(h_state, i):  # To work on both LSTM / GRU, RNN
            if isinstance(h_state, tuple):
                return tuple([select_layer(s, i) for s in h_state])
            else:
                return h_state[i]

        next_hidden = []
        for i, layer in enumerate(self.layers):
            next_hidden_i = layer(inputs, select_layer(hidden, i))
            output = next_hidden_i[0] if isinstance(next_hidden_i, tuple) \
                else next_hidden_i
            if i + 1 < self.num_layers:
                output = self.dropout(output)
            if self.residual and inputs.size(-1) == output.size(-1):
                inputs = output + inputs
            else:
                inputs = output
            next_hidden.append(next_hidden_i)
        if isinstance(hidden, tuple):
            next_hidden = tuple([torch.stack(h) for h in zip(*next_hidden)])
        else:
            next_hidden = torch.stack(next_hidden)
        return inputs, next_hidden 

def __init__(self, input_dim, hidden_dim, layer_num, dropout_rate):
        """Initialize an RNN.
        :param input_dim: dimension of input
        :type input_dim: int
        :param hidden_dim: dimension of filters
        :type hidden_dim: int
        :param layer_num: number of RNN layers
        :type layer_num: int
        :param dropout_rate: dropout rate
        :type dropout_rate: float
        """
        super(RnnModel, self).__init__()
        self.rnn_layer = nn.GRU(
            input_size=input_dim,
            hidden_size=hidden_dim // 2,
            num_layers=layer_num,
            bidirectional=True,
            dropout=dropout_rate,
        ) 

def __init__(self, input_size, hidden_size, n_layers, dropout, n_classes):
    super(GRU, self).__init__()
    self.gru = nn.GRU(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=n_layers,
        batch_first=True,
        dropout=dropout)
    self.fc = nn.Linear(hidden_size, n_classes)

    self.hidden_size = hidden_size
    self.n_layers = n_layers 

def get_gru(input_size, hidden_size, n_layers, dropout, n_classes):
  return GRU(input_size, hidden_size, n_layers, dropout, n_classes) 

def __init__(self):

        super(FConv, self).__init__()

        self.convolution = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=(3,3), padding=(2,2)),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, kernel_size=(3,3), padding=(2,2), dilation=(2,2)),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, kernel_size=(3,3), padding=(2,2), dilation=(4,4)),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, kernel_size=(3,3), padding=(2,2), dilation=(8,8)),
            nn.BatchNorm2d(32),
            nn.ReLU()
        )

        self.fc = nn.Linear(32,1)

        ## Weights initialization
        def _weights_init(m):
            if isinstance(m, nn.Conv2d or nn.Linear or nn.GRU):
                xavier_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d or nn.BatchNorm1d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
        
        self.apply(_weights_init) 

def __init__(self, event_dim, control_dim, init_dim, hidden_dim,
                 gru_layers=3, gru_dropout=0.3):
        super().__init__()

        self.event_dim = event_dim
        self.control_dim = control_dim
        self.init_dim = init_dim
        self.hidden_dim = hidden_dim
        self.gru_layers = gru_layers
        self.concat_dim = event_dim + 1 + control_dim
        self.input_dim = hidden_dim
        self.output_dim = event_dim

        self.primary_event = self.event_dim - 1

        self.inithid_fc = nn.Linear(init_dim, gru_layers * hidden_dim)
        self.inithid_fc_activation = nn.Tanh()

        self.event_embedding = nn.Embedding(event_dim, event_dim)
        self.concat_input_fc = nn.Linear(self.concat_dim, self.input_dim)
        self.concat_input_fc_activation = nn.LeakyReLU(0.1, inplace=True)

        self.gru = nn.GRU(self.input_dim, self.hidden_dim,
                          num_layers=gru_layers, dropout=gru_dropout)
        self.output_fc = nn.Linear(hidden_dim * gru_layers, self.output_dim)
        self.output_fc_activation = nn.Softmax(dim=-1)

        self._initialize_weights() 

def __init__(self, n_classes, n_frames, n_channels, input_size, hidden_size,
               n_layers, dropout, hidden_size_seq, n_layers_seq, dropout_seq,
               encoder_type):
    super(TAD, self).__init__()

    # Visual encoder
    if encoder_type == 'cnn':
      self.embed = get_resnet(n_frames * n_channels, n_classes)
    elif encoder_type == 'rnn':
      self.embed = get_gru(input_size, hidden_size, n_layers, dropout,
                           n_classes)
    else:
      raise NotImplementedError

    # Sequence encoder
    self.rnn = nn.GRU(
        input_size=hidden_size,
        hidden_size=hidden_size_seq,
        num_layers=n_layers_seq,
        batch_first=True,
        dropout=dropout_seq,
        bidirectional=True)
    # Classification layer
    self.fc = nn.Linear(hidden_size_seq,
                        n_classes + 1)  # plus 1 class for background

    self.n_classes = n_classes
    self.hidden_size = hidden_size
    self.hidden_size_seq = hidden_size_seq
    self.n_layers_seq = n_layers_seq
    self.encoder_type = encoder_type