TransF_perLine.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformer模型Encoder原理精讲及其PyTorch逐行实现\n",
    "https://www.bilibili.com/video/BV1cP4y1V7GF?spm_id_from=333.337.search-card.all.click&vd_source=afe449886875b2cd7aa123878846a9f3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='2,3,4,5'\n",
    "import torch\n",
    "print(torch.cuda.device_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA_VISIBLE_DEVICES Count: 4\n",
      "[[4, 11, 5, 10, 6, 1], [8, 3, 12, 5, 7, 9, 2]]\n",
      "[[10, 12, 11, 3, 8, 6, 5, 1], [2, 7, 9, 7, 11, 13, 4, 1]]\n",
      "end!\n",
      "tensor([[ 4, 11,  5, 10,  6,  1,  0,  0],\n",
      "        [ 8,  3, 12,  5,  7,  9,  2,  0]], dtype=torch.int32)\n",
      "tensor([[10, 12, 11,  3,  8,  6,  5,  1,  0,  0],\n",
      "        [ 2,  7,  9,  7, 11, 13,  4,  1,  0,  0]], dtype=torch.int32)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='2,3,4,5'\n",
    "import torch\n",
    "print('CUDA_VISIBLE_DEVICES Count:',torch.cuda.device_count())\n",
    "\n",
    "\n",
    "import torch\n",
    "import numpy\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# 关于word embedding， 以序列建模为例\n",
    "\n",
    "src = ['I have a hd ,','but it is a big dog .']\n",
    "tgt = ['wo you yi ge hd ,','dan shi ta shi yi zhi gou .']\n",
    "src_len = 8  # for pad and position embedding \n",
    "tgt_len = 10  # for pad and position embedding \n",
    "max_pos_len = 20\n",
    "# word_pad = ['P', 'S', 'E']\n",
    "\n",
    "import numpy as np\n",
    "# 从原始的句子生成字典和序号输入\n",
    "def sentence2input(sentence):\n",
    "    '''\n",
    "    input: sentence\n",
    "    output: sequence, word2idx, idx2word\n",
    "    '''\n",
    "    source_split_words = [s.split(' ') for s in sentence]\n",
    "    word = []\n",
    "    for i in [j.split(' ') for j in sentence]:\n",
    "        word.extend(i)\n",
    "    # generate vocabulary =================\n",
    "    vocab = np.array(word)\n",
    "    vocab = np.unique(vocab)\n",
    "    idx2word = dict(enumerate(vocab,start=1))  # 从1开始因为pad填充0\n",
    "    word2idx = {v: k for k, v in idx2word.items()}\n",
    "    # word2idx = {w:i for i, w in enumerate(vocab)} \n",
    "    sequence = []\n",
    "    for n in source_split_words:\n",
    "        bs = [word2idx[w] for w in n]\n",
    "        sequence.append(bs)\n",
    "    return sequence, word2idx, idx2word\n",
    "\n",
    "\n",
    "def pad(input,max_len,pad_value=0):\n",
    "    '''\n",
    "    input: sentence\n",
    "    output: sequence, word2idx, idx2word\n",
    "    '''\n",
    "    import copy\n",
    "    pad_ = []\n",
    "    for i in input:\n",
    "        # print(i)\n",
    "        ii = copy.deepcopy(i)\n",
    "        if len(i)<max_len:\n",
    "            error_len = max_len - len(i)\n",
    "            for _ in range(error_len):\n",
    "                ii.append(pad_value)\n",
    "        pad_.append(ii[:max_len])\n",
    "    pad_ = torch.IntTensor(pad_)\n",
    "    return pad_\n",
    "\n",
    "src_input, src_vocab_word2idx, enc_vocab_idx2word = sentence2input(src)\n",
    "tgt_input, tgt_vocab_word2idx, tgt_vocab_idx2word = sentence2input(tgt)\n",
    "src_vocab_len = len(src_vocab_word2idx)  # 后面会用到\n",
    "tgt_vocab_len = len(tgt_vocab_word2idx)  # 后面会用到\n",
    "print(src_input)\n",
    "print(tgt_input)\n",
    "print('end!')\n",
    "src_input_T = pad(src_input,src_len,pad_value=0)\n",
    "tgt_input_T = pad(tgt_input,tgt_len,pad_value=0)\n",
    "print(src_input_T)\n",
    "print(tgt_input_T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 2.3609,  0.1995, -0.4773, -0.0060, -0.1113,  0.6070,  0.7216, -1.2232],\n",
      "        [ 0.0999,  0.3541, -0.0295,  0.8658, -0.1582,  0.9843,  1.1086, -0.8304],\n",
      "        [ 0.9446, -0.4625, -0.0565,  1.5416,  0.4657, -0.8963, -1.9402, -0.3984],\n",
      "        [ 0.8807,  1.0814, -0.0042,  0.4982,  0.6170, -1.7120,  0.0210,  1.2159],\n",
      "        [ 0.2230,  0.0111, -1.2859, -0.9706, -0.3395,  0.5946,  2.0974, -0.8544],\n",
      "        [ 0.6314,  1.2585, -0.2734, -0.2422,  0.3816, -0.7744,  0.2272,  0.6160],\n",
      "        [-1.5049, -0.1470, -1.6718, -2.2119,  1.1717, -2.1795,  0.1091,  1.5841],\n",
      "        [-0.8387, -0.6542, -1.2333,  0.0528,  1.3309,  0.5169, -0.6209,  0.9065],\n",
      "        [-0.3801, -0.8342, -1.5328, -0.9427,  1.2614,  0.0968, -0.5540,  0.8890],\n",
      "        [-0.7521, -0.0358,  0.6511, -1.1870,  0.6755, -0.7864,  0.7283,  0.1648],\n",
      "        [-0.5014, -1.1260, -0.3892,  0.3321, -0.2281,  0.3304,  2.7161, -0.1130],\n",
      "        [ 0.0074,  0.2132, -1.6190, -0.5002, -1.2446, -2.4929, -1.0454, -0.8808],\n",
      "        [-0.4476,  0.6607, -1.4005, -0.3224,  1.5685, -0.1988, -0.2547,  0.6091]],\n",
      "       requires_grad=True)\n",
      "tensor([[ 4, 11,  5, 10,  6,  1,  0,  0],\n",
      "        [ 8,  3, 12,  5,  7,  9,  2,  0]], dtype=torch.int32)\n",
      "tensor([[[ 0.2230,  0.0111, -1.2859, -0.9706, -0.3395,  0.5946,  2.0974,\n",
      "          -0.8544],\n",
      "         [ 0.0074,  0.2132, -1.6190, -0.5002, -1.2446, -2.4929, -1.0454,\n",
      "          -0.8808],\n",
      "         [ 0.6314,  1.2585, -0.2734, -0.2422,  0.3816, -0.7744,  0.2272,\n",
      "           0.6160],\n",
      "         [-0.5014, -1.1260, -0.3892,  0.3321, -0.2281,  0.3304,  2.7161,\n",
      "          -0.1130],\n",
      "         [-1.5049, -0.1470, -1.6718, -2.2119,  1.1717, -2.1795,  0.1091,\n",
      "           1.5841],\n",
      "         [ 0.0999,  0.3541, -0.0295,  0.8658, -0.1582,  0.9843,  1.1086,\n",
      "          -0.8304],\n",
      "         [ 2.3609,  0.1995, -0.4773, -0.0060, -0.1113,  0.6070,  0.7216,\n",
      "          -1.2232],\n",
      "         [ 2.3609,  0.1995, -0.4773, -0.0060, -0.1113,  0.6070,  0.7216,\n",
      "          -1.2232]],\n",
      "\n",
      "        [[-0.3801, -0.8342, -1.5328, -0.9427,  1.2614,  0.0968, -0.5540,\n",
      "           0.8890],\n",
      "         [ 0.8807,  1.0814, -0.0042,  0.4982,  0.6170, -1.7120,  0.0210,\n",
      "           1.2159],\n",
      "         [-0.4476,  0.6607, -1.4005, -0.3224,  1.5685, -0.1988, -0.2547,\n",
      "           0.6091],\n",
      "         [ 0.6314,  1.2585, -0.2734, -0.2422,  0.3816, -0.7744,  0.2272,\n",
      "           0.6160],\n",
      "         [-0.8387, -0.6542, -1.2333,  0.0528,  1.3309,  0.5169, -0.6209,\n",
      "           0.9065],\n",
      "         [-0.7521, -0.0358,  0.6511, -1.1870,  0.6755, -0.7864,  0.7283,\n",
      "           0.1648],\n",
      "         [ 0.9446, -0.4625, -0.0565,  1.5416,  0.4657, -0.8963, -1.9402,\n",
      "          -0.3984],\n",
      "         [ 2.3609,  0.1995, -0.4773, -0.0060, -0.1113,  0.6070,  0.7216,\n",
      "          -1.2232]]], grad_fn=<EmbeddingBackward0>)\n"
     ]
    }
   ],
   "source": [
    "# 构造word embedding\n",
    "model_dim = 8 # 512\n",
    "\n",
    "src_embedding_table = nn.Embedding(src_vocab_len+1,model_dim)  # 初始化一个embedding类，shape：num_embeddings: int, embedding_dim: int\n",
    "tgt_embedding_table = nn.Embedding(tgt_vocab_len+1,model_dim)  # 调用的是nn.Embedding类的forword方法，直接调用类后面一个括号就是调用该类中的forward方法\n",
    "print(src_embedding_table.weight)\n",
    "\n",
    "print(src_input_T)\n",
    "src_embedding = src_embedding_table(src_input_T)\n",
    "tgt_embedding = tgt_embedding_table(tgt_input_T)\n",
    "print(src_embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 构造position embeddings\n",
    "# src_len = 8 是input长度\n",
    "# 实例化nn.Embedding类\n",
    "# src_pos_embedding_table = nn.Embedding(src_len+1,model_dim)  # 初始化一个embedding类，shape：num_embeddings: int, embedding_dim: int\n",
    "# tgt_pos_embedding_table = nn.Embedding(tgt_len+1,model_dim)  # 调用的是nn.Embedding类的forword方法，直接调用类后面一个括号就是调用该类中的forward方法"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 构造position embeddings\n",
    "$$ PE(pos, 2i) = sin(pos/10000^{2i/d_model})  $$\n",
    "$$ PE(pos, 2i+1) = cos(pos/1000^{2i/d_model})  $$\n",
    "\n",
    " where $pos$ is the position and $i$ is the dimension. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0, 1, 2, 3, 4, 5, 6, 7],\n",
      "        [0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.int32)\n",
      "tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
      "        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], dtype=torch.int32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 8, 8]), torch.Size([2, 10, 8]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 构造position embeddings\n",
    "# src_len = 8 是input长度\n",
    "\n",
    "# 构造全长 position embedding table\n",
    "# max_pos_len = 20\n",
    "pos_mat = torch.arange(max_pos_len).reshape(-1,1)  # tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])\n",
    "i_mat = torch.pow(1000, torch.arange(0,model_dim,2).reshape(1,-1)/model_dim)  # tensor([[  1.0000,   5.6234,  31.6228, 177.8279]])\n",
    "pos_emb_table = torch.zeros(max_pos_len,model_dim)\n",
    "pos_emb_table[:,0::2] = torch.sin(pos_mat/i_mat)\n",
    "pos_emb_table[:,1::2] = torch.cos(pos_mat/i_mat)\n",
    "pos_embedding = nn.Embedding(max_pos_len,model_dim)\n",
    "# pos_embedding.weight  # nn.Embedding.weight随机初始化方式是标准正态分布，即均值μ=0，方差σ=1的正态分布。\n",
    "pos_embedding.weight = nn.Parameter(pos_emb_table,requires_grad=False)  # 这里是修改nn.Embedding类的初始化权重.weight，改为计算出的pos_emb_table\n",
    "#  torch.nn.Parameter是继承自torch.Tensor的子类，其主要作用是作为nn.Module中的可训练参数使用。它与torch.Tensor的区别就是nn.Parameter会自动被认为是module的可训练参数，即加入到parameter()这个迭代器中去；而module中非nn.Parameter()的普通tensor是不在parameter中的。\n",
    "\n",
    "# 获取 position count\n",
    "src_pos = [list(range(src_len)) for _ in src]  # 遍历样本src，src_len=8\n",
    "src_pos = torch.IntTensor(src_pos)\n",
    "print(src_pos)\n",
    "\n",
    "tgt_pos = [list(range(tgt_len)) for _ in src]  # 遍历样本src，tgt_len=10\n",
    "tgt_pos = torch.IntTensor(tgt_pos)\n",
    "print(tgt_pos)\n",
    "\n",
    "src_pos_embedding = pos_embedding(src_pos)  # src 和tgt 输入到一个全长的  position embedding table中\n",
    "tgt_pos_embedding = pos_embedding(tgt_pos)\n",
    "src_pos_embedding.size(), tgt_pos_embedding.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0, 0, 0, 0, 0, 0, 1, 1],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 1]], dtype=torch.int32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True]],\n",
       "\n",
       "        [[False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False,  True]]])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 构造encode 的 self attention mask\n",
    "# mask shape [batchsize, src_len,model_dim]\n",
    "# mask shape [batchsize, tgt_len,model_dim]\n",
    "src_real_len = [len(i.split(' ')) for i in src]\n",
    "valid_encoder_pos = torch.cat([torch.unsqueeze(F.pad(torch.zeros(L,dtype=torch.int32) ,pad=[0,src_len-L],value=1),0) for L in src_real_len],0)\n",
    "print(valid_encoder_pos) # 这是使用unsqueeze的方法 把1维变2维\n",
    "valid_encoder_pos = torch.unsqueeze(valid_encoder_pos,2) #需要再次unsqueeze 把2维变3维，这样才能使用bmm函数，保留batch size\n",
    "# print(valid_encoder_pos) \n",
    "valid_encoder_pos_metr = (torch.bmm(valid_encoder_pos, valid_encoder_pos.transpose(1,2)))  # 模拟QK.T相乘\n",
    "mask_encode_self_attention = valid_encoder_pos_metr.to(torch.bool)\n",
    "mask_encode_self_attention\n",
    "\n",
    "# 下面是reshape的方法\n",
    "# src_real_len = [len(i.split(' ')) for i in src]\n",
    "# valid_encoder_pos = torch.cat([F.pad(torch.zeros(L,dtype=torch.int32) ,pad=[0,src_len-L],value=1) for L in src_real_len],0)\n",
    "# valid_encoder_pos = valid_encoder_pos.reshape(len(src),-1)  # 这是使用reshpe的方法\n",
    "# print(valid_encoder_pos)\n",
    "# mask_encode_self_attention = valid_encoder_pos.to(torch.bool)\n",
    "# mask_encode_self_attention\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1, 1, 1, 1, 1, 1, 0, 0],\n",
      "        [1, 1, 1, 1, 1, 1, 1, 0]], dtype=torch.int32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True]],\n",
       "\n",
       "        [[False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False, False],\n",
       "         [False, False, False, False, False, False, False,  True]]])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 构造encode 的 self-ttention mask\n",
    "# mask shape [batchsize, src_len,model_dim]\n",
    "# mask shape [batchsize, tgt_len,model_dim]\n",
    "src_real_len = [len(i.split(' ')) for i in src]\n",
    "valid_encoder_pos = torch.cat([torch.unsqueeze(F.pad(torch.ones(L,dtype=torch.int32) ,pad=[0,src_len-L],value=0),0) for L in src_real_len],0)\n",
    "print(valid_encoder_pos) # 这是使用unsqueeze的方法 把1维变2维\n",
    "valid_encoder_pos = torch.unsqueeze(valid_encoder_pos,2) #需要再次unsqueeze 把2维变3维，这样才能使用bmm函数，保留batch size\n",
    "# print(valid_encoder_pos) \n",
    "\n",
    "\n",
    "invalid_encoder_pos = 1 - valid_encoder_pos\n",
    "\n",
    "\n",
    "\n",
    "invalid_encoder_pos_matrix = (torch.bmm(invalid_encoder_pos, invalid_encoder_pos.transpose(1,2)))  # 模拟QK.T相乘\n",
    "mask_encode_self_attention = invalid_encoder_pos_matrix.to(torch.bool)\n",
    "mask_encode_self_attention\n",
    "\n",
    "# 下面是reshape的方法\n",
    "# src_real_len = [len(i.split(' ')) for i in src]\n",
    "# valid_encoder_pos = torch.cat([F.pad(torch.zeros(L,dtype=torch.int32) ,pad=[0,src_len-L],value=1) for L in src_real_len],0)\n",
    "# valid_encoder_pos = valid_encoder_pos.reshape(len(src),-1)  # 这是使用reshpe的方法\n",
    "# print(valid_encoder_pos)\n",
    "# mask_encode_self_attention = valid_encoder_pos.to(torch.bool)\n",
    "# mask_encode_self_attention\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0289, 0.0186, 0.1601, 0.2390, 0.3970, 0.1229, 0.0101, 0.0236],\n",
       "         [0.0646, 0.0754, 0.1508, 0.1151, 0.2201, 0.0957, 0.0513, 0.2271],\n",
       "         [0.2041, 0.0515, 0.2022, 0.1600, 0.1166, 0.0529, 0.1599, 0.0527],\n",
       "         [0.0988, 0.0744, 0.1625, 0.1116, 0.0330, 0.1326, 0.3431, 0.0440],\n",
       "         [0.0885, 0.0339, 0.1148, 0.4421, 0.0137, 0.0439, 0.0954, 0.1676],\n",
       "         [0.0755, 0.1598, 0.0075, 0.0681, 0.1574, 0.1667, 0.1825, 0.1825],\n",
       "         [0.0793, 0.0506, 0.1599, 0.4487, 0.0338, 0.2278, 0.0000, 0.0000],\n",
       "         [0.0672, 0.0669, 0.3448, 0.0479, 0.1743, 0.2989, 0.0000, 0.0000]],\n",
       "\n",
       "        [[0.0370, 0.4836, 0.0765, 0.0292, 0.1137, 0.0568, 0.0910, 0.1122],\n",
       "         [0.1186, 0.0192, 0.0872, 0.2018, 0.0471, 0.1275, 0.0802, 0.3184],\n",
       "         [0.0174, 0.4050, 0.0492, 0.1051, 0.1201, 0.1937, 0.0527, 0.0567],\n",
       "         [0.0169, 0.0319, 0.2156, 0.2948, 0.1078, 0.1528, 0.1535, 0.0266],\n",
       "         [0.1405, 0.0512, 0.0181, 0.1362, 0.2283, 0.0783, 0.3102, 0.0371],\n",
       "         [0.3372, 0.0285, 0.0233, 0.1065, 0.2415, 0.1668, 0.0738, 0.0222],\n",
       "         [0.0881, 0.0784, 0.1026, 0.0651, 0.2360, 0.2674, 0.1239, 0.0386],\n",
       "         [0.2856, 0.0685, 0.0712, 0.1143, 0.2582, 0.1697, 0.0325, 0.0000]]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# mask 实施 测试\n",
    "\n",
    "score = torch.randn(len(src),src_len,src_len)\n",
    "score\n",
    "masked_score = score.masked_fill(mask_encode_self_attention, -np.inf)\n",
    "masked_score\n",
    "F.softmax(masked_score,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 8, 1])\n",
      "torch.Size([2, 10, 1])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True]],\n",
       "\n",
       "        [[False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [False, False, False, False, False, False, False,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True]]])"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 构造 intra attention mask \n",
    "src_real_len = [len(i.split(' ')) for i in src]\n",
    "tgt_real_len = [len(i.split(' ')) for i in tgt]\n",
    "valid_encoder_pos = torch.unsqueeze(torch.cat([torch.unsqueeze(F.pad(torch.ones(L,dtype=torch.int32) ,pad=[0,src_len-L],value=0),0) for L in src_real_len],0),2)\n",
    "print(valid_encoder_pos.size())\n",
    "valid_decoder_pos = torch.unsqueeze(torch.cat([torch.unsqueeze(F.pad(torch.ones(L,dtype=torch.int32) ,pad=[0,tgt_len-L],value=0),0) for L in tgt_real_len],0),2)\n",
    "print(valid_decoder_pos.size())\n",
    "valid_cross_pos_matrix = torch.bmm(valid_decoder_pos, valid_encoder_pos.transpose(1,2))  #注意这里 tgt-decode是Q，src-encoed是K，计算公式是Q和K转置矩阵乘\n",
    "valid_cross_pos_matrix.size()\n",
    "invalid_cross_pos_matrix = 1 - valid_cross_pos_matrix\n",
    "mask_cross_self_attention = invalid_cross_pos_matrix.to(torch.bool)\n",
    "mask_cross_self_attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 10, 10])\n"
     ]
    }
   ],
   "source": [
    "# 构造 decode attention mask -- 有点不一样\n",
    "# 构造下三角矩阵 为什么？ 因为decoder的输入要把后面的数据mask掉 每次输入的数据向后移一位 即看到的范围越来越大\n",
    "m_ = [torch.unsqueeze(F.pad(torch.tril(torch.ones((L,L))),(0,(tgt_len-L),0,(tgt_len-L))),0)  for L in tgt_real_len]\n",
    "print(m_[0].size())\n",
    "valide_decoder_tri_matrix = torch.cat(m_)\n",
    "invalide_decoder_tri_matrix = 1 - valide_decoder_tri_matrix\n",
    "invalide_decoder_tri_matrix = invalide_decoder_tri_matrix.to(torch.bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[False,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False, False, False,  True,  True,  True],\n",
       "         [False, False, False, False, False, False, False, False,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
       "\n",
       "        [[False,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False,  True,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False,  True,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False, False,  True,  True,  True,  True],\n",
       "         [False, False, False, False, False, False, False,  True,  True,  True],\n",
       "         [False, False, False, False, False, False, False, False,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]]])"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invalide_decoder_tri_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.4672, 0.5328, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.1145, 0.2677, 0.6178, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.1166, 0.2938, 0.0396, 0.5500, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.2554, 0.0477, 0.1712, 0.2855, 0.2401, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.2036, 0.0844, 0.2609, 0.1679, 0.0591, 0.2242, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.0604, 0.1446, 0.2196, 0.0840, 0.3627, 0.0892, 0.0395, 0.0000, 0.0000,\n",
       "         0.0000],\n",
       "        [0.1445, 0.0139, 0.1103, 0.1852, 0.1359, 0.1485, 0.2378, 0.0239, 0.0000,\n",
       "         0.0000],\n",
       "        [0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000,\n",
       "         0.1000],\n",
       "        [0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000,\n",
       "         0.1000]])"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# mask 实施 测试\n",
    "\n",
    "score = torch.randn(len(src),tgt_len,tgt_len)\n",
    "score\n",
    "masked_score = score.masked_fill(invalide_decoder_tri_matrix, -1e9)\n",
    "# print(masked_score)\n",
    "F.softmax(masked_score,-1)[0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 构造scaled self attention 函数\n",
    "def scale_dot_product_attention(Q,K,V,attention_mask):  # Q K 是 batch的， 同时，也是batch * multi-head的\n",
    "    #  shape of Q K V : [batch_size,num_head,src_len,model_dim/num_head]\n",
    "    score = torch.bmm(Q, K.transpose(-1,-2))/torch.sqrt(model_dim)\n",
    "    masked_score = score.masked_fill(attention_mask, -1e9)\n",
    "    prob = F.softmax(masked_score,-1)\n",
    "    context = torch.bmm(prob,V)\n",
    "    return context\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = nn.Transformer(nhead=2, num_encoder_layers=6)\n",
    "src_ = torch.rand((10,32,512))\n",
    "tgt_ = torch.rand((20,32,512))\n",
    "out = transformer(src_,tgt_)\n",
    "out.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([20, 32, 512])"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# # 4 大类\n",
    "# TransformerEncoder\n",
    "#     TransformerEncoderLayer\n",
    "\n",
    "# TransformerDecoder\n",
    "#     TransformerDe  coderLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.3064, 1.4692, 0.0000],\n",
       "        [2.3875, 0.3449, 2.5509]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Transformer Masked loss\n",
    "# 使用机器翻译任务 模拟\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F \n",
    "\n",
    "# 两个句子，每个句子3个单词，那就是一共6个单词\n",
    "Logits = torch.randn(2,3,4) # bs=2,seqlen=3,vocab_size=class=4 \n",
    "Logits = Logits.transpose(1,2)  # 在torch 官方crossentropy调用的时候需要的维度是bs，C，d 没有经过softmax之前\n",
    "lable = torch.randint(0,4,(2,3))\n",
    "Logits.size(),lable.size()\n",
    "# 这里的交叉熵是6个单词的平均的交叉熵\n",
    "F.cross_entropy(Logits,lable) \n",
    "# 返回每一个单词的交叉熵\n",
    "F.cross_entropy(Logits,lable,reduction='none') \n",
    "# 构造mask len 即假设tgt的有效长度只有2，有1位是padding的，那就需要把这一位mask掉再进行loss掉统计\n",
    "tgt_len = torch.Tensor([2,3]).to(torch.int32) # 假设两个样本长度为2，3\n",
    "\n",
    "valid_mask = torch.cat([torch.unsqueeze(F.pad(torch.ones(l),(0,max(tgt_len)-l)),0) for l in tgt_len])\n",
    "F.cross_entropy(Logits,lable,reduction='none') * valid_mask\n",
    "\n",
    "# 或者\n",
    "lable[0,2] = -100\n",
    "F.cross_entropy(Logits,lable,reduction='none')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[-4.1631e-01, -1.1203e+00,  7.8211e-01],\n",
       "          [-1.0408e+00, -2.0775e+00,  1.8497e+00],\n",
       "          [ 8.0159e-01, -3.4229e-01,  2.1815e+00],\n",
       "          [ 8.0410e-01, -4.7673e-01,  1.9605e+00]],\n",
       " \n",
       "         [[ 4.2155e-01,  6.3033e-01, -1.6153e+00],\n",
       "          [-5.7698e-02, -9.6715e-01, -1.7739e+00],\n",
       "          [-1.4661e-01, -4.5681e-02,  1.0823e-03],\n",
       "          [-4.9995e-01,  8.1735e-01, -4.4952e-01]]]),\n",
       " tensor([[3, 1, 1],\n",
       "         [2, 2, 3]]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Logits,lable\n",
    "# tensor([[[ 1.5035,  1.1617,  0.8024,  0.4093],\n",
    "#          [ 0.0355, -0.5075, -1.2482, -0.7190],\n",
    "#          [ 0.5320,  1.7259,  0.4466,  0.6156]],\n",
    "\n",
    "# tensor([[3, 1, 1],\n",
    "#          [2, 2, 3]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2, 3, 0],\n",
       "        [1, 1, 3]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8964, 2.6555, 0.0000],\n",
       "        [1.5174, 1.7465, 1.1464]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8964, 2.6555, 0.0000],\n",
       "        [1.5174, 1.7465, 0.0000]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('pytorch')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "016fd9bb66e6ad19daf5b878b9ae392875dbee4b6103c0b67cb7b40f3b5f6cf6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}