树深时见鹿，
溪午不闻钟。

import torch

# 创建一个 [3, 10, 20] 的张量
tensor = torch.randn(3, 10, 20)

print(f"张量形状: {tensor.shape}")
print(f"张量维度: {tensor.dim()}D")
print(f"总元素数: {tensor.numel()}")  # 3 × 10 × 20 = 600

# 访问不同层的矩阵
print(f"第0层矩阵形状: {tensor[0].shape}")  # [10, 20]
print(f"第1层矩阵形状: {tensor[1].shape}")  # [10, 20]
print(f"第2层矩阵形状: {tensor[2].shape}")  # [10, 20]

import torch

# 形状: [seq_len, embedding_dim]
sentence = torch.randn(10, 20)  # 10个词，每词20维

# 词嵌入矩阵通常是 [vocab_size, embedding_dim]
vocab_embedding = torch.randn(5000, 20)

# 线性变换矩阵通常是 [input_dim, output_dim]
linear_weight = torch.randn(20, 128)  # 20维输入 -> 128维输出

# 矩阵乘法很自然: sentence @ linear_weight
output = sentence @ linear_weight  # [10, 20] @ [20, 128] = [10, 128]

# 每一行是一个词的向量表示
sentence = torch.tensor([
    [0.1, 0.2, 0.3, ...],  # 第1个词的向量 (20维)
    [0.4, 0.5, 0.6, ...],  # 第2个词的向量 (20维)
    [0.7, 0.8, 0.9, ...],  # 第3个词的向量 (20维)
    # ... 共10行(10个词)
])

# 单个句子: [seq_len, embedding_dim]
single_sentence = torch.randn(10, 20)

# 批处理: [batch_size, seq_len, embedding_dim]
batch_sentences = torch.randn(32, 10, 20)  # 32个句子

# 这样维度扩展很自然，最后一维始终是特征维度

import torch.nn as nn

# LSTM/GRU等RNN层期望输入: [seq_len, batch_size, input_size]
# 或 [batch_size, seq_len, input_size]
lstm = nn.LSTM(input_size=20, hidden_size=128)

# Transformer的Multi-Head Attention
# Q, K, V 矩阵形状: [seq_len, d_model]
attention = nn.MultiheadAttention(embed_dim=20, num_heads=4)

# 表格数据: [样本数, 特征数]
tabular_data = torch.randn(1000, 50)  # 1000个样本，50个特征

# 图像: [batch, channels, height, width]
images = torch.randn(32, 3, 224, 224)

# NLP: [batch, seq_len, embedding_dim]
text = torch.randn(32, 100, 300)  # 保持最后一维是特征维度

import torch

# [10, 20] 张量 = 10个长度为20的向量
sentence = torch.randn(10, 20)

print(f"句子形状: {sentence.shape}")
print(f"词数: {sentence.shape[0]}")        # 10个词
print(f"每个词的向量维度: {sentence.shape[1]}")  # 每个词20维

# 获取单个词向量
word_1 = sentence[0]  # 第1个词的向量，形状: [20]
word_2 = sentence[1]  # 第2个词的向量，形状: [20]
print(f"单个词向量形状: {word_1.shape}")

# 假设句子: "I love machine learning"
sentence_vectors = torch.tensor([
    [0.1, 0.3, -0.2, 0.8, ...],  # "I" 的20维向量
    [0.5, -0.1, 0.7, 0.2, ...], # "love" 的20维向量
    [0.2, 0.9, -0.3, 0.1, ...], # "machine" 的20维向量
    [0.8, 0.4, 0.6, -0.5, ...], # "learning" 的20维向量
    # ... 共10个词
])

# 每一行 = 一个词的向量表示
# 每一列 = 向量的某个维度特征

# 单个句子: [seq_len, embedding_dim]
single_sentence = torch.randn(10, 20)  # 10个词，每词20维

# 批处理: [batch_size, seq_len, embedding_dim]
batch = torch.randn(3, 10, 20)  # 3个句子，每句10个词，每词20维

# 可以这样理解批处理：
sentence_1 = batch[0]  # [10, 20] - 第1个句子的10个词向量
sentence_2 = batch[1]  # [10, 20] - 第2个句子的10个词向量
sentence_3 = batch[2]  # [10, 20] - 第3个句子的10个词向量

# 词嵌入层的输出就是这种格式
embedding = nn.Embedding(vocab_size=5000, embedding_dim=20)
word_ids = torch.tensor([1, 25, 100, 300])  # 4个词的ID
word_vectors = embedding(word_ids)  # [4, 20] - 4个词，每词20维

print(f"4个词向量: {word_vectors.shape}")
print(f"第1个词向量: {word_vectors[0].shape}")  # [20]

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

# 1. 最大池化下采样
input_tensor = torch.randn(1, 1, 8, 8)  # [batch, channel, height, width]
maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
downsampled = maxpool(input_tensor)  # [1, 1, 4, 4]
print(f"原始: {input_tensor.shape} -> 下采样: {downsampled.shape}")

# 2. 平均池化下采样
avgpool = nn.AvgPool2d(kernel_size=2, stride=2)
downsampled_avg = avgpool(input_tensor)  # [1, 1, 4, 4]

# 3. 卷积下采样 (stride > 1)
conv_down = nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1)
downsampled_conv = conv_down(input_tensor)  # [1, 1, 4, 4]

# 输入: [1, 1, 4, 4]
small_tensor = torch.randn(1, 1, 4, 4)

# 1. 最近邻插值上采样
upsampled_nearest = F.interpolate(small_tensor, scale_factor=2, mode='nearest')
print(f"最近邻: {small_tensor.shape} -> {upsampled_nearest.shape}")  # [1, 1, 8, 8]

# 2. 双线性插值上采样
upsampled_bilinear = F.interpolate(small_tensor, scale_factor=2, mode='bilinear', align_corners=False)
print(f"双线性: {small_tensor.shape} -> {upsampled_bilinear.shape}")  # [1, 1, 8, 8]

# 3. 转置卷积上采样
conv_transpose = nn.ConvTranspose2d(1, 1, kernel_size=2, stride=2)
upsampled_conv = conv_transpose(small_tensor)  # [1, 1, 8, 8]

