形状后缀——良好编码风格

# 形状后缀——良好编码风格 来源：https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd Noam Shazeer (https://medium.com/@NoamShazeer?source=post_page---byline--f836e72e24fd---------------------------------------) 变量名应当既简洁又包含信息量。对于张量而言，最关键的莫过于它的维度数以及每个维度代表什么。自2022年起，我们Character.AI一直遵循这一惯例。不妨一试，看看是否会感觉更清晰： - 为逻辑维度规定单字母命名体系，例如用 `B` 代表批量大小，`L` 代表序列长度等，并在你的文件/项目/代码库的某处做文档说明。 - 在已知维度的情况下，张量名称应以由这些字母组成的维度后缀结尾，例如 `input_token_id_BL` 表示一个拥有 batch 和 length 维度的二维张量。 - 就这样简单。你可以在 torch、JAX 等框架中使用形状后缀。参见下面的示例。 ``` """ 带有形状后缀的 Transformer 示例代码。 此代码不完整且可能存在错误。请勿尝试运行它。 其目的是演示形状后缀。 维度键： B: 批量大小 L: 序列长度 M: 记忆长度（被关注序列的长度） D: 模型维度（有时称为 d_model 或 embedding_dim） V: 词汇量大小 F: 前馈子网络隐藏层大小 H: 层的注意力头数 K: 每个注意力键或值的大小（有时称为 d_kv） """ def transformer(input_token_id_BL, params): hidden_BLD = params.embedding_VD[input_token_id_BL] for layer_num in range(params.num_layers): hidden_BLD += attention(hidden_BLD, params.attention_params[i]) hidden_BLD += ffn(hidden_BLD, params.ffn_params[i]) hidden_BLD = layer_norm(hidden_BLD, params.final_layernorm_params) logits_BLV = torch.matmul(hidden_BLD, params.embedding_VD.T) return logits_BLV def ffn(input_BLD, params): input_BLD = layer_norm(input_BLD, params.layernorm_params) hidden_BLF = torch.gelu(torch.matmul(input_BLD, params.w_in_DF)) output_BLD = torch.matmul(hidden_BLF, params.w_out_FD) return output_BLD def attention(input_BLD, params): input_BLD = layer_norm(input_BLD, params.layernorm_params) query_BLHK = torch.einsum('BLD,DHK->BLHK', input_BLD, params.w_q_DHK) key_BMHK = torch.einsum('BLD,DHK->BLHK', input_BLD, params.w_k_DHK) value_BMHK = torch.einsum('BLD,DHK->BLHK', input_BLD, params.w_k_DHK) logits_BHLM = torch.einsum('BLHK,BMHK->BHLM', query_BLHK, key_BMHK) B, L, H, K = query_BLHK.shape() logits_BHLM /= K ** 0.5 masked_out_LM = torch.arange(L).unsqueeze(1) < torch.arange(L).unsqueeze(0) logits_BHLM += torch.where(masked_out_LM, -inf, 0) weights_BHLM = torch.softmax(logits_BHLM) wtd_values_BLHK = torch.einsum('BMHK,BHLM->BLHK', value_BMHK, logits_BHLM) out_BLD = torch.einsum('BLHK,HKD->BLD', wtd_values_BLHK, params.w_o_HKD) return out_BLD ```