PyTorch中基于可变长度视频特征的LSTM训练技术问询

Hey there! Let's break down the correct way to train your LSTM on variable-length video sequences (10-35 frames) with a batch size of 1. First, we'll fix gaps in your current code and walk through key best practices tailored to your use case.

1. Complete LSTM Model Definition #

Your original model code was incomplete—here's a full, functional implementation optimized for video feature processing:

import torch
import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # Use batch_first=True to align with intuitive (batch, seq_len, input_size) tensor format
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        # Final fully connected layer for classification (adjust based on your task)
        self.fc = nn.Linear(hidden_size, num_classes)
    
    def forward(self, x):
        # Initialize hidden and cell states (match the device of your input tensor)
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        
        # Forward pass through LSTM
        # Output shape: (batch_size, seq_len, hidden_size)
        output, _ = self.lstm(x, (h0, c0))
        
        # For video classification, use the last time step's hidden state as the sequence's core representation
        last_timestep_hidden = output[:, -1, :]
        return self.fc(last_timestep_hidden)

2. Corrected Training Loop #

Your original loop was missing critical steps like loss calculation, optimization, and proper tensor formatting. Here's the complete training pipeline:

# Hyperparameters (adjust based on your task requirements)
INPUT_SIZE = 4096
HIDDEN_SIZE = 4096
NUM_LAYERS = 1
NUM_CLASSES = 64  # Matches your original code's 64
NUM_EPOCHS = 10
LEARNING_RATE = 0.001

# Setup device (use GPU if available for faster training)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialize model, loss function, and optimizer
model = LSTMModel(INPUT_SIZE, HIDDEN_SIZE, NUM_LAYERS, NUM_CLASSES).to(device)
criterion = nn.CrossEntropyLoss()  # Use cross-entropy for classification; switch to MSELoss for regression
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Set model to training mode
model.train()

for epoch in range(NUM_EPOCHS):
    epoch_loss = 0.0
    for step, (video_features, label) in enumerate(data_loader):
        # Format input: convert (seq_len, 4096) to (1, seq_len, 4096) (since batch_size=1)
        bx = video_features.unsqueeze(0).to(device)
        label = label.to(device)
        
        # Forward pass through the model
        outputs = model(bx)
        
        # Calculate loss
        loss = criterion(outputs, label)
        
        # Backward pass and optimize
        optimizer.zero_grad()  # Critical: clear accumulated gradients from previous steps
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
        
        # Log progress every 100 steps (adjust frequency as needed)
        if (step + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{NUM_EPOCHS}], Step [{step+1}/{len(data_loader)}], Avg Loss: {epoch_loss/100:.4f}")
            epoch_loss = 0.0

3. Key Notes for Variable-Length Sequences #

• Tensor Format: Always ensure your input follows (batch_size, seq_len, input_size) when using batch_first=True—this makes handling variable lengths seamless with batch size 1.
• Hidden State Initialization: We reinitialize h0 and c0 for every video since each sequence is independent. For continuous time-series data, you could retain states, but this isn't needed for isolated videos.
• No Padding Required (Batch Size 1): Since you're using batch size 1, you don't need to pad sequences to a fixed length. If you later increase batch size, use torch.nn.utils.rnn.pack_padded_sequence to skip padding tokens during LSTM processing and avoid unnecessary computation.
• Task-Specific Output: The code above uses the last time step's hidden state for classification. If you're doing frame-wise prediction (e.g., action detection per frame), you'd use the full output tensor instead.

4. Optional Enhancements #

• Gradient Clipping: Prevent exploding gradients (a common LSTM issue) by adding torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) right after loss.backward().
• Learning Rate Scheduling: Adjust the learning rate over time with torch.optim.lr_scheduler.StepLR to improve convergence.
• Validation Loop: Add a validation phase after each epoch to monitor overfitting—switch the model to eval() mode and wrap the forward pass in torch.no_grad() to disable gradient computation.

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