# 4. 像素重排上采样 (PixelShuffle)
# 先用卷积增加通道数，再重排
conv_expand = nn.Conv2d(1, 4, kernel_size=1)  # 1 -> 4 channels
pixel_shuffle = nn.PixelShuffle(upscale_factor=2)
expanded = conv_expand(small_tensor)  # [1, 4, 4, 4]
upsampled_ps = pixel_shuffle(expanded)  # [1, 1, 8, 8]

# CNN中的特征提取（下采样）
class DownBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.conv(x)
        return self.pool(x)  # 下采样

# U-Net中的上采样
class UpBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_channels, out_channels, 2, stride=2)
        self.conv = nn.Conv2d(out_channels, out_channels, 3, padding=1)

    def forward(self, x):
        x = self.up(x)  # 上采样
        return self.conv(x)

# 1D信号的采样
audio_signal = torch.randn(1, 1, 16000)  # 16kHz音频

# 下采样到8kHz
downsampled_audio = F.interpolate(audio_signal, size=8000, mode='linear')

# 上采样回16kHz
upsampled_audio = F.interpolate(downsampled_audio, size=16000, mode='linear')

# 创建一个简单的图像数据
original = torch.randn(1, 3, 64, 64)  # 64x64 RGB图像

print("=== 下采样示例 ===")
# 下采样到32x32
down_32 = F.interpolate(original, size=32, mode='bilinear', align_corners=False)
print(f"64x64 -> 32x32: {original.shape} -> {down_32.shape}")

# 下采样到16x16
down_16 = F.interpolate(original, size=16, mode='bilinear', align_corners=False)
print(f"64x64 -> 16x16: {original.shape} -> {down_16.shape}")

print("\n=== 上采样示例 ===")
# 从16x16上采样到64x64
up_64 = F.interpolate(down_16, size=64, mode='bilinear', align_corners=False)
print(f"16x16 -> 64x64: {down_16.shape} -> {up_64.shape}")

# 从32x32上采样到128x128
up_128 = F.interpolate(down_32, size=128, mode='bilinear', align_corners=False)
print(f"32x32 -> 128x128: {down_32.shape} -> {up_128.shape}")

import torch
import torch.nn.functional as F

# 图像数据: [batch, channels, height, width]
image = torch.randn(1, 3, 64, 64)  # 最后一个维度是width=64

print("=== 下采样 ===")
# 下采样改变的是height和width，不是最后一个维度
downsampled = F.interpolate(image, size=32, mode='bilinear')
print(f"原始: {image.shape}")        # [1, 3, 64, 64]
print(f"下采样: {downsampled.shape}")  # [1, 3, 32, 32]
# 最后一个维度从64变成32，但倒数第二个维度(height)也从64变成32

print("\n=== 上采样 ===")
upsampled = F.interpolate(downsampled, size=128, mode='bilinear')
print(f"上采样: {upsampled.shape}")    # [1, 3, 128, 128]
# height和width都变成128

# 1D信号: [batch, channels, length]
signal_1d = torch.randn(1, 1, 1000)
down_1d = F.interpolate(signal_1d, size=500)  # 改变length维度
print(f"1D: {signal_1d.shape} -> {down_1d.shape}")

# 2D图像: [batch, channels, height, width]
image_2d = torch.randn(1, 3, 64, 64)
down_2d = F.interpolate(image_2d, size=32)    # 改变height和width
print(f"2D: {image_2d.shape} -> {down_2d.shape}")

# 3D体积: [batch, channels, depth, height, width]
volume_3d = torch.randn(1, 1, 32, 64, 64)
down_3d = F.interpolate(volume_3d, size=16)   # 改变depth, height, width
print(f"3D: {volume_3d.shape} -> {down_3d.shape}")

# ❌ 错误理解：只改变最后一个维度
# [1, 3, 64, 64] -> [1, 3, 64, 32]  # 只改变width

# ✅ 正确理解：改变空间维度
# [1, 3, 64, 64] -> [1, 3, 32, 32]  # 改变height和width

# 有时会同时改变通道数和空间尺寸
conv_down = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1)
input_tensor = torch.randn(1, 3, 64, 64)
output = conv_down(input_tensor)
print(f"卷积下采样: {input_tensor.shape} -> {output.shape}")
# [1, 3, 64, 64] -> [1, 64, 32, 32]
# 通道数: 3->64, 空间尺寸: 64->32

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

# 文本序列: [batch_size, seq_len, embedding_dim]
text = torch.randn(2, 100, 256)  # 2个句子，每句100个词，每词256维

print("=== 序列长度采样 ===")
# 下采样：减少序列长度
# 方法1：池化
pooled = F.avg_pool1d(text.transpose(1, 2), kernel_size=2, stride=2)
pooled = pooled.transpose(1, 2)
print(f"池化下采样: {text.shape} -> {pooled.shape}")  # [2, 50, 256]

# 方法2：步长采样
downsampled = text[:, ::2, :]  # 每隔一个词取一个
print(f"步长下采样: {text.shape} -> {downsampled.shape}")  # [2, 50, 256]

# 这是NLP中更常见的"采样"概念
input_text = torch.randn(2, 100, 256)  # [batch, seq_len, embedding_dim]

print("=== 嵌入维度变化 ===")
# "下采样"：降低嵌入维度
down_proj = nn.Linear(256, 128)
reduced_dim = down_proj(input_text)
print(f"维度下采样: {input_text.shape} -> {reduced_dim.shape}")  # [2, 100, 128]

# "上采样"：增加嵌入维度
up_proj = nn.Linear(128, 512)
expanded_dim = up_proj(reduced_dim)
print(f"维度上采样: {reduced_dim.shape} -> {expanded_dim.shape}")  # [2, 100, 512]

# CNN for NLP中的多尺度卷积
class TextCNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)

        # 不同kernel_size的卷积 = 不同"采样"尺度
        self.conv1 = nn.Conv1d(embed_dim, 100, kernel_size=3)  # 3-gram
        self.conv2 = nn.Conv1d(embed_dim, 100, kernel_size=4)  # 4-gram
        self.conv3 = nn.Conv1d(embed_dim, 100, kernel_size=5)  # 5-gram

        self.fc = nn.Linear(300, num_classes)

    def forward(self, x):
        # x: [batch, seq_len]
        x = self.embedding(x)  # [batch, seq_len, embed_dim]
        x = x.transpose(1, 2)  # [batch, embed_dim, seq_len]

        # 不同尺度的特征提取
        conv1_out = F.relu(self.conv1(x))  # [batch, 100, seq_len-2]
        conv2_out = F.relu(self.conv2(x))  # [batch, 100, seq_len-3]
        conv3_out = F.relu(self.conv3(x))  # [batch, 100, seq_len-4]

        # 全局池化（下采样到固定长度）
        pool1 = F.max_pool1d(conv1_out, conv1_out.size(2)).squeeze(2)
        pool2 = F.max_pool1d(conv2_out, conv2_out.size(2)).squeeze(2)
        pool3 = F.max_pool1d(conv3_out, conv3_out.size(2)).squeeze(2)

        # 拼接多尺度特征
        out = torch.cat([pool1, pool2, pool3], dim=1)  # [batch, 300]
        return self.fc(out)

