Hey there! Let's tackle your two questions about modifying RNN/GRU units in PyTorch—super common when you need custom behavior without sacrificing efficiency.

There are two practical, efficient approaches to customize RNN units without touching C-level backend code:

• Build a custom module with vectorized operations: Reimplement the core RNN logic using PyTorch's optimized tensor operations. This avoids slow Python loops and leverages PyTorch's GPU/CPU acceleration fully.
• Extend existing modules selectively: If you only need to tweak small parts (like an activation function), wrap the built-in unit and override specific forward pass steps. For full control though, a custom implementation is better.
• Golden rule: Always prioritize batch-wise vectorized operations over Python-level for/while loops—PyTorch's tensor operations are optimized to run at near-native speed.

You're right that PyTorch's built-in GRU uses C++ backend code that's not directly editable. But we can reimplement the entire GRU logic using pure PyTorch tensor operations, which are just as efficient and fully customizable.

Core GRU Logic Recap #

GRU relies on three key components: update gate (z), reset gate (r), and candidate hidden state (h̃). The core formulas are:

( z = \sigma(W_z \cdot [h_{prev}, x] + b_z) )
( r = \sigma(W_r \cdot [h_{prev}, x] + b_r) )
( \tilde{h} = \text{tanh}(W_h \cdot [r \odot h_{prev}, x] + b_h) )
( h_{new} = (1 - z) \odot h_{prev} + z \odot \tilde{h} )

Custom GRU Cell Implementation #

This cell handles single-step computation, with all operations vectorized for batch processing (no per-element loops):

import torch
import torch.nn as nn
import torch.nn.functional as F

class CustomGRUCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # Combine weight matrices to minimize tensor operations (faster computation)
        self.weight_ih = nn.Parameter(torch.randn(3 * hidden_size, input_size))
        self.weight_hh = nn.Parameter(torch.randn(3 * hidden_size, hidden_size))
        self.bias_ih = nn.Parameter(torch.randn(3 * hidden_size))
        self.bias_hh = nn.Parameter(torch.randn(3 * hidden_size))

    def forward(self, x, h_prev):
        # x shape: (batch_size, input_size)
        # h_prev shape: (batch_size, hidden_size)
        
        # Compute all linear transformations in one go (z, r, h̃ linear terms)
        gates = torch.mm(x, self.weight_ih.t()) + self.bias_ih + torch.mm(h_prev, self.weight_hh.t()) + self.bias_hh
        
        # Split into update gate, reset gate, and candidate hidden state
        z, r, h_tilde = gates.chunk(3, dim=1)
        
        # Apply activation functions
        z = torch.sigmoid(z)
        r = torch.sigmoid(r)
        h_tilde = torch.tanh(h_tilde)
        
        # Update hidden state
        h_new = (1 - z) * h_prev + z * h_tilde
        
        return h_new

Full GRU Layer for Sequences #

To handle entire input sequences, we stack the custom cells. The time-step loop here is minimal Python overhead—all internal computations are vectorized batch operations:

class CustomGRU(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # Stack multiple GRU cells for deep GRU
        self.layers = nn.ModuleList([
            CustomGRUCell(input_size if i == 0 else hidden_size, hidden_size) 
            for i in range(num_layers)
        ])

    def forward(self, x, h0=None):
        # x shape: (seq_len, batch_size, input_size)
        seq_len, batch_size, _ = x.shape
        device = x.device
        
        # Initialize hidden state if not provided
        if h0 is None:
            h0 = torch.zeros(self.num_layers, batch_size, self.hidden_size, device=device)
            
        h_prev_list = list(h0)
        output_sequence = []
        
        # Process each time step
        for t in range(seq_len):
            xt = x[t]
            # Pass through each layer
            for i, layer in enumerate(self.layers):
                h_prev = h_prev_list[i]
                h_new = layer(xt, h_prev)
                h_prev_list[i] = h_new
                xt = h_new  # Output of current layer is input to next
            output_sequence.append(xt)
            
        # Stack outputs into a tensor and prepare final hidden state
        output = torch.stack(output_sequence, dim=0)
        hn = torch.stack(h_prev_list, dim=0)
        
        return output, hn

Customization Example #

If you want to tweak the GRU logic—say, replace the sigmoid update gate with a GELU activation—you just modify the relevant line in CustomGRUCell:

z = F.gelu(z)  # Instead of torch.sigmoid(z)

This implementation is fully compatible with PyTorch's ecosystem (use it with optimizers, move to GPU, etc.) and maintains efficiency because all core operations use PyTorch's optimized tensor kernels.

内容的提问来源于stack exchange，提问作者erutan