TensorFlow中能否单次实现各通道与不同矩阵相乘？

Absolutely! You can pull this off in one go using batch matrix multiplication (or Einstein summation for more flexibility), which lets you pack all three channel-wise matrix multiplications into a single optimized operation. Here's how it works:

Core Idea #

Your 3D array A (shape (N, D1, 3)) can be thought of as 3 separate (N, D1) matrices (one per the last dimension). Instead of multiplying each with B1, B2, B3 individually, we'll stack the three B matrices into a single tensor, then use a batch-aware multiplication to compute all three products at once.

Step-by-Step Implementation #

Let's use code examples to make this concrete (we'll cover NumPy and PyTorch, but the logic translates to TensorFlow/JAX too):

Example 1: Using NumPy with Einstein Summation

Einstein summation (np.einsum) is great for explicit, readable tensor operations:

import numpy as np

# Define sample dimensions and data
N, D1, D2 = 5, 4, 3
A = np.random.rand(N, D1, 3)  # Shape (N, D1, 3)
B1 = np.random.rand(D1, D2)   # Shape (D1, D2)
B2 = np.random.rand(D1, D2)
B3 = np.random.rand(D1, D2)

# Stack the B matrices along the last dimension to match A's channel axis
B_stack = np.stack([B1, B2, B3], axis=-1)  # Shape (D1, D2, 3)

# Single multiplication to get the (N, D2, 3) result
result = np.einsum('n d1 c, d1 d2 c -> n d2 c', A, B_stack)

# Verify it matches the manual channel-wise multiplication
result_manual = np.stack([
    A[..., 0] @ B1,
    A[..., 1] @ B2,
    A[..., 2] @ B3
], axis=-1)
assert np.allclose(result, result_manual)  # Passes if correct

Example 2: Using PyTorch with Batch Matrix Multiplication (bmm)

PyTorch's torch.bmm is optimized for batch matrix multiplication, which requires inputs to be in (batch_size, m, k) and (batch_size, k, n) format:

import torch

# Same sample data (converted to tensors)
A = torch.rand(N, D1, 3)
B1 = torch.rand(D1, D2)
B2 = torch.rand(D1, D2)
B3 = torch.rand(D1, D2)

# Stack B matrices into a (3, D1, D2) batch
B_batch = torch.stack([B1, B2, B3], dim=0)  # Shape (3, D1, D2)
# Reshape A to match the batch axis order: (3, N, D1)
A_reshaped = A.permute(2, 0, 1)  # Swap dimensions to align with B_batch

# Single batch matrix multiplication
batch_result = torch.bmm(A_reshaped, B_batch)  # Shape (3, N, D2)
# Reshape back to (N, D2, 3)
result = batch_result.permute(1, 2, 0)

# Verify against manual calculation
result_manual = torch.stack([
    A[..., 0] @ B1,
    A[..., 1] @ B2,
    A[..., 2] @ B3
], axis=-1)
assert torch.allclose(result, result_manual)

Why This Works #

Under the hood, frameworks like NumPy/PyTorch optimize these batch operations to run efficiently on CPU/GPU—often faster than looping through each channel and running separate matrix multiplies. It also keeps your code cleaner and avoids redundant boilerplate.

内容的提问来源于stack exchange，提问作者girl-meets-world