# 使用示例
model = TextCNN(vocab_size=10000, embed_dim=256, num_classes=2)
input_ids = torch.randint(0, 10000, (4, 50))  # [batch, seq_len]
output = model(input_ids)
print(f"文本分类输出: {output.shape}")  # [4, 2]

# 位置编码可以看作是对序列位置的"采样"
def positional_encoding(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(0, seq_len).unsqueeze(1).float()

    # 不同频率的"采样"
    div_term = torch.exp(torch.arange(0, d_model, 2).float() *
                        -(torch.log(torch.tensor(10000.0)) / d_model))

    pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位置
    pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位置

    return pe

# 不同序列长度的位置编码
pe_100 = positional_encoding(100, 256)  # 100个位置
pe_50 = positional_encoding(50, 256)    # 50个位置（下采样）
pe_200 = positional_encoding(200, 256)  # 200个位置（上采样）

print(f"位置编码: {pe_100.shape}, {pe_50.shape}, {pe_200.shape}")

# 文本摘要可以看作是序列的"下采样"
class TextSummarizer(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.decoder = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.output_proj = nn.Linear(hidden_dim, vocab_size)

    def forward(self, input_text, target_length):
        # 输入长文本: [batch, long_seq_len]
        # 输出短摘要: [batch, short_seq_len]
        embedded = self.embedding(input_text)

        # 编码器处理长序列
        encoder_out, (h, c) = self.encoder(embedded)

        # 解码器生成短序列（下采样效果）
        decoder_input = torch.zeros(input_text.size(0), target_length,
                                  self.embedding.embedding_dim)
        summary, _ = self.decoder(decoder_input, (h, c))

        return self.output_proj(summary)

# 长文本 -> 短摘要 = 序列长度的"下采样"

import torch
import torch.nn as nn

# 创建线性层：输入维度 -> 输出维度
linear = nn.Linear(in_features=10, out_features=5)

print(f"权重形状: {linear.weight.shape}")  # [5, 10] - [out_features, in_features]
print(f"偏置形状: {linear.bias.shape}")    # [5] - [out_features]

# 线性变换: y = xW^T + b
# 其中:
# x: 输入 [batch_size, in_features] 或 [..., in_features]
# W: 权重矩阵 [out_features, in_features]
# b: 偏置向量 [out_features]
# y: 输出 [batch_size, out_features] 或 [..., out_features]

# 手动实现线性层
def manual_linear(x, weight, bias):
    # x @ weight.T + bias
    return torch.mm(x, weight.t()) + bias

# 验证
x = torch.randn(3, 10)  # 3个样本，每个10维
linear = nn.Linear(10, 5)

# PyTorch实现
output_pytorch = linear(x)

# 手动实现
output_manual = manual_linear(x, linear.weight, linear.bias)

print(f"PyTorch输出: {output_pytorch.shape}")
print(f"手动实现输出: {output_manual.shape}")
print(f"结果相等: {torch.allclose(output_pytorch, output_manual)}")

# 1. 简单的线性变换
linear = nn.Linear(128, 64)  # 128维 -> 64维
input_tensor = torch.randn(32, 128)  # 32个样本
output = linear(input_tensor)  # [32, 64]

print(f"输入: {input_tensor.shape} -> 输出: {output.shape}")

# 2. 多维输入（保持最后一维变换）
text_embeddings = torch.randn(4, 50, 128)  # [batch, seq_len, embed_dim]
linear_text = nn.Linear(128, 256)
transformed = linear_text(text_embeddings)  # [4, 50, 256]

print(f"文本嵌入: {text_embeddings.shape} -> {transformed.shape}")

# 3. 不使用偏置
linear_no_bias = nn.Linear(10, 5, bias=False)
print(f"无偏置层权重: {linear_no_bias.weight.shape}")
print(f"偏置存在: {linear_no_bias.bias is not None}")

class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_classes)
        )

    def forward(self, x):
        return self.layers(x)

# 使用示例
classifier = SimpleClassifier(784, 512, 10)  # MNIST分类器
mnist_image = torch.randn(64, 784)  # 64张28x28的图片展平
predictions = classifier(mnist_image)  # [64, 10]
print(f"分类预测: {predictions.shape}")

class TransformerFFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        return self.linear2(self.relu(self.linear1(x)))

# Transformer的FFN层
ffn = TransformerFFN(d_model=512, d_ff=2048)
sequence = torch.randn(32, 100, 512)  # [batch, seq_len, d_model]
output = ffn(sequence)  # [32, 100, 512]
print(f"FFN输出: {output.shape}")

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Q, K, V的线性投影
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # 线性投影得到Q, K, V
        Q = self.w_q(x)  # [batch, seq_len, d_model]
        K = self.w_k(x)  # [batch, seq_len, d_model]
        V = self.w_v(x)  # [batch, seq_len, d_model]

        # 重塑为多头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k)

        # 简化的注意力计算...
        # 这里省略具体的注意力计算

        return self.w_o(x)  # 输出投影

attention = MultiHeadAttention(d_model=512, num_heads=8)

# 查看默认初始化
linear = nn.Linear(100, 50)
print(f"权重统计: 均值={linear.weight.mean():.4f}, 标准差={linear.weight.std():.4f}")
print(f"偏置统计: 均值={linear.bias.mean():.4f}, 标准差={linear.bias.std():.4f}")

# 自定义初始化
def init_weights(module):
    if isinstance(module, nn.Linear):
        # Xavier初始化
        nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)

linear.apply(init_weights)
print(f"初始化后权重统计: 均值={linear.weight.mean():.4f}, 标准差={linear.weight.std():.4f}")

# 1. 批量处理效率
batch_sizes = [1, 32, 128, 512]
linear = nn.Linear(1024, 512)

