Reformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "import torch\n",
    "import torch.nn as nn"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "emb = nn.Embedding(1000, 100)\n",
    "x = torch.LongTensor([[500, 12, 10],[10, 11, 265]])\n",
    "emb(x).shape"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 100])"
      ]
     },
     "metadata": {},
     "execution_count": 4
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    \"\"\"注意力机制的实现, 输入分别是query, key, value, mask: 掩码张量, \n",
    "       dropout是nn.Dropout层的实例化对象, 默认为None\"\"\"\n",
    "    # 在函数中, 首先取query的最后一维的大小, 一般情况下就等同于我们的词嵌入维度, 命名为d_k\n",
    "    d_k = query.size(-1)\n",
    "    # 按照注意力公式, 将query与key的转置相乘, 这里面key是将最后两个维度进行转置, 再除以缩放系数根号下d_k, 这种计算方法也称为缩放点积注意力计算.\n",
    "    # 得到注意力得分张量scores\n",
    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "\n",
    "    # 接着判断是否使用掩码张量\n",
    "    if mask is not None:\n",
    "        # 使用tensor的masked_fill方法, 将掩码张量和scores张量每个位置一一比较, 如果掩码张量处为0\n",
    "        # 则对应的scores张量用-1e9这个值来替换, 如下演示\n",
    "        scores = scores.masked_fill(mask == 0, -1e9)\n",
    "\n",
    "    # 对scores的最后一维进行softmax操作, 使用F.softmax方法, 第一个参数是softmax对象, 第二个是目标维度.\n",
    "    # 这样获得最终的注意力张量\n",
    "    p_attn = F.softmax(scores, dim = -1)\n",
    "\n",
    "    # 之后判断是否使用dropout进行随机置0\n",
    "    if dropout is not None:\n",
    "        # 将p_attn传入dropout对象中进行'丢弃'处理\n",
    "        p_attn = dropout(p_attn)\n",
    "\n",
    "    # 最后, 根据公式将p_attn与value张量相乘获得最终的query注意力表示, 同时返回注意力张量\n",
    "    return torch.matmul(p_attn, value), p_attn"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.7.6",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.7.6 64-bit ('base': conda)"
  },
  "interpreter": {
   "hash": "5f8e52cd05087fdb2a09ffc6084939d169f5962b33677882bb90cd2973b9af2a"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}