在Theano中从同形方阵列表构建分块对角矩阵（GPU优化）

Got it, let's break this down properly—you're looking to work with block diagonal matrices in Theano's GPU environment to speed up eigenvalue/eigenvector calculations for lots of small, same-shape square matrices. The key here is to leverage Theano's native GPU-optimized functions and avoid unnecessary overhead (like building a huge block diagonal matrix when you don't have to).

First: Skip Building the Block Diagonal Matrix (If You Can) #

Here's a critical insight: The eigenvalues and eigenvectors of a block diagonal matrix are just the union of eigenvalues/eigenvectors from each individual small matrix. That means you don't need to construct the full large matrix to compute what you need—this is way more efficient for GPU processing, since we can batch-process all small matrices directly.

Batch Eigenvalue/Eigenvector Calculation #

First, structure your input as a 3D Theano tensor matrices with shape (num_matrices, mat_size, mat_size) (stack all your small square matrices into this tensor instead of keeping them as a list). Then use Theano's native linear algebra functions, which are optimized for GPU parallelism:

import theano
import theano.tensor as T

# Symbolic input: 3D tensor holding N small MxM matrices
matrices = T.tensor3("matrices", dtype="float32")

# For symmetric/Hermitian matrices (faster, more stable—use this if applicable!)
eigenvalues, eigenvectors = T.nlinalg.eigh(matrices)

# For general square matrices (use only if your matrices aren't symmetric)
# eigenvalues, eigenvectors = T.nlinalg.eig(matrices)

# Compile the GPU-accelerated function
compute_eigs = theano.function([matrices], [eigenvalues, eigenvectors])

• The output eigenvalues will be a (N, M) tensor (each row contains eigenvalues for one small matrix).
• The output eigenvectors will be a (N, M, M) tensor (each slice is the eigenvector matrix for one small matrix).

This approach lets the GPU parallelize computations across all your small matrices simultaneously—way faster than building a giant block diagonal matrix first.

If You Must Build the Block Diagonal Matrix #

If you need the actual block diagonal matrix for other operations (not just eigenvalue calculations), use Theano's native T.extra_ops.block_diag or a tensor-based approach (no Python loops, which kill GPU performance):

Option 1: Using T.extra_ops.block_diag #

import theano.tensor.extra_ops as extra_ops

# Unpack the 3D tensor into individual matrices and build the block diagonal
block_diag_matrix = extra_ops.block_diag(*[matrices[i] for i in range(matrices.shape[0])])

Option 2: Tensor-Manipulation Approach (Better for Large Batches) #

This avoids explicit Python unpacking, which is cleaner for large numbers of matrices:

# Get dimensions
num_mats = matrices.shape[0]
mat_size = matrices.shape[1]
total_size = num_mats * mat_size

# Create a mask to select diagonal block positions
mask = T.repeat(T.eye(num_mats), mat_size, axis=0)
mask = T.repeat(mask, mat_size, axis=1)

# Initialize empty block diagonal matrix and fill in blocks
block_diag_matrix = T.zeros((total_size, total_size), dtype="float32")
block_diag_matrix = T.set_subtensor(block_diag_matrix[mask.nonzero()], matrices.flatten())

Both approaches use Theano's GPU-optimized operations, so no slow CPU-GPU data transfers here.

Performance Tips for Maximum GPU Speed #

• Use float32: GPUs are drastically faster with single-precision floats. Cast your input matrices to float32 unless you absolutely need double-precision (float64).
• Batch smartly: Always use a 3D tensor for your small matrices—Theano can't vectorize operations over a Python list, so stacking into a tensor is non-negotiable for parallelism.
• Prefer eigh over eig: If your matrices are symmetric (or Hermitian), eigh uses specialized GPU kernels that are way faster and more stable than the general eig function.
• Avoid Python loops: Any loop in Python will force data to move back to the CPU, destroying GPU performance. Stick to Theano's symbolic tensor operations.

Quick Test Example #

Let’s verify this with a small batch of symmetric matrices:

import numpy as np

# Create 5 random symmetric 2x2 matrices (float32 for GPU)
np_mats = np.random.randn(5, 2, 2).astype(np.float32)
np_mats = (np_mats + np_mats.transpose(0, 2, 1)) / 2  # Make symmetric

# Compute on GPU
eigs, vecs = compute_eigs(np_mats)

# Compare to numpy's result for the first matrix
first_mat_eigs_np = np.linalg.eigvalsh(np_mats[0])
print("First matrix eigenvalues (NumPy):", first_mat_eigs_np)
print("First matrix eigenvalues (Theano GPU):", eigs[0])

You’ll get identical results, but the Theano version will scale to thousands of matrices way faster on GPU.

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