for batch_size in batch_sizes:
    x = torch.randn(batch_size, 1024)

    # 计时
    import time
    start_time = time.time()
    for _ in range(100):
        _ = linear(x)
    end_time = time.time()

    print(f"批大小 {batch_size:3d}: {(end_time - start_time)*10:.2f}ms/100次")

# 2. GPU加速
if torch.cuda.is_available():
    linear_gpu = linear.cuda()
    x_gpu = torch.randn(512, 1024).cuda()
    output_gpu = linear_gpu(x_gpu)
    print(f"GPU输出: {output_gpu.shape}")

# 1. 维度不匹配
linear = nn.Linear(100, 50)
try:
    wrong_input = torch.randn(32, 200)  # 应该是(..., 100)
    output = linear(wrong_input)
except RuntimeError as e:
    print(f"维度错误: {e}")

# 2. 正确的多维输入
correct_input = torch.randn(32, 10, 100)  # [batch, seq, features]
output = linear(correct_input)  # [32, 10, 50] - 只变换最后一维
print(f"多维输入: {correct_input.shape} -> {output.shape}")

# 3. 内存效率
# 对于大型网络，考虑使用in-place操作或梯度检查点
large_linear = nn.Linear(4096, 4096)
large_input = torch.randn(128, 4096)
# 使用torch.utils.checkpoint可以节省内存

import torch

# 张量的维度数（ndim/rank）
tensor_2d = torch.randn(32, 128)    # 2维张量
tensor_3d = torch.randn(32, 50, 128) # 3维张量
tensor_4d = torch.randn(32, 3, 224, 224) # 4维张量

print(f"2D张量维度数: {tensor_2d.dim()}")  # 2
print(f"3D张量维度数: {tensor_3d.dim()}")  # 3
print(f"4D张量维度数: {tensor_4d.dim()}")  # 4

# 特征维度（通常是最后一维）
print(f"2D张量特征维度: {tensor_2d.shape[-1]}")  # 128
print(f"3D张量特征维度: {tensor_3d.shape[-1]}")  # 128
print(f"4D张量特征维度: {tensor_4d.shape[-1]}")  # 224（宽度，不是特征！）

# NLP: 词嵌入维度
word_embeddings = torch.randn(100, 300)  # 100个词，每词300维
print(f"词嵌入维度: {word_embeddings.shape[-1]}")  # 300维

# 线性层: 输入/输出维度
linear = torch.nn.Linear(512, 256)  # 512维输入 -> 256维输出
print(f"输入维度: {linear.in_features}, 输出维度: {linear.out_features}")

# 图像: 高度和宽度维度
image = torch.randn(1, 3, 224, 224)  # [batch, channels, height, width]
print(f"图像空间维度: {image.shape[2]} x {image.shape[3]}")  # 224 x 224

# 卷积核: 空间维度
conv = torch.nn.Conv2d(3, 64, kernel_size=3)  # 3x3卷积核
print(f"卷积核空间维度: {conv.kernel_size}")  # (3, 3)

# RNN: 序列长度维度
sequence = torch.randn(32, 50, 128)  # [batch, seq_len, features]
print(f"序列长度维度: {sequence.shape[1]}")  # 50
print(f"特征维度: {sequence.shape[2]}")      # 128

import torch.nn as nn

# 例子1: NLP中的"维度"通常指特征维度（最后一维）
text = torch.randn(4, 50, 300)  # [batch, seq_len, embed_dim]
linear_text = nn.Linear(300, 128)  # 改变特征维度：300 -> 128
output_text = linear_text(text)    # [4, 50, 128]
print(f"NLP维度变化: {text.shape} -> {output_text.shape}")
print("这里的'维度'指的是最后一维（特征维度）")

# 例子2: 图像中的"维度"可能指不同的东西
image = torch.randn(2, 3, 64, 64)  # [batch, channels, height, width]

# 通道维度变化
conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
conv_output = conv(image)  # [2, 16, 64, 64]
print(f"通道维度变化: {image.shape[1]} -> {conv_output.shape[1]}")

# 空间维度变化（下采样）
pool = nn.MaxPool2d(2, 2)
pool_output = pool(image)  # [2, 3, 32, 32]
print(f"空间维度变化: {image.shape[2:]} -> {pool_output.shape[2:]}")

# 批次维度（第0维）
batch_size = tensor_3d.shape[0]  # 32

# 特征维度（通常是最后一维）
feature_dim = tensor_3d.shape[-1]  # 128

# 空间维度（图像的H, W）
spatial_dims = tensor_4d.shape[2:]  # (224, 224)

# 时间/序列维度
seq_len = tensor_3d.shape[1]  # 50

print(f"批次维度: {batch_size}")
print(f"特征维度: {feature_dim}")
print(f"空间维度: {spatial_dims}")
print(f"序列维度: {seq_len}")

# Transformer中的多种"维度"
class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048):
        super().__init__()
        # d_model: 模型维度（特征维度）
        # num_heads: 注意力头数
        # d_ff: 前馈网络的隐藏层维度

        self.attention = nn.MultiheadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),      # 特征维度: 512 -> 2048
            nn.ReLU(),
            nn.Linear(d_ff, d_model)       # 特征维度: 2048 -> 512
        )

    def forward(self, x):
        # x: [seq_len, batch, d_model]
        # seq_len: 序列维度
        # batch: 批次维度
        # d_model: 特征维度（最后一维）

        attn_out, _ = self.attention(x, x, x)
        return self.ffn(attn_out)

# 使用
transformer = TransformerBlock()
input_seq = torch.randn(100, 32, 512)  # [seq_len, batch, features]
output = transformer(input_seq)
print(f"Transformer输出: {output.shape}")

import torch
import torch.nn as nn

# Linear层总是对最后一维进行变换
linear = nn.Linear(128, 64)  # 128维特征 -> 64维特征

# 2D输入: [batch_size, features]
input_2d = torch.randn(32, 128)
output_2d = linear(input_2d)  # [32, 64]

# 3D输入: [batch_size, seq_len, features]
input_3d = torch.randn(32, 50, 128)
output_3d = linear(input_3d)  # [32, 50, 64] - 只变换最后一维

# 4D输入: [batch_size, height, width, features]
input_4d = torch.randn(32, 10, 20, 128)
output_4d = linear(input_4d)  # [32, 10, 20, 64] - 只变换最后一维

