CMRC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5f2aa98-2611-41de-9220-908e20f55e03",
   "metadata": {},
   "source": [
    "# CMRC\n",
    "1. 原始官方数据：https://storage.googleapis.com/cluebenchmark/tasks/cmrc2018_public.zip (`ChineseBERT-Paddle/data/CMRC/raw_data`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd5caa29-0a83-47ee-ac3f-1d9a0e5a06e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/data1/workspace/research/ChineseBERT-Paddle/Paddle_ChineseBert/PaddleNLP')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "80ca5170-266b-4277-a23f-02aea00d818d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from paddlenlp.datasets import load_dataset\n",
    "import json\n",
    "\n",
    "def read(data_path):\n",
    "    if not isinstance(data_path, list):\n",
    "        data_path = [data_path]\n",
    "    for file in data_path:\n",
    "        with open(file, \"r\", encoding=\"utf8\") as f:\n",
    "            input_data = json.load(f)[\"data\"]\n",
    "            for entry in input_data:\n",
    "                title = entry.get(\"title\", \"\").strip()\n",
    "                for paragraph in entry[\"paragraphs\"]:\n",
    "                    context = paragraph[\"context\"].strip()\n",
    "                    for qa in paragraph[\"qas\"]:\n",
    "                        qas_id = qa[\"id\"]\n",
    "                        question = qa[\"question\"].strip()\n",
    "                        answer_starts = [\n",
    "                            answer[\"answer_start\"]\n",
    "                            for answer in qa.get(\"answers\", [])\n",
    "                        ]\n",
    "                        answers = [\n",
    "                            answer[\"text\"].strip()\n",
    "                            for answer in qa.get(\"answers\", [])\n",
    "                        ]\n",
    "\n",
    "                        yield {\n",
    "                            'id': qas_id,\n",
    "                            'title': title,\n",
    "                            'context': context,\n",
    "                            'question': question,\n",
    "                            'answers': answers,\n",
    "                            'answer_starts': answer_starts\n",
    "                        }\n",
    "\n",
    "# data_path为read()方法的参数\n",
    "train_ds = load_dataset(read, data_path=['./data/CMRC/raw_data/train.json', './data/CMRC/raw_data/trial.json'], lazy=False)\n",
    "dev_ds = load_dataset(read, data_path='./data/CMRC/raw_data/dev.json', lazy=False)\n",
    "test_ds = load_dataset(read, data_path='./data/CMRC/raw_data/test.json', lazy=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b5fd89a7-3ea4-4b01-8903-e832e8f13675",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11144"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "868c08c7-cf0b-4ee6-bcfb-075d2eee096b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': 'TEST_1089_QUERY_4', 'title': '蓝鸿震', 'context': '蓝鸿震（，人称蓝爷，籍贯广东大埔，香港民政事务局前局长。小学就读德信学校，中学毕业于圣芳济书院。于罗富国教育学院(今香港教育大学，教育学院前身)毕业后，在福华街实用中学执教。汇丰控股首位华人执行董事郑海泉和前立法会议员郑经翰都是其学生。1965年取得伦敦大学文学士学位及特许秘书资历，之后加入医务衞生署，成为伊利沙伯医院的助理医院秘书。1966年转职政务主任，曾出任的高层职位包括：教育署副署长（1987年6月－1988年9月）、港九政务署长（1988年9月－1991年9月）及香港驻东京经济贸易首席代表（1991年10月－1997年9月），官至民政事务局局长。退休后，曾以新世纪论坛成员身份伙拍区议员陈财喜和学协成员杨沁瑜参选2000年香港立法会选举香港岛选区，是继黄钱其濂后第二位角逐立法机构议席的政府高官，但最终落败。现为建制派人士，曾任第十届及第十一届全国政协委员。', 'question': '蓝鸿震曾出任的高层职位包括哪些？', 'answers': ['FAKE_ANSWER_1', 'FAKE_ANSWER_2', 'FAKE_ANSWER_3'], 'answer_starts': [-1, -1, -1]}\n",
      "{'id': 'TEST_1091_QUERY_4', 'title': '仓后街道', 'context': '，是一个已撤销的街道办事处，是江门的发源地，总面积5平方公里，人口6万。仓后街道位于蓬江区南部，为江门繁华中心地带，北与环市街道相连，东南隔江门河与江海区滘头街道相望，东与堤东街道办事处接壤，西与白沙街道办事处毗邻。明、清时，仓后辖地为新会县归德都管辖，民国时，属新会县第二区，1951年属江门市并建区（镇级），驻地在仓后路，故名。1959年夏，仓后与垆顶两区合并为中区公社，1962年改称仓后区，1963年复置仓后办事处，1969年3月改称仓后区革命委员会，1976年11月改称仓后街道办事处。2015年，江门市人民政府批复同意蓬江区调整部分街道行政区划，撤销仓后街道，将其行政区划（五邑大学校园除外）并入白沙街道。', 'question': '2015年，后仓街道并入哪个街道？', 'answers': ['FAKE_ANSWER_1', 'FAKE_ANSWER_2', 'FAKE_ANSWER_3'], 'answer_starts': [-1, -1, -1]}\n",
      "{'id': 'TEST_1092_QUERY_2', 'title': '北街街道 (江门市)', 'context': '，是一个已撤销的街道办事处，总面积5平方公里，人口3.4万。北街街道位于蓬江区东部，位处西江与江门河交汇处，为江门的门户和军事重地，北隔西江与潮连街道相望，东隔江门河与滘北街道相望，西与堤东、环市街道办事处毗邻。北街原为西江河道，约在元末明初淤积成陆，成陆后为新会县归德都白石乡辖地。1904年，英国强迫清政府开辟江门为通商口岸，并在北街建立海关。1912年2月16日15时15分，北街海关升起新会第一面青天白日满地红旗。1913年，新宁铁路在北街建火车站，北街日渐兴盛，逐渐成为新会水陆交通转运地和华洋荟萃之地。1951年江门设市后，1952年初成立北街街道办事处。2015年，江门市人民政府批复同意蓬江区调整部分街道行政区划，撤销北街街道，将其行政区划（港口二路－白石大道－甘棠路－发展大道－西江边－港口二路围闭区域除外）并入白沙街道。', 'question': '北街街道对于江门而言有什么战略意义？', 'answers': ['FAKE_ANSWER_1', 'FAKE_ANSWER_2', 'FAKE_ANSWER_3'], 'answer_starts': [-1, -1, -1]}\n",
      "{'id': 'TEST_1093_QUERY_0', 'title': '米尔斯常数', 'context': '米尔斯常数是使对于所有正整数\"n\"，二重指数函数的整数部分都是素数的最小正实数\"A\"。这个常数以W·H·米尔斯命名，他在1947年证明了这个常数的存在。米尔斯常数的值是未知的，但如果黎曼猜想成立，它的值大约为：由米尔斯常数所产生的素数称为米尔斯素数；如果黎曼猜想成立，这个数列的最初几项为：如果用\"a\"(\"i\")来表示数列中的第\"i\"个素数，则\"a\"(\"i\")可以计算为大于\"a\"(\"i\"\\xa0−1)的最小的素数。为了保证当\"n\" = 1，2，3，……时，\"A\"的整数部分是这个素数数列，必须有\"a\"(\"i\")\\xa0<\\xa0(\"a\"(\"i\"\\xa0−1)\\xa0+\\xa01)。Hoheisel和Ingham的结果保证了在任何两个足够大的立方数之间一定有一个素数，这足以证明这个不等式，如果我们从一个足够大的素数\"a\"(1)开始。从黎曼猜想，可以推出任何两个连续的立方数之间一定有一个素数，这样就可以去掉\"足够大\"的条件，并允许米尔斯素数的数列从\"a\"(1) = 2开始。目前已知最大的米尔斯素数（假设黎曼猜想成立）是：它有20,562位。通过计算米尔斯素数，我们可以近似计算米尔斯常数为：用这个方法计算出米尔斯常数的差不多七千位数。目前还没有闭合公式可以计算米尔斯常数，甚至不知道它是不是有理数。', 'question': '米尔斯常数是指什么？', 'answers': ['FAKE_ANSWER_1', 'FAKE_ANSWER_2', 'FAKE_ANSWER_3'], 'answer_starts': [-1, -1, -1]}\n"
     ]
    }
   ],
   "source": [
    "for example in test_ds:\n",
    "    if example['id'] == 'TEST_1089_QUERY_4':\n",
    "        print(example)\n",
    "    if example['id'] == \"TEST_1091_QUERY_4\":\n",
    "        print(example)\n",
    "    if example['id'] == \"TEST_1092_QUERY_2\":\n",
    "        print(example)\n",
    "    if example['id'] == \"TEST_1093_QUERY_0\":\n",
    "        print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "e017a16e-7033-4795-9ad1-553a43cfa515",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config:\n",
    "    adam_epsilon=1e-08\n",
    "    batch_size=4\n",
    "    device='gpu'\n",
    "    do_lower_case=True\n",
    "    do_predict=True\n",
    "    do_train=True\n",
    "    doc_stride=128\n",
    "    learning_rate=3e-05\n",
    "    logging_steps=10\n",
    "    max_answer_length=30\n",
    "    max_grad_norm=1.0\n",
    "    max_query_length=64\n",
    "    max_seq_length=384\n",
    "    max_steps=-1\n",
    "    model_name_or_path='bert-base-chinese'\n",
    "    model_type='bert'\n",
    "    n_best_size=20\n",
    "    num_train_epochs=10\n",
    "    output_dir='./data/tmp/chinesebert/'\n",
    "    predict_file=None\n",
    "    save_steps=1000\n",
    "    seed=42\n",
    "    task_name='cmrc2018'\n",
    "    train_file=None\n",
    "    verbose=False\n",
    "    warmup_proportion=0.1\n",
    "    weight_decay=0.01\n",
    "\n",
    "args = Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "c04e03f3-9bf0-4e8a-91b6-f10ed5c2775c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import time\n",
    "import json\n",
    "import math\n",
    "\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "import paddle\n",
    "\n",
    "from paddle.io import DataLoader\n",
    "\n",
    "import paddlenlp as ppnlp\n",
    "\n",
    "from paddlenlp.data import Pad, Stack, Tuple, Dict\n",
    "from paddlenlp.transformers import BertForQuestionAnswering, BertTokenizer\n",
    "from paddlenlp.transformers import ErnieForQuestionAnswering, ErnieTokenizer\n",
    "from paddlenlp.transformers import ErnieGramForQuestionAnswering, ErnieGramTokenizer\n",
    "from paddlenlp.transformers import RobertaForQuestionAnswering, RobertaTokenizer\n",
    "from paddlenlp.transformers import GlyceBertForQuestionAnswering, ChineseBertTokenizer\n",
    "from paddlenlp.transformers import LinearDecayWithWarmup\n",
    "from paddlenlp.metrics.squad import squad_evaluate, compute_prediction\n",
    "from paddlenlp.datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "456d12d2-9520-4e7b-ba77-5f17572b6af7",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_CLASSES = {\n",
    "    \"bert\": (BertForQuestionAnswering, BertTokenizer),\n",
    "    \"ernie\": (ErnieForQuestionAnswering, ErnieTokenizer),\n",
    "    \"ernie_gram\": (ErnieGramForQuestionAnswering, ErnieGramTokenizer),\n",
    "    \"roberta\": (RobertaForQuestionAnswering, RobertaTokenizer)\n",
    "}\n",
    "\n",
    "def set_seed(args):\n",
    "    random.seed(args.seed)\n",
    "    np.random.seed(args.seed)\n",
    "    paddle.seed(args.seed)\n",
    "\n",
    "\n",
    "@paddle.no_grad()\n",
    "def evaluate(model, data_loader, args, epoch, dataset_name=\"dev\"):\n",
    "    model.eval()\n",
    "\n",
    "    all_start_logits = []\n",
    "    all_end_logits = []\n",
    "    tic_eval = time.time()\n",
    "\n",
    "    for batch in data_loader:\n",
    "        input_ids, pinyin_ids, token_type_ids = batch\n",
    "        start_logits_tensor, end_logits_tensor = model(input_ids,\n",
    "                                                       pinyin_ids,\n",
    "                                                       token_type_ids)\n",
    "\n",
    "        for idx in range(start_logits_tensor.shape[0]):\n",
    "            if len(all_start_logits) % 1000 == 0 and len(all_start_logits):\n",
    "                print(\"Processing example: %d\" % len(all_start_logits))\n",
    "                print('time per 1000:', time.time() - tic_eval)\n",
    "                tic_eval = time.time()\n",
    "\n",
    "            all_start_logits.append(start_logits_tensor.numpy()[idx])\n",
    "            all_end_logits.append(end_logits_tensor.numpy()[idx])\n",
    "\n",
    "    all_predictions, _, _ = compute_prediction(\n",
    "        data_loader.dataset.data, data_loader.dataset.new_data,\n",
    "        (all_start_logits, all_end_logits), False, args.n_best_size,\n",
    "        args.max_answer_length)\n",
    "\n",
    "    # Can also write all_nbest_json and scores_diff_json files if needed\n",
    "    with open(f'{dataset_name}_prediction_{epoch}.json', \"w\", encoding='utf-8') as writer:\n",
    "        writer.write(\n",
    "            json.dumps(\n",
    "                all_predictions, ensure_ascii=False, indent=4) + \"\\n\")\n",
    "\n",
    "    squad_evaluate(\n",
    "        examples=data_loader.dataset.data,\n",
    "        preds=all_predictions,\n",
    "        is_whitespace_splited=False)\n",
    "\n",
    "    model.train()\n",
    "\n",
    "\n",
    "class CrossEntropyLossForSQuAD(paddle.nn.Layer):\n",
    "    def __init__(self):\n",
    "        super(CrossEntropyLossForSQuAD, self).__init__()\n",
    "\n",
    "    def forward(self, y, label):\n",
    "        start_logits, end_logits = y\n",
    "        start_position, end_position = label\n",
    "        start_position = paddle.unsqueeze(start_position, axis=-1)\n",
    "        end_position = paddle.unsqueeze(end_position, axis=-1)\n",
    "        start_loss = paddle.nn.functional.cross_entropy(\n",
    "            input=start_logits, label=start_position)\n",
    "        end_loss = paddle.nn.functional.cross_entropy(\n",
    "            input=end_logits, label=end_position)\n",
    "        loss = (start_loss + end_loss) / 2\n",
    "        return loss\n",
    "\n",
    "\n",
    "def run(args):\n",
    "    paddle.set_device(args.device)\n",
    "    if paddle.distributed.get_world_size() > 1:\n",
    "        paddle.distributed.init_parallel_env()\n",
    "    rank = paddle.distributed.get_rank()\n",
    "\n",
    "    task_name = args.task_name.lower()\n",
    "    args.model_type = args.model_type.lower()\n",
    "    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]\n",
    "    # tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)\n",
    "    CHINESEBERT_PADDLE_PATH = \"./pretrain_models/paddle/ChineseBERT-large\"\n",
    "    tokenizer = ChineseBertTokenizer(CHINESEBERT_PADDLE_PATH)\n",
    "    set_seed(args)\n",
    "\n",
    "#     train_ds = load_dataset(\n",
    "#         task_name, splits=\"train\", data_files=args.train_file)\n",
    "#     dev_ds = load_dataset(task_name, splits=\"dev\", data_files=args.predict_file)\n",
    "\n",
    "    if rank == 0:\n",
    "        if os.path.exists(args.model_name_or_path):\n",
    "            print(\"init checkpoint from %s\" % args.model_name_or_path)\n",
    "\n",
    "    # model = model_class.from_pretrained(args.model_name_or_path)\n",
    "    model = GlyceBertForQuestionAnswering.from_pretrained(CHINESEBERT_PADDLE_PATH)\n",
    "\n",
    "    if paddle.distributed.get_world_size() > 1:\n",
    "        model = paddle.DataParallel(model)\n",
    "\n",
    "    def prepare_train_features(examples):\n",
    "        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results\n",
    "        # in one example possible giving several features when a context is long, each of those features having a\n",
    "        # context that overlaps a bit the context of the previous feature.\n",
    "        # NOTE: Almost the same functionality as HuggingFace's prepare_train_features function. The main difference is\n",
    "        # that HugggingFace uses ArrowTable as basic data structure, while we use list of dictionary instead.\n",
    "        contexts = [examples[i]['context'] for i in range(len(examples))]\n",
    "        questions = [examples[i]['question'] for i in range(len(examples))]\n",
    "\n",
    "        tokenized_examples = tokenizer(\n",
    "            questions,\n",
    "            contexts,\n",
    "            stride=args.doc_stride,\n",
    "            max_seq_len=args.max_seq_length)\n",
    "\n",
    "        # Let's label those examples!\n",
    "        for i, tokenized_example in enumerate(tokenized_examples):\n",
    "            # We will label impossible answers with the index of the CLS token.\n",
    "            input_ids = tokenized_example[\"input_ids\"]\n",
    "            cls_index = input_ids.index(tokenizer.cls_token_id)\n",
    "\n",
    "            # 引入 pinyin_ids\n",
    "            pinyin_ids = tokenizer.convert_tokens_to_pinyin_ids(tokenizer.convert_ids_to_tokens(input_ids))\n",
    "            tokenized_examples[i][\"pinyin_ids\"] = pinyin_ids\n",
    "\n",
    "            # The offset mappings will give us a map from token to character position in the original context. This will\n",
    "            # help us compute the start_positions and end_positions.\n",
    "            offsets = tokenized_example['offset_mapping']\n",
    "\n",
    "            # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "            sequence_ids = tokenized_example['token_type_ids']\n",
    "\n",
    "            # One example can give several spans, this is the index of the example containing this span of text.\n",
    "            sample_index = tokenized_example['overflow_to_sample']\n",
    "            answers = examples[sample_index]['answers']\n",
    "            answer_starts = examples[sample_index]['answer_starts']\n",
    "\n",
    "            # Start/end character index of the answer in the text.\n",
    "            start_char = answer_starts[0]\n",
    "            end_char = start_char + len(answers[0])\n",
    "\n",
    "            # Start token index of the current span in the text.\n",
    "            token_start_index = 0\n",
    "            while sequence_ids[token_start_index] != 1:\n",
    "                token_start_index += 1\n",
    "\n",
    "            # End token index of the current span in the text.\n",
    "            token_end_index = len(input_ids) - 1\n",
    "            while sequence_ids[token_end_index] != 1:\n",
    "                token_end_index -= 1\n",
    "            # Minus one more to reach actual text\n",
    "            token_end_index -= 1\n",
    "\n",
    "            # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).\n",
    "            if not (offsets[token_start_index][0] <= start_char and\n",
    "                    offsets[token_end_index][1] >= end_char):\n",
    "                tokenized_examples[i][\"start_positions\"] = cls_index\n",
    "                tokenized_examples[i][\"end_positions\"] = cls_index\n",
    "            else:\n",
    "                # Otherwise move the token_start_index and token_end_index to the two ends of the answer.\n",
    "                # Note: we could go after the last offset if the answer is the last word (edge case).\n",
    "                while token_start_index < len(offsets) and offsets[\n",
    "                        token_start_index][0] <= start_char:\n",
    "                    token_start_index += 1\n",
    "                tokenized_examples[i][\"start_positions\"] = token_start_index - 1\n",
    "                while offsets[token_end_index][1] >= end_char:\n",
    "                    token_end_index -= 1\n",
    "                tokenized_examples[i][\"end_positions\"] = token_end_index + 1\n",
    "\n",
    "        return tokenized_examples\n",
    "\n",
    "    def prepare_validation_features(examples):\n",
    "        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results\n",
    "        # in one example possible giving several features when a context is long, each of those features having a\n",
    "        # context that overlaps a bit the context of the previous feature.\n",
    "        # NOTE: Almost the same functionality as HuggingFace's prepare_train_features function. The main difference is\n",
    "        # that HugggingFace uses ArrowTable as basic data structure, while we use list of dictionary instead.\n",
    "        contexts = [examples[i]['context'] for i in range(len(examples))]\n",
    "        questions = [examples[i]['question'] for i in range(len(examples))]\n",
    "\n",
    "        tokenized_examples = tokenizer(\n",
    "            questions,\n",
    "            contexts,\n",
    "            stride=args.doc_stride,\n",
    "            max_seq_len=args.max_seq_length)\n",
    "\n",
    "        # For validation, there is no need to compute start and end positions\n",
    "        for i, tokenized_example in enumerate(tokenized_examples):                    \n",
    "            # 引入 pinyin_ids\n",
    "            input_ids = tokenized_example[\"input_ids\"]\n",
    "            pinyin_ids = tokenizer.convert_tokens_to_pinyin_ids(tokenizer.convert_ids_to_tokens(input_ids))\n",
    "            tokenized_examples[i][\"pinyin_ids\"] = pinyin_ids\n",
    "\n",
    "            # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "            sequence_ids = tokenized_example['token_type_ids']\n",
    "\n",
    "            # One example can give several spans, this is the index of the example containing this span of text.\n",
    "            sample_index = tokenized_example['overflow_to_sample']\n",
    "            tokenized_examples[i][\"example_id\"] = examples[sample_index]['id']\n",
    "\n",
    "            # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token\n",
    "            # position is part of the context or not.\n",
    "            tokenized_examples[i][\"offset_mapping\"] = [\n",
    "                (o if sequence_ids[k] == 1 else None)\n",
    "                for k, o in enumerate(tokenized_example[\"offset_mapping\"])\n",
    "            ]\n",
    "\n",
    "        return tokenized_examples\n",
    "\n",
    "    if args.do_train:\n",
    "        train_ds.map(prepare_train_features, batched=True)\n",
    "        train_batch_sampler = paddle.io.DistributedBatchSampler(\n",
    "            train_ds, batch_size=args.batch_size, shuffle=True)\n",
    "        train_batchify_fn = lambda samples, fn=Dict({\n",
    "            \"input_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"pinyin_ids\":  Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"token_type_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_type_id),\n",
    "            \"start_positions\": Stack(dtype=\"int64\"),\n",
    "            \"end_positions\": Stack(dtype=\"int64\")\n",
    "        }): fn(samples)\n",
    "\n",
    "        train_data_loader = DataLoader(\n",
    "            dataset=train_ds,\n",
    "            batch_sampler=train_batch_sampler,\n",
    "            collate_fn=train_batchify_fn,\n",
    "            return_list=True)\n",
    "\n",
    "    if args.do_predict:\n",
    "        # dev_ds\n",
    "        dev_ds.map(prepare_validation_features, batched=True)\n",
    "        dev_batch_sampler = paddle.io.BatchSampler(\n",
    "            dev_ds, batch_size=args.batch_size, shuffle=False)\n",
    "\n",
    "        dev_batchify_fn = lambda samples, fn=Dict({\n",
    "            \"input_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"pinyin_ids\":  Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"token_type_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_type_id)\n",
    "        }): fn(samples)\n",
    "\n",
    "        dev_data_loader = DataLoader(\n",
    "            dataset=dev_ds,\n",
    "            batch_sampler=dev_batch_sampler,\n",
    "            collate_fn=dev_batchify_fn,\n",
    "            return_list=True)\n",
    "        \n",
    "        # test_ds\n",
    "        test_ds.map(prepare_validation_features, batched=True)\n",
    "        test_batch_sampler = paddle.io.BatchSampler(\n",
    "            test_ds, batch_size=args.batch_size, shuffle=False)\n",
    "\n",
    "        test_batchify_fn = lambda samples, fn=Dict({\n",
    "            \"input_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"pinyin_ids\":  Pad(axis=0, pad_val=tokenizer.pad_token_id),\n",
    "            \"token_type_ids\": Pad(axis=0, pad_val=tokenizer.pad_token_type_id)\n",
    "        }): fn(samples)\n",
    "\n",
    "        test_data_loader = DataLoader(\n",
    "            dataset=test_ds,\n",
    "            batch_sampler=test_batch_sampler,\n",
    "            collate_fn=test_batchify_fn,\n",
    "            return_list=True)\n",
    "\n",
    "\n",
    "        num_training_steps = args.max_steps if args.max_steps > 0 else len(\n",
    "            train_data_loader) * args.num_train_epochs\n",
    "        num_train_epochs = math.ceil(num_training_steps /\n",
    "                                     len(train_data_loader))\n",
    "\n",
    "        lr_scheduler = LinearDecayWithWarmup(\n",
    "            args.learning_rate, num_training_steps, args.warmup_proportion)\n",
    "\n",
    "        # Generate parameter names needed to perform weight decay.\n",
    "        # All bias and LayerNorm parameters are excluded.\n",
    "        decay_params = [\n",
    "            p.name for n, p in model.named_parameters()\n",
    "            if not any(nd in n for nd in [\"bias\", \"norm\"])\n",
    "        ]\n",
    "        optimizer = paddle.optimizer.AdamW(\n",
    "            learning_rate=lr_scheduler,\n",
    "            epsilon=args.adam_epsilon,\n",
    "            parameters=model.parameters(),\n",
    "            weight_decay=args.weight_decay,\n",
    "            apply_decay_param_fun=lambda x: x in decay_params)\n",
    "        criterion = CrossEntropyLossForSQuAD()\n",
    "\n",
    "        global_step = 0\n",
    "        tic_train = time.time()\n",
    "        for epoch in range(num_train_epochs):\n",
    "            for step, batch in enumerate(train_data_loader):\n",
    "                global_step += 1\n",
    "                input_ids, pinyin_ids, token_type_ids, start_positions, end_positions = batch\n",
    "                logits = model(\n",
    "                    input_ids=input_ids, pinyin_ids=pinyin_ids, token_type_ids=token_type_ids)\n",
    "                loss = criterion(logits, (start_positions, end_positions))\n",
    "\n",
    "                if global_step % args.logging_steps == 0:\n",
    "                    print(\n",
    "                        \"global step %d, epoch: %d, batch: %d, loss: %f, speed: %.2f step/s\"\n",
    "                        % (global_step, epoch + 1, step + 1, loss,\n",
    "                           args.logging_steps / (time.time() - tic_train)))\n",
    "                    tic_train = time.time()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                lr_scheduler.step()\n",
    "                optimizer.clear_grad()\n",
    "\n",
    "                if global_step % args.save_steps == 0 or global_step == num_training_steps:\n",
    "                    # if rank == 0:\n",
    "                    #     output_dir = os.path.join(args.output_dir,\n",
    "                    #                               \"model_%d\" % global_step)\n",
    "                    #     if not os.path.exists(output_dir):\n",
    "                    #         os.makedirs(output_dir)\n",
    "                    #     # need better way to get inner model of DataParallel\n",
    "                    #     model_to_save = model._layers if isinstance(\n",
    "                    #         model, paddle.DataParallel) else model\n",
    "                    #     model_to_save.save_pretrained(output_dir)\n",
    "                    #     tokenizer.save_pretrained(output_dir)\n",
    "                    #     print('Saving checkpoint to:', output_dir)\n",
    "                    if global_step == num_training_steps:\n",
    "                        break\n",
    "            if args.do_predict and rank == 0:\n",
    "                # dev_ds\n",
    "                evaluate(model, dev_data_loader, args, epoch+1, dataset_name='dev')\n",
    "                # test_ds\n",
    "                evaluate(model, test_data_loader, args, epoch+1, dataset_name='test')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec223966-6884-484d-9713-2cc36e40749b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "global step 10, epoch: 1, batch: 10, loss: 5.862209, speed: 1.48 step/s\n",
      "global step 20, epoch: 1, batch: 20, loss: 5.988250, speed: 1.72 step/s\n",
      "global step 30, epoch: 1, batch: 30, loss: 6.065759, speed: 1.72 step/s\n",
      "global step 40, epoch: 1, batch: 40, loss: 6.056306, speed: 1.72 step/s\n",
      "global step 50, epoch: 1, batch: 50, loss: 6.070762, speed: 1.72 step/s\n",
      "global step 60, epoch: 1, batch: 60, loss: 5.757662, speed: 1.72 step/s\n",
      "global step 70, epoch: 1, batch: 70, loss: 5.930552, speed: 1.72 step/s\n",
      "global step 80, epoch: 1, batch: 80, loss: 5.934573, speed: 1.72 step/s\n",
      "global step 90, epoch: 1, batch: 90, loss: 5.695853, speed: 1.72 step/s\n",
      "global step 100, epoch: 1, batch: 100, loss: 5.741165, speed: 1.75 step/s\n",
      "global step 110, epoch: 1, batch: 110, loss: 5.724845, speed: 1.73 step/s\n",
      "global step 120, epoch: 1, batch: 120, loss: 5.683299, speed: 1.73 step/s\n",
      "global step 130, epoch: 1, batch: 130, loss: 5.767026, speed: 1.72 step/s\n",
      "global step 140, epoch: 1, batch: 140, loss: 5.903282, speed: 1.74 step/s\n",
      "global step 150, epoch: 1, batch: 150, loss: 5.429982, speed: 1.72 step/s\n",
      "global step 160, epoch: 1, batch: 160, loss: 6.092547, speed: 1.72 step/s\n",
      "global step 170, epoch: 1, batch: 170, loss: 5.683709, speed: 1.72 step/s\n",
      "global step 180, epoch: 1, batch: 180, loss: 5.714303, speed: 1.72 step/s\n",
      "global step 190, epoch: 1, batch: 190, loss: 5.526630, speed: 1.74 step/s\n",
      "global step 200, epoch: 1, batch: 200, loss: 5.100922, speed: 1.72 step/s\n",
      "global step 210, epoch: 1, batch: 210, loss: 5.573108, speed: 1.72 step/s\n",
      "global step 220, epoch: 1, batch: 220, loss: 4.878587, speed: 1.72 step/s\n",
      "global step 230, epoch: 1, batch: 230, loss: 4.557806, speed: 1.74 step/s\n",
      "global step 240, epoch: 1, batch: 240, loss: 4.142955, speed: 1.73 step/s\n",
      "global step 250, epoch: 1, batch: 250, loss: 4.480210, speed: 1.72 step/s\n",
      "global step 260, epoch: 1, batch: 260, loss: 4.575274, speed: 1.72 step/s\n",
      "global step 270, epoch: 1, batch: 270, loss: 4.490275, speed: 1.71 step/s\n",
      "global step 280, epoch: 1, batch: 280, loss: 4.134420, speed: 1.73 step/s\n",
      "global step 290, epoch: 1, batch: 290, loss: 3.722122, speed: 1.72 step/s\n",
      "global step 300, epoch: 1, batch: 300, loss: 5.098068, speed: 1.72 step/s\n",
      "global step 310, epoch: 1, batch: 310, loss: 5.379973, speed: 1.74 step/s\n",
      "global step 320, epoch: 1, batch: 320, loss: 4.761113, speed: 1.72 step/s\n",
      "global step 330, epoch: 1, batch: 330, loss: 2.730665, speed: 1.73 step/s\n",
      "global step 340, epoch: 1, batch: 340, loss: 4.916422, speed: 1.72 step/s\n",
      "global step 350, epoch: 1, batch: 350, loss: 3.179594, speed: 1.72 step/s\n",
      "global step 360, epoch: 1, batch: 360, loss: 3.913397, speed: 1.72 step/s\n",
      "global step 370, epoch: 1, batch: 370, loss: 4.896363, speed: 1.73 step/s\n",
      "global step 380, epoch: 1, batch: 380, loss: 1.821494, speed: 1.72 step/s\n",
      "global step 390, epoch: 1, batch: 390, loss: 1.869850, speed: 1.72 step/s\n",
      "global step 400, epoch: 1, batch: 400, loss: 2.251956, speed: 1.72 step/s\n",
      "global step 410, epoch: 1, batch: 410, loss: 3.121856, speed: 1.72 step/s\n",
      "global step 420, epoch: 1, batch: 420, loss: 2.462332, speed: 1.72 step/s\n",
      "global step 430, epoch: 1, batch: 430, loss: 1.510234, speed: 1.72 step/s\n",
      "global step 440, epoch: 1, batch: 440, loss: 2.390119, speed: 1.75 step/s\n",
      "global step 450, epoch: 1, batch: 450, loss: 0.790545, speed: 1.72 step/s\n",
      "global step 460, epoch: 1, batch: 460, loss: 3.288955, speed: 1.73 step/s\n",
      "global step 470, epoch: 1, batch: 470, loss: 1.373342, speed: 1.71 step/s\n",
      "global step 480, epoch: 1, batch: 480, loss: 1.173578, speed: 1.72 step/s\n",
      "global step 490, epoch: 1, batch: 490, loss: 1.632465, speed: 1.72 step/s\n",
      "global step 500, epoch: 1, batch: 500, loss: 1.392969, speed: 1.72 step/s\n",
      "global step 510, epoch: 1, batch: 510, loss: 1.471222, speed: 1.72 step/s\n",
      "global step 520, epoch: 1, batch: 520, loss: 2.843056, speed: 1.72 step/s\n",
      "global step 530, epoch: 1, batch: 530, loss: 0.540803, speed: 1.72 step/s\n",
      "global step 540, epoch: 1, batch: 540, loss: 1.130098, speed: 1.72 step/s\n",
      "global step 550, epoch: 1, batch: 550, loss: 0.781961, speed: 1.72 step/s\n",
      "global step 560, epoch: 1, batch: 560, loss: 1.249907, speed: 1.74 step/s\n",
      "global step 570, epoch: 1, batch: 570, loss: 2.875434, speed: 1.74 step/s\n",
      "global step 580, epoch: 1, batch: 580, loss: 0.058912, speed: 1.72 step/s\n",
      "global step 590, epoch: 1, batch: 590, loss: 1.204504, speed: 1.72 step/s\n",
      "global step 600, epoch: 1, batch: 600, loss: 0.980775, speed: 1.72 step/s\n",
      "global step 610, epoch: 1, batch: 610, loss: 2.336946, speed: 1.72 step/s\n",
      "global step 620, epoch: 1, batch: 620, loss: 2.387858, speed: 1.72 step/s\n",
      "global step 630, epoch: 1, batch: 630, loss: 1.567626, speed: 1.72 step/s\n",
      "global step 640, epoch: 1, batch: 640, loss: 1.602778, speed: 1.72 step/s\n",
      "global step 650, epoch: 1, batch: 650, loss: 0.779575, speed: 1.72 step/s\n",
      "global step 660, epoch: 1, batch: 660, loss: 1.918970, speed: 1.72 step/s\n",
      "global step 670, epoch: 1, batch: 670, loss: 0.382029, speed: 1.71 step/s\n",
      "global step 680, epoch: 1, batch: 680, loss: 1.399765, speed: 1.72 step/s\n",
      "global step 690, epoch: 1, batch: 690, loss: 1.322743, speed: 1.72 step/s\n",
      "global step 700, epoch: 1, batch: 700, loss: 2.641803, speed: 1.72 step/s\n",
      "global step 710, epoch: 1, batch: 710, loss: 0.693250, speed: 1.74 step/s\n",
      "global step 720, epoch: 1, batch: 720, loss: 1.819213, speed: 1.72 step/s\n",
      "global step 730, epoch: 1, batch: 730, loss: 2.312699, speed: 1.74 step/s\n",
      "global step 740, epoch: 1, batch: 740, loss: 1.795611, speed: 1.72 step/s\n",
      "global step 750, epoch: 1, batch: 750, loss: 1.719385, speed: 1.73 step/s\n",
      "global step 760, epoch: 1, batch: 760, loss: 1.430498, speed: 1.73 step/s\n",
      "global step 770, epoch: 1, batch: 770, loss: 2.978318, speed: 1.72 step/s\n",
      "global step 780, epoch: 1, batch: 780, loss: 1.960914, speed: 1.72 step/s\n",
      "global step 790, epoch: 1, batch: 790, loss: 1.114838, speed: 1.72 step/s\n",
      "global step 800, epoch: 1, batch: 800, loss: 0.602157, speed: 1.72 step/s\n",
      "global step 810, epoch: 1, batch: 810, loss: 2.250195, speed: 1.72 step/s\n",
      "global step 820, epoch: 1, batch: 820, loss: 1.121379, speed: 1.72 step/s\n",
      "global step 830, epoch: 1, batch: 830, loss: 1.916674, speed: 1.72 step/s\n",
      "global step 840, epoch: 1, batch: 840, loss: 1.688372, speed: 1.72 step/s\n",
      "global step 850, epoch: 1, batch: 850, loss: 2.680009, speed: 1.72 step/s\n",
      "global step 860, epoch: 1, batch: 860, loss: 1.210648, speed: 1.73 step/s\n",
      "global step 870, epoch: 1, batch: 870, loss: 0.804460, speed: 1.72 step/s\n",
      "global step 880, epoch: 1, batch: 880, loss: 1.568301, speed: 1.72 step/s\n",
      "global step 890, epoch: 1, batch: 890, loss: 1.062510, speed: 1.75 step/s\n",
      "global step 900, epoch: 1, batch: 900, loss: 1.278451, speed: 1.71 step/s\n",
      "global step 910, epoch: 1, batch: 910, loss: 0.840074, speed: 1.72 step/s\n",
      "global step 920, epoch: 1, batch: 920, loss: 4.632985, speed: 1.72 step/s\n",
      "global step 930, epoch: 1, batch: 930, loss: 0.936441, speed: 1.72 step/s\n",
      "global step 940, epoch: 1, batch: 940, loss: 0.088207, speed: 1.72 step/s\n",
      "global step 950, epoch: 1, batch: 950, loss: 1.533095, speed: 1.73 step/s\n",
      "global step 960, epoch: 1, batch: 960, loss: 0.669909, speed: 1.72 step/s\n",
      "global step 970, epoch: 1, batch: 970, loss: 0.168159, speed: 1.72 step/s\n",
      "global step 980, epoch: 1, batch: 980, loss: 0.722996, speed: 1.72 step/s\n",
      "global step 990, epoch: 1, batch: 990, loss: 0.373739, speed: 1.72 step/s\n",
      "global step 1000, epoch: 1, batch: 1000, loss: 0.984697, speed: 1.72 step/s\n",
      "global step 1010, epoch: 1, batch: 1010, loss: 0.242844, speed: 1.72 step/s\n",
      "global step 1020, epoch: 1, batch: 1020, loss: 0.204567, speed: 1.72 step/s\n",
      "global step 1030, epoch: 1, batch: 1030, loss: 0.381579, speed: 1.72 step/s\n",
      "global step 1040, epoch: 1, batch: 1040, loss: 1.316459, speed: 1.72 step/s\n",
      "global step 1050, epoch: 1, batch: 1050, loss: 0.973426, speed: 1.72 step/s\n",
      "global step 1060, epoch: 1, batch: 1060, loss: 0.108155, speed: 1.73 step/s\n",
      "global step 1070, epoch: 1, batch: 1070, loss: 1.428588, speed: 1.72 step/s\n",
      "global step 1080, epoch: 1, batch: 1080, loss: 0.064006, speed: 1.72 step/s\n",
      "global step 1090, epoch: 1, batch: 1090, loss: 1.367014, speed: 1.72 step/s\n",
      "global step 1100, epoch: 1, batch: 1100, loss: 0.772743, speed: 1.70 step/s\n",
      "global step 1110, epoch: 1, batch: 1110, loss: 1.895742, speed: 1.71 step/s\n",
      "global step 1120, epoch: 1, batch: 1120, loss: 1.673788, speed: 1.72 step/s\n",
      "global step 1130, epoch: 1, batch: 1130, loss: 2.921703, speed: 1.72 step/s\n",
      "global step 1140, epoch: 1, batch: 1140, loss: 0.845236, speed: 1.72 step/s\n",
      "global step 1150, epoch: 1, batch: 1150, loss: 2.085048, speed: 1.72 step/s\n",
      "global step 1160, epoch: 1, batch: 1160, loss: 2.527676, speed: 1.72 step/s\n",
      "global step 1170, epoch: 1, batch: 1170, loss: 0.218560, speed: 1.72 step/s\n",
      "global step 1180, epoch: 1, batch: 1180, loss: 0.792902, speed: 1.72 step/s\n",
      "global step 1190, epoch: 1, batch: 1190, loss: 2.704049, speed: 1.74 step/s\n",
      "global step 1200, epoch: 1, batch: 1200, loss: 1.568636, speed: 1.72 step/s\n",
      "global step 1210, epoch: 1, batch: 1210, loss: 0.656125, speed: 1.74 step/s\n",
      "global step 1220, epoch: 1, batch: 1220, loss: 0.709616, speed: 1.72 step/s\n",
      "global step 1230, epoch: 1, batch: 1230, loss: 1.723013, speed: 1.71 step/s\n",
      "global step 1240, epoch: 1, batch: 1240, loss: 1.169493, speed: 1.72 step/s\n",
      "global step 1250, epoch: 1, batch: 1250, loss: 0.540070, speed: 1.73 step/s\n",
      "global step 1260, epoch: 1, batch: 1260, loss: 1.748898, speed: 1.73 step/s\n",
      "global step 1270, epoch: 1, batch: 1270, loss: 1.988709, speed: 1.72 step/s\n",
      "global step 1280, epoch: 1, batch: 1280, loss: 1.462082, speed: 1.72 step/s\n",
      "global step 1290, epoch: 1, batch: 1290, loss: 1.098573, speed: 1.72 step/s\n",
      "global step 1300, epoch: 1, batch: 1300, loss: 1.124013, speed: 1.70 step/s\n",
      "global step 1310, epoch: 1, batch: 1310, loss: 1.179720, speed: 1.72 step/s\n",
      "global step 1320, epoch: 1, batch: 1320, loss: 0.257377, speed: 1.72 step/s\n",
      "global step 1330, epoch: 1, batch: 1330, loss: 0.872777, speed: 1.72 step/s\n",
      "global step 1340, epoch: 1, batch: 1340, loss: 2.260149, speed: 1.72 step/s\n",
      "global step 1350, epoch: 1, batch: 1350, loss: 1.699663, speed: 1.72 step/s\n",
      "global step 1360, epoch: 1, batch: 1360, loss: 1.378320, speed: 1.72 step/s\n",
      "global step 1370, epoch: 1, batch: 1370, loss: 2.160755, speed: 1.72 step/s\n",
      "global step 1380, epoch: 1, batch: 1380, loss: 0.975598, speed: 1.74 step/s\n",
      "global step 1390, epoch: 1, batch: 1390, loss: 1.246903, speed: 1.72 step/s\n",
      "global step 1400, epoch: 1, batch: 1400, loss: 2.016550, speed: 1.72 step/s\n",
      "global step 1410, epoch: 1, batch: 1410, loss: 0.367073, speed: 1.72 step/s\n",
      "global step 1420, epoch: 1, batch: 1420, loss: 2.301661, speed: 1.72 step/s\n",
      "global step 1430, epoch: 1, batch: 1430, loss: 2.400827, speed: 1.72 step/s\n",
      "global step 1440, epoch: 1, batch: 1440, loss: 1.161201, speed: 1.72 step/s\n",
      "global step 1450, epoch: 1, batch: 1450, loss: 1.609808, speed: 1.74 step/s\n",
      "global step 1460, epoch: 1, batch: 1460, loss: 1.060877, speed: 1.72 step/s\n",
      "global step 1470, epoch: 1, batch: 1470, loss: 1.278297, speed: 1.72 step/s\n",
      "global step 1480, epoch: 1, batch: 1480, loss: 0.892756, speed: 1.74 step/s\n",
      "global step 1490, epoch: 1, batch: 1490, loss: 0.282515, speed: 1.74 step/s\n",
      "global step 1500, epoch: 1, batch: 1500, loss: 2.338285, speed: 1.72 step/s\n",
      "global step 1510, epoch: 1, batch: 1510, loss: 1.809365, speed: 1.72 step/s\n",
      "global step 1520, epoch: 1, batch: 1520, loss: 0.847434, speed: 1.72 step/s\n",
      "global step 1530, epoch: 1, batch: 1530, loss: 0.013191, speed: 1.72 step/s\n",
      "global step 1540, epoch: 1, batch: 1540, loss: 0.373365, speed: 1.72 step/s\n",
      "global step 1550, epoch: 1, batch: 1550, loss: 0.452228, speed: 1.72 step/s\n",
      "global step 1560, epoch: 1, batch: 1560, loss: 1.542564, speed: 1.72 step/s\n",
      "global step 1570, epoch: 1, batch: 1570, loss: 0.511348, speed: 1.72 step/s\n",
      "global step 1580, epoch: 1, batch: 1580, loss: 0.556282, speed: 1.72 step/s\n",
      "global step 1590, epoch: 1, batch: 1590, loss: 1.131314, speed: 1.72 step/s\n",
      "global step 1600, epoch: 1, batch: 1600, loss: 0.727029, speed: 1.72 step/s\n",
      "global step 1610, epoch: 1, batch: 1610, loss: 1.392577, speed: 1.73 step/s\n",
      "global step 1620, epoch: 1, batch: 1620, loss: 0.803226, speed: 1.72 step/s\n",
      "global step 1630, epoch: 1, batch: 1630, loss: 3.058621, speed: 1.72 step/s\n",
      "global step 1640, epoch: 1, batch: 1640, loss: 0.965555, speed: 1.73 step/s\n",
      "global step 1650, epoch: 1, batch: 1650, loss: 1.229284, speed: 1.72 step/s\n",
      "global step 1660, epoch: 1, batch: 1660, loss: 1.279098, speed: 1.72 step/s\n",
      "global step 1670, epoch: 1, batch: 1670, loss: 0.749234, speed: 1.72 step/s\n",
      "global step 1680, epoch: 1, batch: 1680, loss: 0.640299, speed: 1.72 step/s\n",
      "global step 1690, epoch: 1, batch: 1690, loss: 2.294492, speed: 1.72 step/s\n",
      "global step 1700, epoch: 1, batch: 1700, loss: 1.268882, speed: 1.72 step/s\n",
      "global step 1710, epoch: 1, batch: 1710, loss: 1.195714, speed: 1.73 step/s\n",
      "global step 1720, epoch: 1, batch: 1720, loss: 0.711590, speed: 1.72 step/s\n",
      "global step 1730, epoch: 1, batch: 1730, loss: 1.050802, speed: 1.72 step/s\n",
      "global step 1740, epoch: 1, batch: 1740, loss: 1.652377, speed: 1.72 step/s\n",
      "global step 1750, epoch: 1, batch: 1750, loss: 1.481072, speed: 1.72 step/s\n",
      "global step 1760, epoch: 1, batch: 1760, loss: 2.397473, speed: 1.72 step/s\n",
      "global step 1770, epoch: 1, batch: 1770, loss: 1.415377, speed: 1.74 step/s\n",
      "global step 1780, epoch: 1, batch: 1780, loss: 1.397924, speed: 1.73 step/s\n",
      "global step 1790, epoch: 1, batch: 1790, loss: 1.414526, speed: 1.73 step/s\n",
      "global step 1800, epoch: 1, batch: 1800, loss: 0.684959, speed: 1.72 step/s\n",
      "global step 1810, epoch: 1, batch: 1810, loss: 0.030000, speed: 1.72 step/s\n",
      "global step 1820, epoch: 1, batch: 1820, loss: 0.800208, speed: 1.72 step/s\n",
      "global step 1830, epoch: 1, batch: 1830, loss: 0.955548, speed: 1.73 step/s\n",
      "global step 1840, epoch: 1, batch: 1840, loss: 1.792126, speed: 1.73 step/s\n",
      "global step 1850, epoch: 1, batch: 1850, loss: 0.938019, speed: 1.75 step/s\n",
      "global step 1860, epoch: 1, batch: 1860, loss: 1.654095, speed: 1.72 step/s\n",
      "global step 1870, epoch: 1, batch: 1870, loss: 1.241892, speed: 1.72 step/s\n",
      "global step 1880, epoch: 1, batch: 1880, loss: 2.267979, speed: 1.73 step/s\n",
      "global step 1890, epoch: 1, batch: 1890, loss: 0.885959, speed: 1.72 step/s\n",
      "global step 1900, epoch: 1, batch: 1900, loss: 1.274768, speed: 1.73 step/s\n",
      "global step 1910, epoch: 1, batch: 1910, loss: 1.919451, speed: 1.72 step/s\n",
      "global step 1920, epoch: 1, batch: 1920, loss: 0.887836, speed: 1.72 step/s\n",
      "global step 1930, epoch: 1, batch: 1930, loss: 0.438459, speed: 1.72 step/s\n",
      "global step 1940, epoch: 1, batch: 1940, loss: 1.517845, speed: 1.72 step/s\n",
      "global step 1950, epoch: 1, batch: 1950, loss: 0.909352, speed: 1.72 step/s\n",
      "global step 1960, epoch: 1, batch: 1960, loss: 0.913879, speed: 1.72 step/s\n",
      "global step 1970, epoch: 1, batch: 1970, loss: 1.051879, speed: 1.72 step/s\n",
      "global step 1980, epoch: 1, batch: 1980, loss: 1.188096, speed: 1.72 step/s\n",
      "global step 1990, epoch: 1, batch: 1990, loss: 0.355597, speed: 1.72 step/s\n",
      "global step 2000, epoch: 1, batch: 2000, loss: 1.628312, speed: 1.72 step/s\n",
      "global step 2010, epoch: 1, batch: 2010, loss: 1.512648, speed: 1.75 step/s\n",
      "global step 2020, epoch: 1, batch: 2020, loss: 0.909932, speed: 1.73 step/s\n",
      "global step 2030, epoch: 1, batch: 2030, loss: 3.522049, speed: 1.73 step/s\n",
      "global step 2040, epoch: 1, batch: 2040, loss: 1.975844, speed: 1.72 step/s\n",
      "global step 2050, epoch: 1, batch: 2050, loss: 0.641623, speed: 1.72 step/s\n",
      "global step 2060, epoch: 1, batch: 2060, loss: 1.002260, speed: 1.73 step/s\n",
      "global step 2070, epoch: 1, batch: 2070, loss: 0.962765, speed: 1.74 step/s\n",
      "global step 2080, epoch: 1, batch: 2080, loss: 1.311398, speed: 1.72 step/s\n",
      "global step 2090, epoch: 1, batch: 2090, loss: 0.078177, speed: 1.72 step/s\n",
      "global step 2100, epoch: 1, batch: 2100, loss: 0.192193, speed: 1.72 step/s\n",
      "global step 2110, epoch: 1, batch: 2110, loss: 0.177712, speed: 1.72 step/s\n",
      "global step 2120, epoch: 1, batch: 2120, loss: 1.654738, speed: 1.72 step/s\n",
      "global step 2130, epoch: 1, batch: 2130, loss: 0.390278, speed: 1.73 step/s\n",
      "global step 2140, epoch: 1, batch: 2140, loss: 0.209590, speed: 1.72 step/s\n",
      "global step 2150, epoch: 1, batch: 2150, loss: 1.226096, speed: 1.73 step/s\n",
      "global step 2160, epoch: 1, batch: 2160, loss: 1.281403, speed: 1.72 step/s\n",
      "global step 2170, epoch: 1, batch: 2170, loss: 0.882523, speed: 1.69 step/s\n",
      "global step 2180, epoch: 1, batch: 2180, loss: 1.777527, speed: 1.72 step/s\n",
      "global step 2190, epoch: 1, batch: 2190, loss: 2.045994, speed: 1.72 step/s\n",
      "global step 2200, epoch: 1, batch: 2200, loss: 0.456940, speed: 1.72 step/s\n",
      "global step 2210, epoch: 1, batch: 2210, loss: 0.849245, speed: 1.73 step/s\n",
      "global step 2220, epoch: 1, batch: 2220, loss: 0.578574, speed: 1.72 step/s\n",
      "global step 2230, epoch: 1, batch: 2230, loss: 1.419772, speed: 1.72 step/s\n",
      "global step 2240, epoch: 1, batch: 2240, loss: 0.227779, speed: 1.73 step/s\n",
      "global step 2250, epoch: 1, batch: 2250, loss: 0.009989, speed: 1.72 step/s\n",
      "global step 2260, epoch: 1, batch: 2260, loss: 0.505905, speed: 1.76 step/s\n",
      "global step 2270, epoch: 1, batch: 2270, loss: 1.166158, speed: 1.72 step/s\n",
      "global step 2280, epoch: 1, batch: 2280, loss: 1.859774, speed: 1.72 step/s\n",
      "global step 2290, epoch: 1, batch: 2290, loss: 0.991390, speed: 1.72 step/s\n",
      "global step 2300, epoch: 1, batch: 2300, loss: 0.634597, speed: 1.72 step/s\n",
      "global step 2310, epoch: 1, batch: 2310, loss: 4.203133, speed: 1.72 step/s\n",
      "global step 2320, epoch: 1, batch: 2320, loss: 0.890760, speed: 1.72 step/s\n",
      "global step 2330, epoch: 1, batch: 2330, loss: 0.521168, speed: 1.72 step/s\n",
      "global step 2340, epoch: 1, batch: 2340, loss: 0.412193, speed: 1.72 step/s\n",
      "global step 2350, epoch: 1, batch: 2350, loss: 0.624411, speed: 1.74 step/s\n",
      "global step 2360, epoch: 1, batch: 2360, loss: 0.838720, speed: 1.72 step/s\n",
      "global step 2370, epoch: 1, batch: 2370, loss: 0.901794, speed: 1.72 step/s\n",
      "global step 2380, epoch: 1, batch: 2380, loss: 1.184186, speed: 1.72 step/s\n",
      "global step 2390, epoch: 1, batch: 2390, loss: 0.577278, speed: 1.72 step/s\n",
      "global step 2400, epoch: 1, batch: 2400, loss: 3.794066, speed: 1.72 step/s\n",
      "global step 2410, epoch: 1, batch: 2410, loss: 1.829729, speed: 1.72 step/s\n",
      "global step 2420, epoch: 1, batch: 2420, loss: 0.343415, speed: 1.72 step/s\n",
      "global step 2430, epoch: 1, batch: 2430, loss: 0.665191, speed: 1.72 step/s\n",
      "global step 2440, epoch: 1, batch: 2440, loss: 2.877231, speed: 1.73 step/s\n",
      "global step 2450, epoch: 1, batch: 2450, loss: 0.492252, speed: 1.72 step/s\n",
      "global step 2460, epoch: 1, batch: 2460, loss: 0.056425, speed: 1.72 step/s\n",
      "global step 2470, epoch: 1, batch: 2470, loss: 0.006542, speed: 1.72 step/s\n",
      "global step 2480, epoch: 1, batch: 2480, loss: 1.120047, speed: 1.69 step/s\n",
      "global step 2490, epoch: 1, batch: 2490, loss: 0.273900, speed: 1.72 step/s\n",
      "global step 2500, epoch: 1, batch: 2500, loss: 1.991983, speed: 1.74 step/s\n",
      "global step 2510, epoch: 1, batch: 2510, loss: 1.398640, speed: 1.72 step/s\n",
      "global step 2520, epoch: 1, batch: 2520, loss: 1.206679, speed: 1.72 step/s\n",
      "global step 2530, epoch: 1, batch: 2530, loss: 1.299921, speed: 1.72 step/s\n",
      "global step 2540, epoch: 1, batch: 2540, loss: 0.429937, speed: 1.74 step/s\n",
      "global step 2550, epoch: 1, batch: 2550, loss: 0.850352, speed: 1.72 step/s\n",
      "global step 2560, epoch: 1, batch: 2560, loss: 1.274604, speed: 1.72 step/s\n",
      "global step 2570, epoch: 1, batch: 2570, loss: 0.357030, speed: 1.73 step/s\n",
      "global step 2580, epoch: 1, batch: 2580, loss: 1.992625, speed: 1.72 step/s\n",
      "global step 2590, epoch: 1, batch: 2590, loss: 0.104905, speed: 1.73 step/s\n",
      "global step 2600, epoch: 1, batch: 2600, loss: 1.465252, speed: 1.72 step/s\n",
      "global step 2610, epoch: 1, batch: 2610, loss: 0.337938, speed: 1.65 step/s\n",
      "global step 2620, epoch: 1, batch: 2620, loss: 1.290479, speed: 1.72 step/s\n",
      "global step 2630, epoch: 1, batch: 2630, loss: 0.667659, speed: 1.72 step/s\n",
      "global step 2640, epoch: 1, batch: 2640, loss: 1.668470, speed: 1.72 step/s\n",
      "global step 2650, epoch: 1, batch: 2650, loss: 1.058464, speed: 1.72 step/s\n",
      "global step 2660, epoch: 1, batch: 2660, loss: 0.627809, speed: 1.72 step/s\n",
      "global step 2670, epoch: 1, batch: 2670, loss: 0.971619, speed: 1.72 step/s\n",
      "global step 2680, epoch: 1, batch: 2680, loss: 0.986101, speed: 1.73 step/s\n",
      "global step 2690, epoch: 1, batch: 2690, loss: 1.329489, speed: 1.72 step/s\n",
      "global step 2700, epoch: 1, batch: 2700, loss: 0.836564, speed: 1.72 step/s\n",
      "global step 2710, epoch: 1, batch: 2710, loss: 2.342359, speed: 1.72 step/s\n",
      "global step 2720, epoch: 1, batch: 2720, loss: 1.052953, speed: 1.72 step/s\n",
      "global step 2730, epoch: 1, batch: 2730, loss: 1.162145, speed: 1.75 step/s\n",
      "global step 2740, epoch: 1, batch: 2740, loss: 0.699144, speed: 1.74 step/s\n",
      "global step 2750, epoch: 1, batch: 2750, loss: 1.170300, speed: 1.74 step/s\n",
      "global step 2760, epoch: 1, batch: 2760, loss: 2.115952, speed: 1.72 step/s\n",
      "global step 2770, epoch: 1, batch: 2770, loss: 1.075383, speed: 1.74 step/s\n",
      "global step 2780, epoch: 1, batch: 2780, loss: 1.597697, speed: 1.72 step/s\n",
      "global step 2790, epoch: 1, batch: 2790, loss: 0.536114, speed: 1.72 step/s\n",
      "global step 2800, epoch: 1, batch: 2800, loss: 0.825203, speed: 1.72 step/s\n",
      "global step 2810, epoch: 1, batch: 2810, loss: 0.461950, speed: 1.72 step/s\n",
      "global step 2820, epoch: 1, batch: 2820, loss: 1.585045, speed: 1.72 step/s\n",
      "global step 2830, epoch: 1, batch: 2830, loss: 0.464989, speed: 1.72 step/s\n",
      "global step 2840, epoch: 1, batch: 2840, loss: 0.321604, speed: 1.72 step/s\n",
      "global step 2850, epoch: 1, batch: 2850, loss: 1.529008, speed: 1.71 step/s\n",
      "global step 2860, epoch: 1, batch: 2860, loss: 1.763955, speed: 1.72 step/s\n",
      "global step 2870, epoch: 1, batch: 2870, loss: 1.024062, speed: 1.73 step/s\n",
      "global step 2880, epoch: 1, batch: 2880, loss: 1.408616, speed: 1.72 step/s\n",
      "global step 2890, epoch: 1, batch: 2890, loss: 0.783753, speed: 1.72 step/s\n",
      "global step 2900, epoch: 1, batch: 2900, loss: 0.420857, speed: 1.72 step/s\n",
      "global step 2910, epoch: 1, batch: 2910, loss: 0.713272, speed: 1.72 step/s\n",
      "global step 2920, epoch: 1, batch: 2920, loss: 0.733993, speed: 1.72 step/s\n",
      "global step 2930, epoch: 1, batch: 2930, loss: 1.801777, speed: 1.72 step/s\n",
      "global step 2940, epoch: 1, batch: 2940, loss: 0.922909, speed: 1.73 step/s\n",
      "global step 2950, epoch: 1, batch: 2950, loss: 0.097187, speed: 1.73 step/s\n",
      "global step 2960, epoch: 1, batch: 2960, loss: 1.058692, speed: 1.73 step/s\n",
      "global step 2970, epoch: 1, batch: 2970, loss: 1.391064, speed: 1.73 step/s\n",
      "global step 2980, epoch: 1, batch: 2980, loss: 2.375133, speed: 1.72 step/s\n",
      "global step 2990, epoch: 1, batch: 2990, loss: 1.198325, speed: 1.72 step/s\n",
      "global step 3000, epoch: 1, batch: 3000, loss: 0.942199, speed: 1.72 step/s\n",
      "global step 3010, epoch: 1, batch: 3010, loss: 1.534578, speed: 1.72 step/s\n",
      "global step 3020, epoch: 1, batch: 3020, loss: 1.005218, speed: 1.72 step/s\n",
      "global step 3030, epoch: 1, batch: 3030, loss: 1.081325, speed: 1.72 step/s\n",
      "global step 3040, epoch: 1, batch: 3040, loss: 1.368807, speed: 1.72 step/s\n",
      "global step 3050, epoch: 1, batch: 3050, loss: 1.482757, speed: 1.69 step/s\n",
      "global step 3060, epoch: 1, batch: 3060, loss: 0.795191, speed: 1.69 step/s\n",
      "global step 3070, epoch: 1, batch: 3070, loss: 0.695311, speed: 1.72 step/s\n",
      "global step 3080, epoch: 1, batch: 3080, loss: 0.540442, speed: 1.72 step/s\n",
      "global step 3090, epoch: 1, batch: 3090, loss: 0.273761, speed: 1.74 step/s\n",
      "global step 3100, epoch: 1, batch: 3100, loss: 1.530443, speed: 1.71 step/s\n",
      "global step 3110, epoch: 1, batch: 3110, loss: 0.926209, speed: 1.72 step/s\n",
      "global step 3120, epoch: 1, batch: 3120, loss: 1.233856, speed: 1.72 step/s\n",
      "global step 3130, epoch: 1, batch: 3130, loss: 1.071731, speed: 1.73 step/s\n",
      "global step 3140, epoch: 1, batch: 3140, loss: 1.304787, speed: 1.72 step/s\n",
      "global step 3150, epoch: 1, batch: 3150, loss: 0.977641, speed: 1.73 step/s\n",
      "global step 3160, epoch: 1, batch: 3160, loss: 2.106482, speed: 1.72 step/s\n",
      "global step 3170, epoch: 1, batch: 3170, loss: 2.366071, speed: 1.73 step/s\n",
      "global step 3180, epoch: 1, batch: 3180, loss: 0.043187, speed: 1.72 step/s\n",
      "global step 3190, epoch: 1, batch: 3190, loss: 0.566984, speed: 1.72 step/s\n",
      "global step 3200, epoch: 1, batch: 3200, loss: 0.539253, speed: 1.73 step/s\n",
      "global step 3210, epoch: 1, batch: 3210, loss: 1.912734, speed: 1.72 step/s\n",
      "global step 3220, epoch: 1, batch: 3220, loss: 1.131119, speed: 1.72 step/s\n",
      "global step 3230, epoch: 1, batch: 3230, loss: 1.778761, speed: 1.72 step/s\n",
      "global step 3240, epoch: 1, batch: 3240, loss: 1.472757, speed: 1.75 step/s\n",
      "global step 3250, epoch: 1, batch: 3250, loss: 2.014839, speed: 1.69 step/s\n",
      "global step 3260, epoch: 1, batch: 3260, loss: 0.928782, speed: 1.72 step/s\n",
      "global step 3270, epoch: 1, batch: 3270, loss: 1.250610, speed: 1.72 step/s\n",
      "global step 3280, epoch: 1, batch: 3280, loss: 1.410649, speed: 1.72 step/s\n",
      "global step 3290, epoch: 1, batch: 3290, loss: 2.907534, speed: 1.72 step/s\n",
      "global step 3300, epoch: 1, batch: 3300, loss: 1.229034, speed: 1.72 step/s\n",
      "global step 3310, epoch: 1, batch: 3310, loss: 0.729436, speed: 1.71 step/s\n",
      "global step 3320, epoch: 1, batch: 3320, loss: 1.054258, speed: 1.72 step/s\n",
      "global step 3330, epoch: 1, batch: 3330, loss: 1.023691, speed: 1.72 step/s\n",
      "global step 3340, epoch: 1, batch: 3340, loss: 0.757898, speed: 1.72 step/s\n",
      "global step 3350, epoch: 1, batch: 3350, loss: 0.867092, speed: 1.72 step/s\n",
      "global step 3360, epoch: 1, batch: 3360, loss: 0.537523, speed: 1.72 step/s\n",
      "global step 3370, epoch: 1, batch: 3370, loss: 0.576178, speed: 1.74 step/s\n",
      "global step 3380, epoch: 1, batch: 3380, loss: 0.437656, speed: 1.72 step/s\n",
      "global step 3390, epoch: 1, batch: 3390, loss: 0.323663, speed: 1.72 step/s\n",
      "global step 3400, epoch: 1, batch: 3400, loss: 0.076632, speed: 1.72 step/s\n",
      "global step 3410, epoch: 1, batch: 3410, loss: 0.187218, speed: 1.72 step/s\n",
      "global step 3420, epoch: 1, batch: 3420, loss: 1.196574, speed: 1.72 step/s\n",
      "global step 3430, epoch: 1, batch: 3430, loss: 1.125070, speed: 1.73 step/s\n",
      "global step 3440, epoch: 1, batch: 3440, loss: 1.591521, speed: 1.65 step/s\n",
      "global step 3450, epoch: 1, batch: 3450, loss: 0.507517, speed: 1.72 step/s\n",
      "global step 3460, epoch: 1, batch: 3460, loss: 0.853512, speed: 1.72 step/s\n",
      "global step 3470, epoch: 1, batch: 3470, loss: 2.340531, speed: 1.72 step/s\n",
      "global step 3480, epoch: 1, batch: 3480, loss: 1.574912, speed: 1.72 step/s\n",
      "global step 3490, epoch: 1, batch: 3490, loss: 2.454221, speed: 1.72 step/s\n",
      "global step 3500, epoch: 1, batch: 3500, loss: 0.231459, speed: 1.72 step/s\n",
      "global step 3510, epoch: 1, batch: 3510, loss: 0.551065, speed: 1.72 step/s\n",
      "global step 3520, epoch: 1, batch: 3520, loss: 0.312794, speed: 1.71 step/s\n",
      "global step 3530, epoch: 1, batch: 3530, loss: 0.348645, speed: 1.74 step/s\n",
      "global step 3540, epoch: 1, batch: 3540, loss: 1.728759, speed: 1.73 step/s\n",
      "global step 3550, epoch: 1, batch: 3550, loss: 1.275519, speed: 1.72 step/s\n",
      "global step 3560, epoch: 1, batch: 3560, loss: 3.091007, speed: 1.72 step/s\n",
      "global step 3570, epoch: 1, batch: 3570, loss: 0.705055, speed: 1.72 step/s\n",
      "global step 3580, epoch: 1, batch: 3580, loss: 2.014466, speed: 1.72 step/s\n",
      "global step 3590, epoch: 1, batch: 3590, loss: 0.361126, speed: 1.72 step/s\n",
      "global step 3600, epoch: 1, batch: 3600, loss: 0.948080, speed: 1.73 step/s\n",
      "global step 3610, epoch: 1, batch: 3610, loss: 0.227062, speed: 1.72 step/s\n",
      "global step 3620, epoch: 1, batch: 3620, loss: 1.434412, speed: 1.72 step/s\n",
      "global step 3630, epoch: 1, batch: 3630, loss: 0.481256, speed: 1.72 step/s\n",
      "global step 3640, epoch: 1, batch: 3640, loss: 0.896167, speed: 1.72 step/s\n",
      "global step 3650, epoch: 1, batch: 3650, loss: 0.552973, speed: 1.68 step/s\n",
      "global step 3660, epoch: 1, batch: 3660, loss: 0.475757, speed: 1.72 step/s\n",
      "global step 3670, epoch: 1, batch: 3670, loss: 1.123086, speed: 1.72 step/s\n",
      "global step 3680, epoch: 1, batch: 3680, loss: 1.058467, speed: 1.73 step/s\n",
      "global step 3690, epoch: 1, batch: 3690, loss: 1.110027, speed: 1.72 step/s\n",
      "global step 3700, epoch: 1, batch: 3700, loss: 0.306349, speed: 1.72 step/s\n",
      "global step 3710, epoch: 1, batch: 3710, loss: 0.656565, speed: 1.72 step/s\n",
      "global step 3720, epoch: 1, batch: 3720, loss: 2.211752, speed: 1.72 step/s\n",
      "global step 3730, epoch: 1, batch: 3730, loss: 2.581747, speed: 1.71 step/s\n",
      "global step 3740, epoch: 1, batch: 3740, loss: 0.193741, speed: 1.74 step/s\n",
      "global step 3750, epoch: 1, batch: 3750, loss: 0.765175, speed: 1.74 step/s\n",
      "global step 3760, epoch: 1, batch: 3760, loss: 0.485724, speed: 1.72 step/s\n",
      "global step 3770, epoch: 1, batch: 3770, loss: 0.795506, speed: 1.73 step/s\n",
      "global step 3780, epoch: 1, batch: 3780, loss: 1.336757, speed: 1.72 step/s\n",
      "global step 3790, epoch: 1, batch: 3790, loss: 1.023834, speed: 1.72 step/s\n",
      "global step 3800, epoch: 1, batch: 3800, loss: 0.685625, speed: 1.72 step/s\n",
      "global step 3810, epoch: 1, batch: 3810, loss: 0.843983, speed: 1.72 step/s\n",
      "global step 3820, epoch: 1, batch: 3820, loss: 1.314924, speed: 1.72 step/s\n",
      "global step 3830, epoch: 1, batch: 3830, loss: 0.097517, speed: 1.72 step/s\n",
      "global step 3840, epoch: 1, batch: 3840, loss: 0.992321, speed: 1.72 step/s\n",
      "global step 3850, epoch: 1, batch: 3850, loss: 0.252778, speed: 1.72 step/s\n",
      "global step 3860, epoch: 1, batch: 3860, loss: 0.875240, speed: 1.72 step/s\n",
      "global step 3870, epoch: 1, batch: 3870, loss: 1.773454, speed: 1.69 step/s\n",
      "global step 3880, epoch: 1, batch: 3880, loss: 0.383844, speed: 1.73 step/s\n",
      "global step 3890, epoch: 1, batch: 3890, loss: 1.293478, speed: 1.73 step/s\n",
      "global step 3900, epoch: 1, batch: 3900, loss: 0.758856, speed: 1.72 step/s\n",
      "global step 3910, epoch: 1, batch: 3910, loss: 0.495044, speed: 1.72 step/s\n",
      "global step 3920, epoch: 1, batch: 3920, loss: 0.996757, speed: 1.72 step/s\n",
      "global step 3930, epoch: 1, batch: 3930, loss: 1.600040, speed: 1.71 step/s\n",
      "global step 3940, epoch: 1, batch: 3940, loss: 0.504818, speed: 1.74 step/s\n",
      "global step 3950, epoch: 1, batch: 3950, loss: 0.535672, speed: 1.72 step/s\n",
      "global step 3960, epoch: 1, batch: 3960, loss: 1.033683, speed: 1.72 step/s\n",
      "global step 3970, epoch: 1, batch: 3970, loss: 0.788233, speed: 1.73 step/s\n",
      "global step 3980, epoch: 1, batch: 3980, loss: 1.345911, speed: 1.72 step/s\n",
      "global step 3990, epoch: 1, batch: 3990, loss: 0.478652, speed: 1.72 step/s\n",
      "global step 4000, epoch: 1, batch: 4000, loss: 2.203691, speed: 1.72 step/s\n",
      "global step 4010, epoch: 1, batch: 4010, loss: 2.278192, speed: 1.73 step/s\n",
      "global step 4020, epoch: 1, batch: 4020, loss: 1.358571, speed: 1.73 step/s\n",
      "global step 4030, epoch: 1, batch: 4030, loss: 0.609052, speed: 1.72 step/s\n",
      "global step 4040, epoch: 1, batch: 4040, loss: 3.375820, speed: 1.72 step/s\n",
      "global step 4050, epoch: 1, batch: 4050, loss: 1.249080, speed: 1.72 step/s\n",
      "global step 4060, epoch: 1, batch: 4060, loss: 0.484693, speed: 1.70 step/s\n",
      "global step 4070, epoch: 1, batch: 4070, loss: 1.421540, speed: 1.71 step/s\n",
      "global step 4080, epoch: 1, batch: 4080, loss: 0.805228, speed: 1.73 step/s\n",
      "global step 4090, epoch: 1, batch: 4090, loss: 0.733736, speed: 1.76 step/s\n",
      "global step 4100, epoch: 1, batch: 4100, loss: 0.797075, speed: 1.72 step/s\n",
      "global step 4110, epoch: 1, batch: 4110, loss: 0.957244, speed: 1.73 step/s\n",
      "global step 4120, epoch: 1, batch: 4120, loss: 1.780301, speed: 1.72 step/s\n",
      "global step 4130, epoch: 1, batch: 4130, loss: 0.667152, speed: 1.72 step/s\n",
      "global step 4140, epoch: 1, batch: 4140, loss: 0.603766, speed: 1.71 step/s\n",
      "global step 4150, epoch: 1, batch: 4150, loss: 0.997567, speed: 1.72 step/s\n",
      "global step 4160, epoch: 1, batch: 4160, loss: 0.670215, speed: 1.72 step/s\n",
      "global step 4170, epoch: 1, batch: 4170, loss: 1.291621, speed: 1.72 step/s\n",
      "global step 4180, epoch: 1, batch: 4180, loss: 0.686680, speed: 1.73 step/s\n",
      "global step 4190, epoch: 1, batch: 4190, loss: 0.552944, speed: 1.72 step/s\n",
      "global step 4200, epoch: 1, batch: 4200, loss: 0.874456, speed: 1.72 step/s\n",
      "global step 4210, epoch: 1, batch: 4210, loss: 0.120371, speed: 1.72 step/s\n",
      "global step 4220, epoch: 1, batch: 4220, loss: 2.027211, speed: 1.72 step/s\n",
      "global step 4230, epoch: 1, batch: 4230, loss: 0.816092, speed: 1.72 step/s\n",
      "global step 4240, epoch: 1, batch: 4240, loss: 0.642388, speed: 1.74 step/s\n",
      "global step 4250, epoch: 1, batch: 4250, loss: 0.791541, speed: 1.72 step/s\n",
      "global step 4260, epoch: 1, batch: 4260, loss: 0.169543, speed: 1.70 step/s\n",
      "global step 4270, epoch: 1, batch: 4270, loss: 1.091181, speed: 1.72 step/s\n",
      "global step 4280, epoch: 1, batch: 4280, loss: 0.802800, speed: 1.72 step/s\n",
      "global step 4290, epoch: 1, batch: 4290, loss: 0.223272, speed: 1.72 step/s\n",
      "global step 4300, epoch: 1, batch: 4300, loss: 0.193769, speed: 1.72 step/s\n",
      "global step 4310, epoch: 1, batch: 4310, loss: 1.050884, speed: 1.72 step/s\n",
      "global step 4320, epoch: 1, batch: 4320, loss: 0.121795, speed: 1.72 step/s\n",
      "global step 4330, epoch: 1, batch: 4330, loss: 0.301659, speed: 1.72 step/s\n",
      "global step 4340, epoch: 1, batch: 4340, loss: 0.913653, speed: 1.72 step/s\n",
      "global step 4350, epoch: 1, batch: 4350, loss: 0.792180, speed: 1.71 step/s\n",
      "global step 4360, epoch: 1, batch: 4360, loss: 1.999292, speed: 1.72 step/s\n",
      "global step 4370, epoch: 1, batch: 4370, loss: 0.796821, speed: 1.72 step/s\n",
      "global step 4380, epoch: 1, batch: 4380, loss: 1.144489, speed: 1.72 step/s\n",
      "global step 4390, epoch: 1, batch: 4390, loss: 0.567343, speed: 1.72 step/s\n",
      "global step 4400, epoch: 1, batch: 4400, loss: 0.424459, speed: 1.72 step/s\n",
      "global step 4410, epoch: 1, batch: 4410, loss: 0.668780, speed: 1.73 step/s\n",
      "global step 4420, epoch: 1, batch: 4420, loss: 1.519226, speed: 1.72 step/s\n",
      "global step 4430, epoch: 1, batch: 4430, loss: 1.158512, speed: 1.72 step/s\n",
      "global step 4440, epoch: 1, batch: 4440, loss: 2.083773, speed: 1.72 step/s\n",
      "global step 4450, epoch: 1, batch: 4450, loss: 1.166136, speed: 1.71 step/s\n",
      "global step 4460, epoch: 1, batch: 4460, loss: 1.223029, speed: 1.71 step/s\n",
      "global step 4470, epoch: 1, batch: 4470, loss: 0.255541, speed: 1.72 step/s\n",
      "global step 4480, epoch: 1, batch: 4480, loss: 0.846860, speed: 1.72 step/s\n",
      "global step 4490, epoch: 1, batch: 4490, loss: 1.686024, speed: 1.72 step/s\n",
      "global step 4500, epoch: 1, batch: 4500, loss: 0.134985, speed: 1.72 step/s\n",
      "global step 4510, epoch: 1, batch: 4510, loss: 0.682375, speed: 1.72 step/s\n",
      "global step 4520, epoch: 1, batch: 4520, loss: 1.006267, speed: 1.72 step/s\n",
      "global step 4530, epoch: 1, batch: 4530, loss: 0.856327, speed: 1.73 step/s\n",
      "global step 4540, epoch: 1, batch: 4540, loss: 0.030287, speed: 1.74 step/s\n",
      "global step 4550, epoch: 1, batch: 4550, loss: 0.582583, speed: 1.72 step/s\n",
      "global step 4560, epoch: 1, batch: 4560, loss: 0.841819, speed: 1.71 step/s\n",
      "global step 4570, epoch: 1, batch: 4570, loss: 0.531618, speed: 1.72 step/s\n",
      "global step 4580, epoch: 1, batch: 4580, loss: 1.106761, speed: 1.72 step/s\n",
      "global step 4590, epoch: 1, batch: 4590, loss: 1.898680, speed: 1.72 step/s\n",
      "global step 4600, epoch: 1, batch: 4600, loss: 0.459157, speed: 1.72 step/s\n",
      "global step 4610, epoch: 1, batch: 4610, loss: 0.212909, speed: 1.72 step/s\n",
      "global step 4620, epoch: 1, batch: 4620, loss: 3.816391, speed: 1.72 step/s\n",
      "global step 4630, epoch: 1, batch: 4630, loss: 1.181067, speed: 1.72 step/s\n",
      "global step 4640, epoch: 1, batch: 4640, loss: 1.259422, speed: 1.73 step/s\n",
      "global step 4650, epoch: 1, batch: 4650, loss: 0.635612, speed: 1.73 step/s\n",
      "global step 4660, epoch: 1, batch: 4660, loss: 0.597528, speed: 1.72 step/s\n",
      "global step 4670, epoch: 1, batch: 4670, loss: 0.828895, speed: 1.75 step/s\n",
      "global step 4680, epoch: 1, batch: 4680, loss: 0.298996, speed: 1.70 step/s\n",
      "global step 4690, epoch: 1, batch: 4690, loss: 0.419783, speed: 1.72 step/s\n",
      "global step 4700, epoch: 1, batch: 4700, loss: 0.802472, speed: 1.72 step/s\n",
      "global step 4710, epoch: 1, batch: 4710, loss: 0.872005, speed: 1.75 step/s\n",
      "global step 4720, epoch: 1, batch: 4720, loss: 1.593771, speed: 1.72 step/s\n",
      "global step 4730, epoch: 1, batch: 4730, loss: 0.560774, speed: 1.73 step/s\n",
      "global step 4740, epoch: 1, batch: 4740, loss: 0.641503, speed: 1.75 step/s\n",
      "global step 4750, epoch: 1, batch: 4750, loss: 2.316868, speed: 1.72 step/s\n",
      "global step 4760, epoch: 1, batch: 4760, loss: 1.412192, speed: 1.72 step/s\n",
      "global step 4770, epoch: 1, batch: 4770, loss: 0.151460, speed: 1.70 step/s\n",
      "global step 4780, epoch: 1, batch: 4780, loss: 0.875939, speed: 1.72 step/s\n",
      "global step 4790, epoch: 1, batch: 4790, loss: 1.150535, speed: 1.72 step/s\n",
      "global step 4800, epoch: 1, batch: 4800, loss: 0.998106, speed: 1.72 step/s\n",
      "global step 4810, epoch: 1, batch: 4810, loss: 0.889561, speed: 1.72 step/s\n",
      "global step 4820, epoch: 1, batch: 4820, loss: 1.235574, speed: 1.72 step/s\n",
      "global step 4830, epoch: 1, batch: 4830, loss: 1.684402, speed: 1.73 step/s\n",
      "global step 4840, epoch: 1, batch: 4840, loss: 0.571494, speed: 1.72 step/s\n",
      "global step 4850, epoch: 1, batch: 4850, loss: 1.646423, speed: 1.72 step/s\n",
      "global step 4860, epoch: 1, batch: 4860, loss: 0.653645, speed: 1.72 step/s\n",
      "global step 4870, epoch: 1, batch: 4870, loss: 0.945711, speed: 1.72 step/s\n",
      "global step 4880, epoch: 1, batch: 4880, loss: 1.259786, speed: 1.72 step/s\n",
      "global step 4890, epoch: 1, batch: 4890, loss: 0.192439, speed: 1.68 step/s\n",
      "global step 4900, epoch: 1, batch: 4900, loss: 1.143809, speed: 1.66 step/s\n",
      "global step 4910, epoch: 1, batch: 4910, loss: 0.005814, speed: 1.72 step/s\n",
      "global step 4920, epoch: 1, batch: 4920, loss: 0.315462, speed: 1.73 step/s\n",
      "global step 4930, epoch: 1, batch: 4930, loss: 0.087004, speed: 1.72 step/s\n",
      "global step 4940, epoch: 1, batch: 4940, loss: 0.993625, speed: 1.72 step/s\n",
      "global step 4950, epoch: 1, batch: 4950, loss: 1.045401, speed: 1.73 step/s\n",
      "global step 4960, epoch: 1, batch: 4960, loss: 1.814744, speed: 1.72 step/s\n",
      "global step 4970, epoch: 1, batch: 4970, loss: 0.251808, speed: 1.71 step/s\n",
      "global step 4980, epoch: 1, batch: 4980, loss: 0.125141, speed: 1.72 step/s\n",
      "global step 4990, epoch: 1, batch: 4990, loss: 0.015593, speed: 1.73 step/s\n",
      "global step 5000, epoch: 1, batch: 5000, loss: 0.621671, speed: 1.72 step/s\n",
      "global step 5010, epoch: 1, batch: 5010, loss: 0.616060, speed: 1.72 step/s\n",
      "global step 5020, epoch: 1, batch: 5020, loss: 0.253167, speed: 1.72 step/s\n",
      "global step 5030, epoch: 1, batch: 5030, loss: 1.530154, speed: 1.75 step/s\n",
      "global step 5040, epoch: 1, batch: 5040, loss: 1.787961, speed: 1.72 step/s\n",
      "global step 5050, epoch: 1, batch: 5050, loss: 0.583159, speed: 1.72 step/s\n",
      "global step 5060, epoch: 1, batch: 5060, loss: 1.742193, speed: 1.72 step/s\n",
      "global step 5070, epoch: 1, batch: 5070, loss: 1.465203, speed: 1.73 step/s\n",
      "global step 5080, epoch: 1, batch: 5080, loss: 1.072043, speed: 1.72 step/s\n",
      "global step 5090, epoch: 1, batch: 5090, loss: 0.960962, speed: 1.73 step/s\n",
      "global step 5100, epoch: 1, batch: 5100, loss: 0.648536, speed: 1.72 step/s\n",
      "global step 5110, epoch: 1, batch: 5110, loss: 1.243581, speed: 1.72 step/s\n",
      "global step 5120, epoch: 1, batch: 5120, loss: 0.067307, speed: 1.72 step/s\n",
      "global step 5130, epoch: 1, batch: 5130, loss: 1.411190, speed: 1.72 step/s\n",
      "global step 5140, epoch: 1, batch: 5140, loss: 0.882007, speed: 1.72 step/s\n",
      "global step 5150, epoch: 1, batch: 5150, loss: 1.859406, speed: 1.72 step/s\n",
      "global step 5160, epoch: 1, batch: 5160, loss: 0.059077, speed: 1.72 step/s\n",
      "global step 5170, epoch: 1, batch: 5170, loss: 1.592047, speed: 1.72 step/s\n",
      "global step 5180, epoch: 1, batch: 5180, loss: 0.642611, speed: 1.72 step/s\n",
      "global step 5190, epoch: 1, batch: 5190, loss: 2.095972, speed: 1.72 step/s\n",
      "global step 5200, epoch: 1, batch: 5200, loss: 0.620648, speed: 1.74 step/s\n",
      "global step 5210, epoch: 1, batch: 5210, loss: 1.884748, speed: 1.76 step/s\n",
      "global step 5220, epoch: 1, batch: 5220, loss: 0.603033, speed: 1.73 step/s\n",
      "global step 5230, epoch: 1, batch: 5230, loss: 1.721477, speed: 1.72 step/s\n",
      "global step 5240, epoch: 1, batch: 5240, loss: 2.185949, speed: 1.72 step/s\n",
      "global step 5250, epoch: 1, batch: 5250, loss: 0.762883, speed: 1.72 step/s\n",
      "global step 5260, epoch: 1, batch: 5260, loss: 2.082008, speed: 1.72 step/s\n",
      "global step 5270, epoch: 1, batch: 5270, loss: 1.730993, speed: 1.72 step/s\n",
      "global step 5280, epoch: 1, batch: 5280, loss: 0.578763, speed: 1.72 step/s\n",
      "global step 5290, epoch: 1, batch: 5290, loss: 0.091789, speed: 1.72 step/s\n",
      "global step 5300, epoch: 1, batch: 5300, loss: 0.520061, speed: 1.72 step/s\n",
      "global step 5310, epoch: 1, batch: 5310, loss: 1.057713, speed: 1.72 step/s\n",
      "global step 5320, epoch: 1, batch: 5320, loss: 0.883028, speed: 1.72 step/s\n",
      "global step 5330, epoch: 1, batch: 5330, loss: 0.456703, speed: 1.72 step/s\n",
      "global step 5340, epoch: 1, batch: 5340, loss: 0.414846, speed: 1.74 step/s\n",
      "global step 5350, epoch: 1, batch: 5350, loss: 0.828094, speed: 1.72 step/s\n",
      "global step 5360, epoch: 1, batch: 5360, loss: 0.655888, speed: 1.72 step/s\n",
      "global step 5370, epoch: 1, batch: 5370, loss: 2.091935, speed: 1.72 step/s\n"
     ]
    }
   ],
   "source": [
    "run(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91235079-ae36-41f6-ad37-49bf2a74db49",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}