print(f"2D: {input_2d.shape} -> {output_2d.shape}")
print(f"3D: {input_3d.shape} -> {output_3d.shape}")
print(f"4D: {input_4d.shape} -> {output_4d.shape}")

# CNN中特征维度是通道维度，不是最后一维！
conv2d = nn.Conv2d(3, 64, kernel_size=3, padding=1)

# 输入: [batch, channels, height, width]
image = torch.randn(32, 3, 224, 224)
conv_output = conv2d(image)  # [32, 64, 224, 224]

print(f"CNN特征维度变化: {image.shape[1]} -> {conv_output.shape[1]}")
print(f"空间维度保持: {image.shape[2:]} -> {conv_output.shape[2:]}")

# 1D CNN也是如此
conv1d = nn.Conv1d(128, 256, kernel_size=3, padding=1)
sequence = torch.randn(32, 128, 100)  # [batch, channels, seq_len]
conv1d_output = conv1d(sequence)  # [32, 256, 100]
print(f"1D CNN: {sequence.shape} -> {conv1d_output.shape}")

# MultiheadAttention对最后一维进行变换
attention = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# 输入: [seq_len, batch, embed_dim] 或 [batch, seq_len, embed_dim]
query = torch.randn(50, 32, 512)    # [seq_len, batch, features]
key = torch.randn(50, 32, 512)
value = torch.randn(50, 32, 512)

attn_output, _ = attention(query, key, value)
print(f"Attention: {query.shape} -> {attn_output.shape}")
# 最后一维（特征维度）保持512不变

# Embedding - 输出的最后一维是特征维度
embedding = nn.Embedding(10000, 300)  # vocab_size=10000, embed_dim=300
token_ids = torch.randint(0, 10000, (32, 50))  # [batch, seq_len]
embedded = embedding(token_ids)  # [32, 50, 300] - 最后一维是特征

# LayerNorm - 对最后一维进行归一化
layer_norm = nn.LayerNorm(300)  # 对300维特征进行归一化
normalized = layer_norm(embedded)  # [32, 50, 300]

# Dropout - 保持所有维度不变
dropout = nn.Dropout(0.1)
dropped = dropout(embedded)  # [32, 50, 300]

# LSTM - 最后一维是特征维度
lstm = nn.LSTM(input_size=300, hidden_size=256, batch_first=True)
lstm_input = torch.randn(32, 50, 300)  # [batch, seq, features]
lstm_output, _ = lstm(lstm_input)  # [32, 50, 256] - 输出特征维度256

print(f"Embedding: {token_ids.shape} -> {embedded.shape}")
print(f"LSTM: {lstm_input.shape} -> {lstm_output.shape}")

# BatchNorm的特征维度位置取决于输入格式
# BatchNorm1d: 特征维度在第1维
bn1d = nn.BatchNorm1d(128)
input_1d = torch.randn(32, 128, 100)  # [batch, features, seq_len]
bn1d_output = bn1d(input_1d)  # 对第1维（128）进行归一化

# BatchNorm2d: 特征维度在第1维（通道维度）
bn2d = nn.BatchNorm2d(64)
input_2d = torch.randn(32, 64, 56, 56)  # [batch, channels, H, W]
bn2d_output = bn2d(input_2d)  # 对第1维（64个通道）进行归一化

print(f"BatchNorm1d特征维度: 第{1}维")
print(f"BatchNorm2d特征维度: 第{1}维")

class TransformerEncoder(nn.Module):
    def __init__(self, d_model=512, nhead=8, dim_feedforward=2048):
        super().__init__()
        # 所有层都以最后一维作为特征维度
        self.self_attn = nn.MultiheadAttention(d_model, nhead)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, src):
        # src: [seq_len, batch, d_model] - 最后一维是特征维度

        # Self-attention
        attn_output, _ = self.self_attn(src, src, src)
        src = self.norm1(src + attn_output)  # 残差连接

        # Feed-forward
        ff_output = self.linear2(torch.relu(self.linear1(src)))
        src = self.norm2(src + ff_output)  # 残差连接

        return src  # [seq_len, batch, d_model]

# 使用
encoder = TransformerEncoder()
input_seq = torch.randn(100, 32, 512)  # [seq_len, batch, features]
output = encoder(input_seq)
print(f"Transformer编码器: {input_seq.shape} -> {output.shape}")

import torch
import torch.nn as nn

# NLP中的标准张量格式
# [batch_size, sequence_length, feature_dimension]
text_tensor = torch.randn(32, 100, 512)  # 32个句子，每句100个词，每词512维

print(f"批次维度: {text_tensor.shape[0]}")      # 32 - 有多少个句子
print(f"序列长度维度: {text_tensor.shape[1]}")   # 100 - 每句有多少个词
print(f"特征维度: {text_tensor.shape[2]}")      # 512 - 每个词的向量维度

# 1. 词嵌入层：词ID -> 词向量
vocab_size, embed_dim = 10000, 512
embedding = nn.Embedding(vocab_size, embed_dim)

word_ids = torch.randint(0, vocab_size, (32, 100))  # [batch, seq_len]
word_vectors = embedding(word_ids)  # [32, 100, 512]
print(f"Embedding: {word_ids.shape} -> {word_vectors.shape}")

# 2. 线性变换：改变特征维度
linear1 = nn.Linear(512, 256)  # 512维 -> 256维
linear2 = nn.Linear(256, 128)  # 256维 -> 128维

transformed1 = linear1(word_vectors)  # [32, 100, 256]
transformed2 = linear2(transformed1)  # [32, 100, 128]

print(f"Linear1: {word_vectors.shape} -> {transformed1.shape}")
print(f"Linear2: {transformed1.shape} -> {transformed2.shape}")

# 3. LayerNorm：对最后一维进行归一化
layer_norm = nn.LayerNorm(128)
normalized = layer_norm(transformed2)  # [32, 100, 128]
print(f"LayerNorm: {transformed2.shape} -> {normalized.shape}")

# 4. 多头注意力：特征维度保持不变
attention = nn.MultiheadAttention(embed_dim=128, num_heads=8, batch_first=True)
attn_out, _ = attention(normalized, normalized, normalized)  # [32, 100, 128]
print(f"Attention: {normalized.shape} -> {attn_out.shape}")

class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size=10000, d_model=512, nhead=8, num_classes=2):
        super().__init__()
        # 所有操作都只改变最后一维
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Parameter(torch.randn(1000, d_model))  # 位置编码

        # Transformer层
        self.attention = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        # Feed-forward层
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * 4),  # 512 -> 2048
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)   # 2048 -> 512
        )

        # 分类头
        self.classifier = nn.Linear(d_model, num_classes)  # 512 -> 2

    def forward(self, x):
        # x: [batch, seq_len] - 词ID序列
        batch_size, seq_len = x.shape

        # 词嵌入: [batch, seq_len] -> [batch, seq_len, d_model]
        x = self.embedding(x)

        # 位置编码: 只对最后一维相加
        x = x + self.pos_encoding[:seq_len]

        # Self-attention: 保持 [batch, seq_len, d_model]
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_output)

        # Feed-forward: 只变换最后一维
        ff_output = self.ff(x)
        x = self.norm2(x + ff_output)  # [batch, seq_len, d_model]

        # 分类: 通常取[CLS]位置或平均池化
        # 这里简单取第一个位置
        cls_output = x[:, 0, :]  # [batch, d_model]

        # 最后的分类层: [batch, d_model] -> [batch, num_classes]
        return self.classifier(cls_output)

# 使用示例
model = SimpleTransformer()
input_ids = torch.randint(0, 10000, (16, 128))  # 16个句子，每句128个词
output = model(input_ids)  # [16, 2] - 16个句子的分类结果

print(f"模型输入: {input_ids.shape}")
print(f"模型输出: {output.shape}")

# 原因1: 保持序列结构
# 我们想保持"这是第几个句子"和"这是句子中的第几个词"的信息
text = torch.randn(4, 50, 300)  # 4个句子，每句50词，每词300维

# 如果改变序列长度维度，会破坏词的顺序关系
# 如果改变批次维度，会混淆不同的句子

# 原因2: 数学操作的自然性
linear = nn.Linear(300, 128)
# 矩阵乘法: [4, 50, 300] @ [300, 128] -> [4, 50, 128]
# 相当于对每个词向量分别进行线性变换

# 原因3: 并行计算效率
# GPU可以并行处理所有批次和所有位置的词向量变换

# 虽然通常不改变序列长度，但有些操作会间接影响
# 1. 池化操作
pooled = torch.mean(text, dim=1)  # [4, 50, 300] -> [4, 300]
print(f"平均池化: {text.shape} -> {pooled.shape}")

# 2. 注意力机制（Cross-attention）
query = torch.randn(4, 20, 300)    # 查询序列长度20
key_value = torch.randn(4, 50, 300) # 键值序列长度50
# 输出序列长度通常跟查询一致

# 3. 文本生成（自回归）
# 每次生成一个新词，序列长度逐渐增加

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

# 多头注意力的参数
d_model = 512    # 模型维度（输入特征维度）
num_heads = 8    # 注意力头数
d_k = d_model // num_heads  # 每个头的维度 = 512 // 8 = 64

print(f"模型维度 d_model: {d_model}")
print(f"注意力头数 num_heads: {num_heads}")
print(f"每头维度 d_k: {d_k}")

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 线性投影层
        self.w_q = nn.Linear(d_model, d_model)  # Q投影
        self.w_k = nn.Linear(d_model, d_model)  # K投影
        self.w_v = nn.Linear(d_model, d_model)  # V投影
        self.w_o = nn.Linear(d_model, d_model)  # 输出投影

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        print(f"=== 输入 ===")
        print(f"输入 x: {x.shape}")  # [batch_size, seq_len, d_model]
        print(f"含义: [批次={batch_size}, 序列长度={seq_len}, 特征维度={self.d_model}]")

        # 步骤1: 线性投影得到Q, K, V
        Q = self.w_q(x)  # [batch_size, seq_len, d_model]
        K = self.w_k(x)  # [batch_size, seq_len, d_model]
        V = self.w_v(x)  # [batch_size, seq_len, d_model]

        print(f"\n=== 线性投影后 ===")
        print(f"Q: {Q.shape}")
        print(f"K: {K.shape}")
        print(f"V: {V.shape}")
        print(f"含义: 每个都是 [批次, 序列长度, 模型维度]")

        return Q, K, V

# 测试
mha = MultiHeadAttention()
input_tensor = torch.randn(2, 10, 512)  # 2个句子，每句10个词，每词512维
Q, K, V = mha(input_tensor)

def reshape_for_multihead(tensor, batch_size, seq_len, num_heads, d_k):
    """将张量重塑为多头格式"""
    print(f"\n=== 重塑为多头 ===")
    print(f"重塑前: {tensor.shape}")  # [batch_size, seq_len, d_model]

    # 步骤1: 重塑为 [batch_size, seq_len, num_heads, d_k]
    reshaped = tensor.view(batch_size, seq_len, num_heads, d_k)
    print(f"重塑为4D: {reshaped.shape}")
    print(f"含义: [批次={batch_size}, 序列={seq_len}, 头数={num_heads}, 每头维度={d_k}]")

    # 步骤2: 转置为 [batch_size, num_heads, seq_len, d_k]
    transposed = reshaped.transpose(1, 2)
    print(f"转置后: {transposed.shape}")
    print(f"含义: [批次={batch_size}, 头数={num_heads}, 序列={seq_len}, 每头维度={d_k}]")

    return transposed

# 应用到Q, K, V
batch_size, seq_len = 2, 10
num_heads, d_k = 8, 64

Q_multi = reshape_for_multihead(Q, batch_size, seq_len, num_heads, d_k)
K_multi = reshape_for_multihead(K, batch_size, seq_len, num_heads, d_k)
V_multi = reshape_for_multihead(V, batch_size, seq_len, num_heads, d_k)

def compute_attention_scores(Q, K, V):
    """计算多头注意力分数"""
    print(f"\n=== 计算注意力分数 ===")

    # Q: [batch_size, num_heads, seq_len, d_k]
    # K: [batch_size, num_heads, seq_len, d_k]
    batch_size, num_heads, seq_len, d_k = Q.shape

    print(f"Q形状: {Q.shape}")
    print(f"K形状: {K.shape}")

    # 步骤1: Q @ K^T
    # Q: [batch, heads, seq_len, d_k] @ K^T: [batch, heads, d_k, seq_len]
    # = [batch, heads, seq_len, seq_len]
    scores = torch.matmul(Q, K.transpose(-2, -1))
    print(f"注意力分数 Q@K^T: {scores.shape}")
    print(f"含义: [批次={batch_size}, 头数={num_heads}, 查询长度={seq_len}, 键长度={seq_len}]")
    print("每个位置对每个位置的注意力原始分数")

    # 步骤2: 缩放
    scores = scores / math.sqrt(d_k)
    print(f"缩放后分数: {scores.shape} (数值除以√{d_k})")

    # 步骤3: Softmax
    attention_weights = F.softmax(scores, dim=-1)
    print(f"注意力权重: {attention_weights.shape}")
    print("含义: 每行和为1的概率分布")

    # 步骤4: 加权求和 V
    # attention_weights: [batch, heads, seq_len, seq_len]
    # V: [batch, heads, seq_len, d_k]
    # 结果: [batch, heads, seq_len, d_k]
    attention_output = torch.matmul(attention_weights, V)
    print(f"注意力输出: {attention_output.shape}")
    print(f"含义: [批次, 头数, 序列长度, 每头特征维度]")

    return attention_output, attention_weights

# 计算注意力
attention_output, attention_weights = compute_attention_scores(Q_multi, K_multi, V_multi)

def merge_multihead_output(attention_output):
    """合并多头注意力输出"""
    print(f"\n=== 合并多头输出 ===")

    # attention_output: [batch_size, num_heads, seq_len, d_k]
    batch_size, num_heads, seq_len, d_k = attention_output.shape
    print(f"多头输出: {attention_output.shape}")

    # 步骤1: 转置回 [batch_size, seq_len, num_heads, d_k]
    transposed = attention_output.transpose(1, 2)
    print(f"转置后: {transposed.shape}")

    # 步骤2: 重塑为 [batch_size, seq_len, d_model]
    # d_model = num_heads * d_k
    d_model = num_heads * d_k
    merged = transposed.contiguous().view(batch_size, seq_len, d_model)
    print(f"合并后: {merged.shape}")
    print(f"含义: [批次={batch_size}, 序列={seq_len}, 模型维度={d_model}]")

    return merged

# 合并输出
merged_output = merge_multihead_output(attention_output)

class CompleteMultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # 1. 线性投影
        Q = self.w_q(x)  # [batch, seq_len, d_model]
        K = self.w_k(x)
        V = self.w_v(x)

        # 2. 重塑为多头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        # 现在都是 [batch, num_heads, seq_len, d_k]

        # 3. 计算注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        # scores: [batch, num_heads, seq_len, seq_len]

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attention_weights = F.softmax(scores, dim=-1)
        attention_output = torch.matmul(attention_weights, V)
        # attention_output: [batch, num_heads, seq_len, d_k]

        # 4. 合并多头
        attention_output = attention_output.transpose(1, 2).contiguous()
        # [batch, seq_len, num_heads, d_k]

        attention_output = attention_output.view(batch_size, seq_len, self.d_model)
        # [batch, seq_len, d_model]

        # 5. 最终线性投影
        output = self.w_o(attention_output)
        # [batch, seq_len, d_model]

        return output, attention_weights

# 完整测试
print("=== 完整多头注意力测试 ===")
complete_mha = CompleteMultiHeadAttention(d_model=512, num_heads=8)
input_text = torch.randn(4, 20, 512)  # 4个句子，每句20词，每词512维

output, weights = complete_mha(input_text)
print(f"输入: {input_text.shape}")
print(f"输出: {output.shape}")
print(f"注意力权重: {weights.shape}")

print("=== 多头注意力维度变化总结 ===")
print("假设: batch_size=4, seq_len=20, d_model=512, num_heads=8, d_k=64")
print()
print("1. 输入:")
print("   x: [4, 20, 512] - [批次, 序列长度, 模型维度]")
print()
print("2. 线性投影:")
print("   Q, K, V: [4, 20, 512] - 保持输入形状")
print()
print("3. 重塑为多头:")
print("   Q, K, V: [4, 20, 512] -> [4, 20, 8, 64] -> [4, 8, 20, 64]")
print("   含义: [批次, 头数, 序列长度, 每头维度]")
print()
print("4. 注意力计算:")
print("   scores: [4, 8, 20, 64] @ [4, 8, 64, 20] = [4, 8, 20, 20]")
print("   含义: [批次, 头数, 查询位置, 键位置]")
print()
print("5. 加权求和:")
print("   output: [4, 8, 20, 20] @ [4, 8, 20, 64] = [4, 8, 20, 64]")
print("   含义: [批次, 头数, 序列长度, 每头维度]")
print()
print("6. 合并多头:")
print("   output: [4, 8, 20, 64] -> [4, 20, 8, 64] -> [4, 20, 512]")
print("   含义: [批次, 序列长度, 模型维度]")
print()
print("7. 输出投影:")
print("   final: [4, 20, 512] -> [4, 20, 512]")
print("   含义: 与输入形状完全相同")

import torch
import torch.nn as nn

# 假设参数
batch_size, seq_len, d_model, num_heads = 2, 4, 8, 2
d_k = d_model // num_heads  # 4

# 原始的Q矩阵
Q = torch.randn(batch_size, seq_len, d_model)
print(f"原始Q: {Q.shape}")  # [2, 4, 8]

# 重塑为多头格式
Q_reshaped = Q.view(batch_size, seq_len, num_heads, d_k)
print(f"重塑后Q: {Q_reshaped.shape}")  # [2, 4, 2, 4]
print(f"含义: [批次, 序列位置, 头数, 每头维度]")

# 关键：交换维度
Q_transposed = Q_reshaped.transpose(1, 2)
print(f"交换后Q: {Q_transposed.shape}")  # [2, 2, 4, 4]
print(f"含义: [批次, 头数, 序列位置, 每头维度]")

print("=== 不交换维度的问题 ===")

# 如果不交换维度，Q和K的形状是：
Q_no_transpose = torch.randn(2, 4, 2, 4)  # [batch, seq, heads, d_k]
K_no_transpose = torch.randn(2, 4, 2, 4)  # [batch, seq, heads, d_k]

print(f"不交换的Q: {Q_no_transpose.shape}")
print(f"不交换的K: {K_no_transpose.shape}")

# 尝试计算注意力分数 Q @ K^T
try:
    # 这样无法直接进行批量矩阵乘法！
    # 因为我们需要每个头分别计算，但头维度在中间
    scores_wrong = torch.matmul(Q_no_transpose, K_no_transpose.transpose(-1, -2))
    print("错误的计算方式会导致维度混乱")
except Exception as e:
    print(f"错误: {e}")

print("\n=== 交换维度的好处 ===")

# 交换后的Q和K
Q_correct = torch.randn(2, 2, 4, 4)  # [batch, heads, seq, d_k]
K_correct = torch.randn(2, 2, 4, 4)  # [batch, heads, seq, d_k]

print(f"交换后的Q: {Q_correct.shape}")
print(f"交换后的K: {K_correct.shape}")

# 现在可以轻松进行批量矩阵乘法
scores_correct = torch.matmul(Q_correct, K_correct.transpose(-1, -2))
print(f"正确的注意力分数: {scores_correct.shape}")  # [2, 2, 4, 4]
print("含义: [批次, 头数, 查询位置, 键位置]")

def visualize_computation():
    """可视化并行计算过程"""

    batch_size, num_heads, seq_len, d_k = 2, 2, 3, 4

    Q = torch.randn(batch_size, num_heads, seq_len, d_k)
    K = torch.randn(batch_size, num_heads, seq_len, d_k)

    print("=== 并行计算可视化 ===")
    print(f"Q形状: {Q.shape} - [批次, 头数, 序列, 特征]")
    print(f"K形状: {K.shape}")

    # 一次性计算所有头的注意力分数
    scores = torch.matmul(Q, K.transpose(-1, -2))
    print(f"分数形状: {scores.shape} - [批次, 头数, 查询位置, 键位置]")

    print("\n这意味着:")
    print(f"- 批次1头1: Q[0,0] @ K[0,0].T = scores[0,0] 形状{scores[0,0].shape}")
    print(f"- 批次1头2: Q[0,1] @ K[0,1].T = scores[0,1] 形状{scores[0,1].shape}")
    print(f"- 批次2头1: Q[1,0] @ K[1,0].T = scores[1,0] 形状{scores[1,0].shape}")
    print(f"- 批次2头2: Q[1,1] @ K[1,1].T = scores[1,1] 形状{scores[1,1].shape}")
    print("所有计算同时并行进行！")

visualize_computation()

def compare_computation_methods():
    """对比循环计算和并行计算"""

    batch_size, seq_len, d_model, num_heads = 2, 4, 8, 2
    d_k = d_model // num_heads

    # 原始输入
    Q_orig = torch.randn(batch_size, seq_len, d_model)
    K_orig = torch.randn(batch_size, seq_len, d_model)

    print("=== 方法1: 循环计算（如果不交换维度）===")

    # 重塑但不交换
    Q_no_swap = Q_orig.view(batch_size, seq_len, num_heads, d_k)
    K_no_swap = K_orig.view(batch_size, seq_len, num_heads, d_k)

    # 需要用循环分别计算每个头
    attention_scores_loop = []
    for batch in range(batch_size):
        batch_scores = []
        for head in range(num_heads):
            # 提取单个头的数据: [seq_len, d_k]
            q_head = Q_no_swap[batch, :, head, :]  # [4, 4]
            k_head = K_no_swap[batch, :, head, :]  # [4, 4]

            # 计算这个头的注意力分数
            score_head = torch.matmul(q_head, k_head.transpose(-1, -2))  # [4, 4]
            batch_scores.append(score_head)
        attention_scores_loop.append(torch.stack(batch_scores))

    loop_result = torch.stack(attention_scores_loop)  # [2, 2, 4, 4]
    print(f"循环计算结果: {loop_result.shape}")

    print("\n=== 方法2: 并行计算（交换维度）===")

    # 重塑并交换维度
    Q_swap = Q_orig.view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)
    K_swap = K_orig.view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)

    # 一次性并行计算所有头
    parallel_result = torch.matmul(Q_swap, K_swap.transpose(-1, -2))
    print(f"并行计算结果: {parallel_result.shape}")

    # 验证结果相同
    print(f"结果是否相同: {torch.allclose(loop_result, parallel_result)}")

    print(f"\n优势对比:")
    print(f"- 循环方法: 需要 {batch_size * num_heads} 次单独的矩阵乘法")
    print(f"- 并行方法: 只需要 1 次批量矩阵乘法")
    print(f"- GPU利用率: 并行方法能充分利用GPU的并行计算能力")

compare_computation_methods()

def memory_layout_explanation():
    """解释内存布局优化"""

    print("=== 内存布局优化 ===")

    # 交换前: [batch, seq, heads, d_k]
    tensor_before = torch.randn(2, 4, 2, 4)
    print(f"交换前形状: {tensor_before.shape}")
    print("内存中的数据排列（概念性）:")
    print("batch0_seq0_head0_features, batch0_seq0_head1_features,")
    print("batch0_seq1_head0_features, batch0_seq1_head1_features, ...")

    # 交换后: [batch, heads, seq, d_k]
    tensor_after = tensor_before.transpose(1, 2)
    print(f"\n交换后形状: {tensor_after.shape}")
    print("内存中的数据排列（概念性）:")
    print("batch0_head0_seq0_features, batch0_head0_seq1_features, ...")
    print("batch0_head1_seq0_features, batch0_head1_seq1_features, ...")

    print("\n优势:")
    print("1. 每个头的数据连续存储，访问效率更高")
    print("2. 批量矩阵乘法时，数据局部性更好")
    print("3. GPU的SIMD指令能更好地并行处理")

memory_layout_explanation()

import time

def performance_comparison():
    """性能对比实验"""

    # 较大的尺寸以体现性能差异
    batch_size, seq_len, d_model, num_heads = 32, 128, 512, 8
    d_k = d_model // num_heads

    Q = torch.randn(batch_size, seq_len, d_model)
    K = torch.randn(batch_size, seq_len, d_model)

    print("=== 性能对比实验 ===")
    print(f"测试规模: batch={batch_size}, seq={seq_len}, d_model={d_model}, heads={num_heads}")

    # 方法1: 正确的并行计算
    Q_parallel = Q.view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)
    K_parallel = K.view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)

    start_time = time.time()
    for _ in range(100):
        scores_parallel = torch.matmul(Q_parallel, K_parallel.transpose(-1, -2))
    parallel_time = time.time() - start_time

    print(f"并行计算时间: {parallel_time:.4f}秒 (100次)")
    print(f"平均每次: {parallel_time/100*1000:.2f}毫秒")

    # 结果形状
    print(f"输出形状: {scores_parallel.shape}")

performance_comparison()