{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting Up Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import utils as u\n",
    "import logging\n",
    "import math\n",
    "import random\n",
    "import itertools\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from fairseq import (\n",
    "    checkpoint_utils,\n",
    "    distributed_utils,\n",
    "    options,\n",
    "    quantization_utils,\n",
    "    tasks,\n",
    "    utils,\n",
    ")\n",
    "from fairseq.logging import meters, metrics, progress_bar\n",
    "from fairseq.trainer import Trainer\n",
    "from fairseq.model_parallel.megatron_trainer import MegatronTrainer\n",
    "from fairseq.data import (\n",
    "    AppendTokenDataset,\n",
    "    ConcatDataset,\n",
    "    data_utils,\n",
    "    encoders,\n",
    "    indexed_dataset,\n",
    "    LanguagePairDataset,\n",
    "    PrependTokenDataset,\n",
    "    StripTokenDataset,\n",
    "    TruncateDataset,\n",
    "    iterators\n",
    ")\n",
    "from fairseq.models.fairseq_encoder import EncoderOut\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['mynotebook.py', '../../data/data-bin/iwslt16.dnmt.fr-en/', '--task', 'translation_han', '--max-sentences', '10', '--arch', 'han_transformer_iwslt_fr_en', '--distributed-world-size', '1', '--pretrained-transformer-checkpoint', '../../checkpoints/transformer_iwslt16/checkpoint_best.pt', '--freeze-transfo-params']\n",
    "\n",
    "\n",
    "logging.basicConfig(\n",
    "    format='%(asctime)s | %(levelname)s | %(name)s | %(message)s',\n",
    "    datefmt='%Y-%m-%d %H:%M:%S',\n",
    "    level=logging.INFO,\n",
    "    stream=sys.stdout,\n",
    ")\n",
    "logger = logging.getLogger('fairseq_cli.train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def distributed_main(i, args, start_rank=0):\n",
    "    args.device_id = i\n",
    "    if args.distributed_rank is None:  # torch.multiprocessing.spawn\n",
    "        args.distributed_rank = start_rank + i\n",
    "    main(args, init_distributed=True)\n",
    "\n",
    "def cli_main(modify_parser=None):\n",
    "    parser = options.get_training_parser()\n",
    "    args = options.parse_args_and_arch(parser, modify_parser=modify_parser)\n",
    "\n",
    "    if args.distributed_init_method is None:\n",
    "        distributed_utils.infer_init_method(args)\n",
    "\n",
    "    if args.distributed_init_method is not None:\n",
    "        # distributed training\n",
    "        if torch.cuda.device_count() > 1 and not args.distributed_no_spawn:\n",
    "            start_rank = args.distributed_rank\n",
    "            args.distributed_rank = None  # assign automatically\n",
    "            torch.multiprocessing.spawn(\n",
    "                fn=distributed_main,\n",
    "                args=(args, start_rank),\n",
    "                nprocs=torch.cuda.device_count(),\n",
    "            )\n",
    "        else:\n",
    "            distributed_main(args.device_id, args)\n",
    "    elif args.distributed_world_size > 1:\n",
    "        if not getattr(args, 'tpu', False):\n",
    "            # fallback for single node with multiple GPUs\n",
    "            assert args.distributed_world_size <= torch.cuda.device_count()\n",
    "            port = random.randint(10000, 20000)\n",
    "            args.distributed_init_method = 'tcp://localhost:{port}'.format(port=port)\n",
    "            args.distributed_rank = None  # set based on device id\n",
    "            torch.multiprocessing.spawn(\n",
    "                fn=distributed_main,\n",
    "                args=(args, ),\n",
    "                nprocs=args.distributed_world_size,\n",
    "            )\n",
    "        else:\n",
    "            import torch_xla.distributed.xla_multiprocessing as xmp\n",
    "            torch.multiprocessing.set_sharing_strategy('file_system')\n",
    "            xmp.spawn(\n",
    "                fn=distributed_main,\n",
    "                args=(args, ),\n",
    "                nprocs=8,  # use all 8 TPU cores\n",
    "            )\n",
    "    else:\n",
    "        # single GPU training\n",
    "        main_out = main(args)\n",
    "\n",
    "    return main_out\n",
    "\n",
    "def main(args, init_distributed=False):\n",
    "    utils.import_user_module(args)\n",
    "\n",
    "    assert args.max_tokens is not None or args.max_sentences is not None, \\\n",
    "        'Must specify batch size either with --max-tokens or --max-sentences'\n",
    "    metrics.reset()\n",
    "\n",
    "    # Initialize CUDA and distributed training\n",
    "    if torch.cuda.is_available() and not args.cpu and not getattr(args, 'tpu', False):\n",
    "        torch.cuda.set_device(args.device_id)\n",
    "    np.random.seed(args.seed)\n",
    "    utils.set_torch_seed(args.seed)\n",
    "    if init_distributed:\n",
    "        args.distributed_rank = distributed_utils.distributed_init(args)\n",
    "\n",
    "    if distributed_utils.is_master(args):\n",
    "        checkpoint_utils.verify_checkpoint_directory(args.save_dir)\n",
    "\n",
    "    # Print args\n",
    "    logger.info(args)\n",
    "\n",
    "    # Setup task, e.g., translation, language modeling, etc.\n",
    "    task = tasks.setup_task(args)\n",
    "\n",
    "    # Load valid dataset (we load training data below, based on the latest checkpoint)\n",
    "    for valid_sub_split in args.valid_subset.split(','):\n",
    "        task.load_dataset(valid_sub_split, combine=False, epoch=1)\n",
    "\n",
    "    # Build model and criterion\n",
    "    model = task.build_model(args)\n",
    "    criterion = task.build_criterion(args)\n",
    "    logger.info(model)\n",
    "    logger.info('model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))\n",
    "    logger.info('num. model params: {} (num. trained: {})'.format(\n",
    "        sum(p.numel() for p in model.parameters()),\n",
    "        sum(p.numel() for p in model.parameters() if p.requires_grad),\n",
    "    ))\n",
    "\n",
    "    # (optionally) Configure quantization\n",
    "    if args.quantization_config_path is not None:\n",
    "        quantizer = quantization_utils.Quantizer(\n",
    "            config_path=args.quantization_config_path,\n",
    "            max_epoch=args.max_epoch,\n",
    "            max_update=args.max_update,\n",
    "        )\n",
    "    else:\n",
    "        quantizer = None\n",
    "\n",
    "    # Build trainer\n",
    "    if args.model_parallel_size == 1:\n",
    "        trainer = Trainer(args, task, model, criterion, quantizer)\n",
    "    else:\n",
    "        trainer = MegatronTrainer(args, task, model, criterion)\n",
    "\n",
    "    logger.info('training on {} devices (GPUs/TPUs)'.format(args.distributed_world_size))\n",
    "    logger.info('max tokens per GPU = {} and max sentences per GPU = {}'.format(\n",
    "        args.max_tokens,\n",
    "        args.max_sentences,\n",
    "    ))\n",
    "\n",
    "    # Load the latest checkpoint if one is available and restore the\n",
    "    # corresponding train iterator\n",
    "    extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)\n",
    "    if args.tpu:\n",
    "        import torch_xla.core.xla_model as xm\n",
    "        xm.rendezvous('load_checkpoint')  # wait for all workers\n",
    "        xm.mark_step()\n",
    "\n",
    "    # Train until the learning rate gets too small\n",
    "    max_epoch = args.max_epoch or math.inf\n",
    "    max_update = args.max_update or math.inf\n",
    "    lr = trainer.get_lr()\n",
    "\n",
    "    return trainer, args\n",
    "\n",
    "def train(args, trainer, task, epoch_itr, max_update=math.inf):\n",
    "    \"\"\"Train the model for one epoch and return validation losses.\"\"\"\n",
    "    # Initialize data iterator\n",
    "    itr = epoch_itr.next_epoch_itr(\n",
    "        fix_batches_to_gpus=args.fix_batches_to_gpus,\n",
    "        shuffle=(epoch_itr.next_epoch_idx > args.curriculum),\n",
    "    )\n",
    "\n",
    "    trainer.begin_epoch(epoch_itr.epoch)\n",
    "\n",
    "    return\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-10-06 16:47:31 | INFO | fairseq_cli.train | Namespace(activation_dropout=0.0, activation_fn='relu', adaptive_input=False, adaptive_softmax_cutoff=None, adaptive_softmax_dropout=0, all_gather_list_size=16384, arch='han_transformer_iwslt_fr_en', attention_dropout=0.0, best_checkpoint_metric='loss', bf16=False, bpe=None, broadcast_buffers=False, bucket_cap_mb=25, checkpoint_suffix='', clip_norm=0.0, cpu=False, criterion='cross_entropy', cross_self_attention=False, curriculum=0, data='../../data/data-bin/iwslt16.dnmt.fr-en/', data_buffer_size=10, dataset_impl=None, ddp_backend='c10d', decoder_attention_heads=8, decoder_embed_dim=512, decoder_embed_path=None, decoder_ffn_embed_dim=1024, decoder_input_dim=512, decoder_layerdrop=0, decoder_layers=6, decoder_layers_to_keep=None, decoder_learned_pos=False, decoder_normalize_before=False, decoder_output_dim=512, device_id=0, disable_validation=False, distributed_backend='nccl', distributed_init_method=None, distributed_no_spawn=False, distributed_port=-1, distributed_rank=0, distributed_world_size=1, distributed_wrapper='DDP', dropout=0.2, empty_cache_freq=0, encoder_attention_heads=8, encoder_embed_dim=512, encoder_embed_path=None, encoder_ffn_embed_dim=1024, encoder_layerdrop=0, encoder_layers=6, encoder_layers_to_keep=None, encoder_learned_pos=False, encoder_normalize_before=False, eval_bleu=False, eval_bleu_args=None, eval_bleu_detok='space', eval_bleu_detok_args=None, eval_bleu_print_samples=False, eval_bleu_remove_bpe=None, eval_tokenized_bleu=False, fast_stat_sync=False, find_unused_parameters=False, fix_batches_to_gpus=False, fixed_validation_seed=None, force_anneal=None, fp16=False, fp16_init_scale=128, fp16_no_flatten_grads=False, fp16_scale_tolerance=0.0, fp16_scale_window=None, freeze_transfo_params=True, han_heads=8, keep_best_checkpoints=-1, keep_interval_updates=-1, keep_last_epochs=-1, layernorm_embedding=False, left_pad_source='True', left_pad_target='False', load_alignments=False, localsgd_frequency=3, log_format=None, log_interval=100, lr=[0.25], lr_scheduler='fixed', lr_shrink=0.1, max_epoch=0, max_sentences=10, max_sentences_valid=10, max_source_positions=1024, max_target_positions=1024, max_tokens=None, max_tokens_valid=None, max_update=0, maximize_best_checkpoint_metric=False, memory_efficient_bf16=False, memory_efficient_fp16=False, min_loss_scale=0.0001, min_lr=-1, model_parallel_size=1, momentum=0.99, n_context_sents=3, no_cross_attention=False, no_epoch_checkpoints=False, no_last_checkpoints=False, no_progress_bar=False, no_save=False, no_save_optimizer_state=False, no_scale_embedding=False, no_seed_provided=True, no_token_positional_embeddings=False, nprocs_per_node=2, num_batch_buckets=0, num_workers=1, optimizer='nag', optimizer_overrides='{}', patience=-1, pretrained_transformer_checkpoint='../../checkpoints/transformer_iwslt16/checkpoint_best.pt', profile=False, quant_noise_pq=0, quant_noise_pq_block_size=8, quant_noise_scalar=0, quantization_config_path=None, required_batch_size_multiple=8, reset_dataloader=False, reset_lr_scheduler=False, reset_meters=False, reset_optimizer=False, restore_file='checkpoint_last.pt', save_dir='checkpoints', save_interval=1, save_interval_updates=0, seed=1, sentence_avg=False, share_all_embeddings=True, share_decoder_input_output_embed=False, skip_invalid_size_inputs_valid_test=False, slowmo_algorithm='LocalSGD', slowmo_momentum=None, source_lang=None, stop_time_hours=0, target_lang=None, task='translation_han', tensorboard_logdir='', threshold_loss_scale=None, tie_adaptive_weights=False, tokenizer=None, tpu=False, train_subset='train', truncate_source=False, update_freq=[1], upsample_primary=1, use_bmuf=False, user_dir=None, valid_subset='valid', validate_interval=1, warmup_updates=0, weight_decay=0.0)\n",
      "2020-10-06 16:47:31 | INFO | fairseq.tasks.translation | [fr] dictionary: 31624 types\n",
      "2020-10-06 16:47:31 | INFO | fairseq.tasks.translation | [en] dictionary: 31624 types\n",
      "2020-10-06 16:47:31 | INFO | fairseq.data.data_utils | loaded 887 examples from: ../../data/data-bin/iwslt16.dnmt.fr-en/valid.fr-en.fr\n",
      "2020-10-06 16:47:31 | INFO | fairseq.data.data_utils | loaded 887 examples from: ../../data/data-bin/iwslt16.dnmt.fr-en/valid.fr-en.en\n",
      "2020-10-06 16:47:31 | INFO | fairseq.tasks.translation_han | ../../data/data-bin/iwslt16.dnmt.fr-en/ valid fr-en 887 examples\n",
      "2020-10-06 16:47:37 | INFO | fairseq.models.han_transformer | loaded pre-trained Transformer model from ../../checkpoints/transformer_iwslt16/checkpoint_best.pt\n",
      "2020-10-06 16:47:37 | INFO | fairseq.models.han_transformer | past context sentences modeled: 3 \n",
      "2020-10-06 16:47:38 | INFO | fairseq_cli.train | HanTransformerModel(\n",
      "  (encoder): HanEncoder(\n",
      "    (cache): HiddenStatesCache()\n",
      "    (transformer_encoder): TransformerEncoder(\n",
      "      (dropout_module): FairseqDropout()\n",
      "      (embed_tokens): Embedding(31624, 512, padding_idx=1)\n",
      "      (embed_positions): SinusoidalPositionalEmbedding()\n",
      "      (layers): ModuleList(\n",
      "        (0): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "        (1): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "        (2): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "        (3): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "        (4): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "        (5): TransformerEncoderLayer(\n",
      "          (self_attn): MultiheadAttention(\n",
      "            (dropout_module): FairseqDropout()\n",
      "            (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "            (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          )\n",
      "          (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (activation_dropout_module): FairseqDropout()\n",
      "          (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "          (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "          (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (layer_norm_query_word): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "    (layer_norm_query_sent): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "    (layer_norm_word_level): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "    (layer_norm_sentence_level): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "    (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "    (word_attn): MultiheadAttention(\n",
      "      (dropout_module): FairseqDropout()\n",
      "      (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "    )\n",
      "    (sent_attn): MultiheadAttention(\n",
      "      (dropout_module): FairseqDropout()\n",
      "      (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "      (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "    )\n",
      "    (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "    (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "    (linear): Linear(in_features=1024, out_features=512, bias=True)\n",
      "    (sigmoid): Sigmoid()\n",
      "  )\n",
      "  (decoder): TransformerDecoder(\n",
      "    (dropout_module): FairseqDropout()\n",
      "    (embed_tokens): Embedding(31624, 512, padding_idx=1)\n",
      "    (embed_positions): SinusoidalPositionalEmbedding()\n",
      "    (layers): ModuleList(\n",
      "      (0): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (1): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (2): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (3): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (4): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (5): TransformerDecoderLayer(\n",
      "        (dropout_module): FairseqDropout()\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (activation_dropout_module): FairseqDropout()\n",
      "        (self_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (encoder_attn): MultiheadAttention(\n",
      "          (dropout_module): FairseqDropout()\n",
      "          (k_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (v_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (q_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "          (out_proj): Linear(in_features=512, out_features=512, bias=True)\n",
      "        )\n",
      "        (encoder_attn_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "        (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "    )\n",
      "    (output_projection): Linear(in_features=512, out_features=31624, bias=False)\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-10-06 16:47:38 | INFO | fairseq_cli.train | model han_transformer_iwslt_fr_en, criterion CrossEntropyCriterion\n",
      "2020-10-06 16:47:38 | INFO | fairseq_cli.train | num. model params: 51416064 (num. trained: 3681280)\n",
      "2020-10-06 16:47:50 | INFO | fairseq.trainer | detected shared parameter: encoder.transformer_encoder.embed_tokens.weight <- decoder.embed_tokens.weight\n",
      "2020-10-06 16:47:50 | INFO | fairseq.trainer | detected shared parameter: encoder.transformer_encoder.embed_tokens.weight <- decoder.output_projection.weight\n",
      "2020-10-06 16:47:50 | INFO | fairseq.utils | ***********************CUDA enviroments for all 1 workers***********************\n",
      "2020-10-06 16:47:50 | INFO | fairseq.utils | rank   0: capabilities =  6.1  ; total memory = 23.877 GB ; name = Quadro P6000                            \n",
      "2020-10-06 16:47:50 | INFO | fairseq.utils | ***********************CUDA enviroments for all 1 workers***********************\n",
      "2020-10-06 16:47:50 | INFO | fairseq_cli.train | training on 1 devices (GPUs/TPUs)\n",
      "2020-10-06 16:47:50 | INFO | fairseq_cli.train | max tokens per GPU = None and max sentences per GPU = 10\n",
      "2020-10-06 16:47:50 | INFO | fairseq.trainer | no existing checkpoint found checkpoints/checkpoint_last.pt\n",
      "2020-10-06 16:47:50 | INFO | fairseq.trainer | loading train data for epoch 1\n",
      "2020-10-06 16:47:50 | INFO | fairseq.data.data_utils | loaded 215578 examples from: ../../data/data-bin/iwslt16.dnmt.fr-en/train.fr-en.fr\n",
      "2020-10-06 16:47:50 | INFO | fairseq.data.data_utils | loaded 215578 examples from: ../../data/data-bin/iwslt16.dnmt.fr-en/train.fr-en.en\n",
      "2020-10-06 16:47:50 | INFO | fairseq.tasks.translation_han | ../../data/data-bin/iwslt16.dnmt.fr-en/ train fr-en 215578 examples\n"
     ]
    }
   ],
   "source": [
    "trainer, args = cli_main()\n",
    "dataset = trainer.task.datasets['train']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# doc_heads_src = np.concatenate(([0],np.where(dataset.src.sizes == 1)[0]+1))\n",
    "# doc_heads_tgt = np.concatenate(([0],np.where(dataset.tgt.sizes == 1)[0]+1))\n",
    "# assert all(doc_heads_src == doc_heads_tgt)\n",
    "# decode = dataset.src_dict.string\n",
    "\n",
    "# print('DOCUMENT HEADS \\n--------------')\n",
    "# for h in doc_heads_src:\n",
    "#     src_head = dataset.src.__getitem__(h)\n",
    "#     tgt_head = dataset.tgt.__getitem__(h)\n",
    "#     print(decode(src_head))\n",
    "#     print(decode(tgt_head), '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-07-29 12:46:00 | INFO | fairseq.trainer | loading train data for epoch 1\n",
      "2020-07-29 12:46:00 | INFO | fairseq.data.data_utils | loaded 215578 examples from: ../data/data-bin/iwslt16.dnmt.fr_en/train.fr-en.fr\n",
      "2020-07-29 12:46:00 | INFO | fairseq.data.data_utils | loaded 215578 examples from: ../data/data-bin/iwslt16.dnmt.fr_en/train.fr-en.en\n",
      "2020-07-29 12:46:00 | INFO | fairseq.tasks.translation_han | ../data/data-bin/iwslt16.dnmt.fr_en/ train fr-en 215578 examples\n"
     ]
    }
   ],
   "source": [
    "epoch_itr = trainer.get_train_iterator(epoch=1)\n",
    "itr = epoch_itr.next_epoch_itr(shuffle=False)\n",
    "itr = iterators.GroupedIterator(itr, 1)\n",
    "cnt = 0\n",
    "for samples in itr:\n",
    "    break\n",
    "    # first_src_tokens = samples[0]['net_input']['src_tokens'][0]\n",
    "    # first_tgt_tokens = samples[0]['target'][0]\n",
    "    # print(decode(first_src_tokens), len(first_src_tokens), '({} tokens)'.format(len(first_src_tokens)))\n",
    "    # print(decode(first_tgt_tokens), '({} tokens)'.format(len(first_tgt_tokens)), end='\\n\\n')\n",
    "    # print('-------------------------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieve Encoded Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(\n",
    "#     {\n",
    "#     \"src_tokens\": src_tokens,\n",
    "#     \"src_lengths\": src_lengths,\n",
    "#     \"prev_output_tokens\": prev_output_tokens,\n",
    "#     \"doc_heads\": doc_heads,\n",
    "#     \"id\": id,\n",
    "#     \"sort_order\": sort_order,\n",
    "#     \"encoder_out\": encoder_out\n",
    "#     },\n",
    "#     \"$HOME/dev/fairseq/sh/mytensors.pt\"\n",
    "# )\n",
    "\n",
    "def sort_back(t, sort_order):\n",
    "    \"\"\"\n",
    "    Sort back to original order a Tensor t\n",
    "    that had been previously sorted according to sort_order.\n",
    "    \"\"\"\n",
    "    _, ord = sort_order.sort()\n",
    "    return t.index_select(0, ord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dict of saved tensors\n",
    "t = torch.load(\"mytensors.pt\", map_location=torch.device(device))\n",
    "id = t['id']\n",
    "src_tokens = t['src_tokens']\n",
    "doc_heads = t['doc_heads']\n",
    "sort_order = t['sort_order']\n",
    "encoder_out = t['encoder_out'] # T x B x C = 91 x 164 x 100\n",
    "\n",
    "K=3\n",
    "# e_out_mod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Update Cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'max(id)' = max(id)\n",
      "'cache.id' = cache.id\n",
      "'cache.h.shape' = cache.h.shape\n",
      "'cache.h_padding_mask.shape' = cache.h_padding_mask.shape\n"
     ]
    }
   ],
   "source": [
    "class HiddenStatesCache(nn.Module):\n",
    "\n",
    "    def __init__(self, cache_size):\n",
    "        super().__init__()\n",
    "        self.K = cache_size\n",
    "        self.register_buffer('id', torch.tensor([]))\n",
    "        self.register_buffer('h', torch.tensor([]))\n",
    "        self.register_buffer('h_padding_mask', torch.tensor([]))\n",
    "\n",
    "    def update_cache(self, id, h, h_padding_mask, sort_order, doc_heads):\n",
    "        # TODO(lo) manage updating for end of corpus since\n",
    "        # the line past the last one is not listed as head:\n",
    "        # there might be problems passing from training corpus to valid corpus?\n",
    "\n",
    "        # sort id in the original order; retrieve K last ones\n",
    "        cid = utils.sort_back(id, sort_order)[-self.K:]\n",
    "        # check whether some ids correspond to document heads\n",
    "        where_heads = torch.where(cid[:, None] == doc_heads)[0]\n",
    "\n",
    "        if where_heads.shape[0] > 0:\n",
    "            # only save to cache from the last head on\n",
    "            cid = cid[where_heads[-1]:]\n",
    "        \n",
    "        if cid[:, None] in (doc_heads - 1):\n",
    "            # empty cache if the first sentence of the next batch is an head\n",
    "            # note: this will not work for the last batch #TODO(lo)\n",
    "            cid = id.new_tensor([])\n",
    "\n",
    "        # update cached id\n",
    "        self.id = cid.detach()\n",
    "        # update cached hidden states\n",
    "        self.h = h[:, torch.where(self.id[:, None] == id)[1], :].detach()\n",
    "        self.h_padding_mask = h_padding_mask[torch.where(\n",
    "            self.id[:, None] == id)[1], :].detach()\n",
    "\n",
    "\n",
    "# test class\n",
    "doc_heads = torch.FloatTensor([810, 815]).to(device)\n",
    "cache = HiddenStatesCache(cache_size=K)\n",
    "cache.update_cache(\n",
    "    id=id,\n",
    "    sort_order=sort_order,\n",
    "    h=encoder_out.encoder_out,\n",
    "    h_padding_mask=encoder_out.encoder_padding_mask,\n",
    "    doc_heads=doc_heads)\n",
    "\n",
    "u.debug('max(id)')\n",
    "u.debug('cache.id')\n",
    "u.debug('cache.h.shape')\n",
    "u.debug('cache.h_padding_mask.shape')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Retrieve Context Id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
      "         True, True, True, True, True, True, True]], device='cuda:0')\n",
      "tensor([815, 810], device='cuda:0')\n",
      "tensor([815, 810, 811, 816], device='cuda:0')\n",
      "tensor([815, 810, 811, 816, 817, 812], device='cuda:0')\n",
      "'id.shape' = id.shape\n",
      "'expanded_id.shape' = expanded_id.shape\n",
      "'expanded_encoder_out.shape' = expanded_encoder_out.shape\n",
      "'expanded_encoder_padding_mask.shape' = expanded_encoder_padding_mask.shape\n",
      "'where_context[0].shape' = where_context[0].shape\n"
     ]
    }
   ],
   "source": [
    "class TestContext(object):\n",
    "    def __init__(self, cache, K):\n",
    "        self.cache = cache\n",
    "        self.K = K\n",
    "\n",
    "    def expand_id(self, id):\n",
    "        # expand by prepending cache\n",
    "        expanded_id = torch.cat((self.cache.id, id), dim=0)\n",
    "        # expand by prepending dummy index\n",
    "        dummy_id = id.new_tensor([-1])\n",
    "        expanded_id = torch.cat((dummy_id, expanded_id), dim=0)\n",
    "        return expanded_id\n",
    "\n",
    "    def expand_h(self, h):\n",
    "        # expand by prepending cache\n",
    "        self.cache.h, h, _ = utils.pad_smaller(self.cache.h, h, dim=0, value=0)\n",
    "        expanded_h = torch.cat((self.cache.h, h), dim=1)\n",
    "        # expand by prepending dummy context\n",
    "        dummy_context = expanded_h.new_empty((expanded_h.shape[0], 1, expanded_h.shape[2]))\n",
    "        expanded_h = torch.cat((dummy_context, expanded_h), dim=1)\n",
    "        return expanded_h\n",
    "    \n",
    "    def expand_h_padding_mask(self, h_padding_mask):\n",
    "        # expand by prepending cache\n",
    "        self.cache.h_padding_mask, h_padding_mask, _ = utils.pad_smaller(\n",
    "            self.cache.h_padding_mask, h_padding_mask, dim=1, value=1\n",
    "        )\n",
    "        expanded_h_padding_mask = torch.cat(\n",
    "            (self.cache.h_padding_mask, h_padding_mask),\n",
    "            dim=0\n",
    "        )\n",
    "        # expand by prepending dummy context padding\n",
    "        dummy_padding = expanded_h_padding_mask.new_ones((1, expanded_h_padding_mask.shape[-1]))\n",
    "        print(dummy_padding)\n",
    "        expanded_h_padding_mask = torch.cat((dummy_padding, expanded_h_padding_mask), dim=0)\n",
    "        return expanded_h_padding_mask\n",
    "\n",
    "    def get_context_position(self, id, expanded_id, doc_heads):\n",
    "        # retrieve context positions\n",
    "        where_no_context = id.new_tensor([])\n",
    "        where_context = []\n",
    "        for k in range(1, self.K+1):\n",
    "            # select ids of sentences that do not have kth context\n",
    "            no_context_curr = torch.where(id[:,None]==doc_heads+(k-1))[0]\n",
    "            where_no_context = torch.cat((where_no_context, no_context_curr))\n",
    "            print(id[where_no_context])\n",
    "            # id of kth context sentences\n",
    "            idk = (id - k)\n",
    "            idk[where_no_context] = -1\n",
    "            # retrieve position of the kth context sentence\n",
    "            where_context.append(torch.where(idk[:,None]==expanded_id)[1])\n",
    "        return where_context\n",
    "\n",
    "# test\n",
    "tc = TestContext(cache, K)\n",
    "# B -> K + B\n",
    "expanded_id = tc.expand_id(id)\n",
    "# T x B x C -> T x (K + B) x C\n",
    "expanded_encoder_out = tc.expand_h(encoder_out.encoder_out)\n",
    "# B x T -> (B + K) x T\n",
    "expanded_encoder_padding_mask = tc.expand_h_padding_mask(\n",
    "    encoder_out.encoder_padding_mask)\n",
    "# get positions\n",
    "where_context = tc.get_context_position(\n",
    "    id, expanded_id, doc_heads)\n",
    "\n",
    "u.debug('id.shape')\n",
    "u.debug('expanded_id.shape')\n",
    "u.debug('expanded_encoder_out.shape')\n",
    "u.debug('expanded_encoder_padding_mask.shape')\n",
    "u.debug('where_context[0].shape') # this will not work properly because in the cache we have the same ids as in the current batch. See test of the cache above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add masking of dummy context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., -inf, 0., -inf, 0., 0., 0., 0., -inf, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., -inf, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn_mask = torch.zeros(where_context[2].shape).to(device)\n",
    "attn_mask.masked_fill((where_context[2] == 0), float('-inf'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "wc = torch.stack((where_context[2].type_as(cache.h), torch.zeros(where_context[2].shape, device=where_context[2].device).masked_fill((where_context[2] == 0), float('-inf'))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attn_mask.shape = torch.Size([316, 5, 6])\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (158) must match the size of tensor b (316) at non-singleton dimension 0",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-14-e51cddb4e904>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdebug\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mattn_mask\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mwc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mwhere_context\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mattn_mask\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmasked_fill\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'-inf'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m: The size of tensor a (158) must match the size of tensor b (316) at non-singleton dimension 0"
     ]
    }
   ],
   "source": [
    "attn_mask = torch.zeros(where_context[2].shape[0]*2, 5, 6).to(device)\n",
    "u.debug(attn_mask.shape)\n",
    "wc = (where_context[2] == 0).view(-1, 1, 1)\n",
    "attn_mask.masked_fill(wc, float('-inf'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-7.4466e+09,  4.5853e-41,  7.5929e-06,  3.0642e-41,         nan,\n",
       "           0.0000e+00],\n",
       "         [ 4.4721e+21,  2.8799e+32,  1.7438e+02,  4.1641e+12,  7.5338e+28,\n",
       "           7.0254e+28],\n",
       "         [ 5.2478e-08,  1.2088e+33,  6.0676e-43,  0.0000e+00,  7.5815e-06,\n",
       "           3.0642e-41],\n",
       "         [ 7.5931e-06,  3.0642e-41,         nan,  4.5081e+16,  1.8314e+25,\n",
       "           6.9768e+22],\n",
       "         [ 8.5305e+02,  3.0331e+32,  3.9415e+12,  1.0333e+18,  8.9683e-44,\n",
       "           0.0000e+00]],\n",
       "\n",
       "        [[       -inf,        -inf,        -inf,        -inf,        -inf,\n",
       "                 -inf],\n",
       "         [       -inf,        -inf,        -inf,        -inf,        -inf,\n",
       "                 -inf],\n",
       "         [       -inf,        -inf,        -inf,        -inf,        -inf,\n",
       "                 -inf],\n",
       "         [       -inf,        -inf,        -inf,        -inf,        -inf,\n",
       "                 -inf],\n",
       "         [       -inf,        -inf,        -inf,        -inf,        -inf,\n",
       "                 -inf]],\n",
       "\n",
       "        [[ 8.9683e-44,  0.0000e+00,  8.9683e-44,  0.0000e+00,  7.5885e-06,\n",
       "           3.0642e-41],\n",
       "         [ 3.0829e-44,  0.0000e+00,         nan,  1.4804e+37,  4.4721e+21,\n",
       "           2.8799e+32],\n",
       "         [ 1.1160e+04,  7.1443e+31,  1.8888e+31,  3.9752e+28,  5.2478e-08,\n",
       "           1.2088e+33],\n",
       "         [ 2.4803e-43,  0.0000e+00,  7.5815e-06,  3.0642e-41, -7.4464e+09,\n",
       "           4.5853e-41],\n",
       "         [        nan,  4.5082e+16,  1.8314e+25,  6.9768e+22,  8.5305e+02,\n",
       "           7.5516e+31]]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "p = torch.empty(size=(3,5,6))\n",
    "t = torch.tensor([False, True, False]).view(-1,1,1)\n",
    "p.masked_fill(mask=t, value=float('-inf'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 1],\n",
       "        [1, 1],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "t = torch.tensor([[1, 1], [0, 0], [0, 0]])\n",
    "t.repeat_interleave(repeats=2, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.2574, 0.7313])\n",
      "tensor([0.7786, 0.6751])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor(1.4537)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = nn.Sigmoid()\n",
    "input = torch.randn(2)\n",
    "output = m(input)\n",
    "print(input)\n",
    "print(output)\n",
    "sum(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload pre-trained sentence-level checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-10-06 17:00:16 | INFO | fairseq.models.han_transformer | loaded pre-trained Transformer model from ../../checkpoints/transformer_iwslt16/checkpoint_best.pt\n",
      "2020-10-06 17:00:16 | INFO | fairseq.models.han_transformer | past context sentences modeled: 3 \n"
     ]
    }
   ],
   "source": [
    "model = trainer.task.build_model(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoder.transformer_encoder.version\n",
      "encoder.transformer_encoder.embed_tokens.weight\n",
      "encoder.transformer_encoder.embed_positions._float_tensor\n",
      "encoder.transformer_encoder.layers.0.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.0.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.0.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.0.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.0.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.0.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.0.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.0.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.0.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.0.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.0.fc1.weight\n",
      "encoder.transformer_encoder.layers.0.fc1.bias\n",
      "encoder.transformer_encoder.layers.0.fc2.weight\n",
      "encoder.transformer_encoder.layers.0.fc2.bias\n",
      "encoder.transformer_encoder.layers.0.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.0.final_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.1.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.1.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.1.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.1.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.1.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.1.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.1.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.1.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.1.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.1.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.1.fc1.weight\n",
      "encoder.transformer_encoder.layers.1.fc1.bias\n",
      "encoder.transformer_encoder.layers.1.fc2.weight\n",
      "encoder.transformer_encoder.layers.1.fc2.bias\n",
      "encoder.transformer_encoder.layers.1.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.1.final_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.2.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.2.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.2.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.2.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.2.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.2.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.2.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.2.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.2.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.2.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.2.fc1.weight\n",
      "encoder.transformer_encoder.layers.2.fc1.bias\n",
      "encoder.transformer_encoder.layers.2.fc2.weight\n",
      "encoder.transformer_encoder.layers.2.fc2.bias\n",
      "encoder.transformer_encoder.layers.2.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.2.final_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.3.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.3.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.3.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.3.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.3.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.3.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.3.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.3.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.3.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.3.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.3.fc1.weight\n",
      "encoder.transformer_encoder.layers.3.fc1.bias\n",
      "encoder.transformer_encoder.layers.3.fc2.weight\n",
      "encoder.transformer_encoder.layers.3.fc2.bias\n",
      "encoder.transformer_encoder.layers.3.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.3.final_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.4.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.4.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.4.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.4.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.4.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.4.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.4.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.4.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.4.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.4.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.4.fc1.weight\n",
      "encoder.transformer_encoder.layers.4.fc1.bias\n",
      "encoder.transformer_encoder.layers.4.fc2.weight\n",
      "encoder.transformer_encoder.layers.4.fc2.bias\n",
      "encoder.transformer_encoder.layers.4.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.4.final_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.5.self_attn.k_proj.weight\n",
      "encoder.transformer_encoder.layers.5.self_attn.k_proj.bias\n",
      "encoder.transformer_encoder.layers.5.self_attn.v_proj.weight\n",
      "encoder.transformer_encoder.layers.5.self_attn.v_proj.bias\n",
      "encoder.transformer_encoder.layers.5.self_attn.q_proj.weight\n",
      "encoder.transformer_encoder.layers.5.self_attn.q_proj.bias\n",
      "encoder.transformer_encoder.layers.5.self_attn.out_proj.weight\n",
      "encoder.transformer_encoder.layers.5.self_attn.out_proj.bias\n",
      "encoder.transformer_encoder.layers.5.self_attn_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.5.self_attn_layer_norm.bias\n",
      "encoder.transformer_encoder.layers.5.fc1.weight\n",
      "encoder.transformer_encoder.layers.5.fc1.bias\n",
      "encoder.transformer_encoder.layers.5.fc2.weight\n",
      "encoder.transformer_encoder.layers.5.fc2.bias\n",
      "encoder.transformer_encoder.layers.5.final_layer_norm.weight\n",
      "encoder.transformer_encoder.layers.5.final_layer_norm.bias\n",
      "encoder.layer_norm_query_word.weight\n",
      "encoder.layer_norm_query_word.bias\n",
      "encoder.layer_norm_query_sent.weight\n",
      "encoder.layer_norm_query_sent.bias\n",
      "encoder.layer_norm_word_level.weight\n",
      "encoder.layer_norm_word_level.bias\n",
      "encoder.layer_norm_sentence_level.weight\n",
      "encoder.layer_norm_sentence_level.bias\n",
      "encoder.final_layer_norm.weight\n",
      "encoder.final_layer_norm.bias\n",
      "encoder.word_attn.k_proj.weight\n",
      "encoder.word_attn.k_proj.bias\n",
      "encoder.word_attn.v_proj.weight\n",
      "encoder.word_attn.v_proj.bias\n",
      "encoder.word_attn.q_proj.weight\n",
      "encoder.word_attn.q_proj.bias\n",
      "encoder.word_attn.out_proj.weight\n",
      "encoder.word_attn.out_proj.bias\n",
      "encoder.sent_attn.k_proj.weight\n",
      "encoder.sent_attn.k_proj.bias\n",
      "encoder.sent_attn.v_proj.weight\n",
      "encoder.sent_attn.v_proj.bias\n",
      "encoder.sent_attn.q_proj.weight\n",
      "encoder.sent_attn.q_proj.bias\n",
      "encoder.sent_attn.out_proj.weight\n",
      "encoder.sent_attn.out_proj.bias\n",
      "encoder.fc1.weight\n",
      "encoder.fc1.bias\n",
      "encoder.fc2.weight\n",
      "encoder.fc2.bias\n",
      "encoder.linear.weight\n",
      "encoder.linear.bias\n",
      "decoder.version\n",
      "decoder.embed_tokens.weight\n",
      "decoder.embed_positions._float_tensor\n",
      "decoder.layers.0.self_attn.k_proj.weight\n",
      "decoder.layers.0.self_attn.k_proj.bias\n",
      "decoder.layers.0.self_attn.v_proj.weight\n",
      "decoder.layers.0.self_attn.v_proj.bias\n",
      "decoder.layers.0.self_attn.q_proj.weight\n",
      "decoder.layers.0.self_attn.q_proj.bias\n",
      "decoder.layers.0.self_attn.out_proj.weight\n",
      "decoder.layers.0.self_attn.out_proj.bias\n",
      "decoder.layers.0.self_attn_layer_norm.weight\n",
      "decoder.layers.0.self_attn_layer_norm.bias\n",
      "decoder.layers.0.encoder_attn.k_proj.weight\n",
      "decoder.layers.0.encoder_attn.k_proj.bias\n",
      "decoder.layers.0.encoder_attn.v_proj.weight\n",
      "decoder.layers.0.encoder_attn.v_proj.bias\n",
      "decoder.layers.0.encoder_attn.q_proj.weight\n",
      "decoder.layers.0.encoder_attn.q_proj.bias\n",
      "decoder.layers.0.encoder_attn.out_proj.weight\n",
      "decoder.layers.0.encoder_attn.out_proj.bias\n",
      "decoder.layers.0.encoder_attn_layer_norm.weight\n",
      "decoder.layers.0.encoder_attn_layer_norm.bias\n",
      "decoder.layers.0.fc1.weight\n",
      "decoder.layers.0.fc1.bias\n",
      "decoder.layers.0.fc2.weight\n",
      "decoder.layers.0.fc2.bias\n",
      "decoder.layers.0.final_layer_norm.weight\n",
      "decoder.layers.0.final_layer_norm.bias\n",
      "decoder.layers.1.self_attn.k_proj.weight\n",
      "decoder.layers.1.self_attn.k_proj.bias\n",
      "decoder.layers.1.self_attn.v_proj.weight\n",
      "decoder.layers.1.self_attn.v_proj.bias\n",
      "decoder.layers.1.self_attn.q_proj.weight\n",
      "decoder.layers.1.self_attn.q_proj.bias\n",
      "decoder.layers.1.self_attn.out_proj.weight\n",
      "decoder.layers.1.self_attn.out_proj.bias\n",
      "decoder.layers.1.self_attn_layer_norm.weight\n",
      "decoder.layers.1.self_attn_layer_norm.bias\n",
      "decoder.layers.1.encoder_attn.k_proj.weight\n",
      "decoder.layers.1.encoder_attn.k_proj.bias\n",
      "decoder.layers.1.encoder_attn.v_proj.weight\n",
      "decoder.layers.1.encoder_attn.v_proj.bias\n",
      "decoder.layers.1.encoder_attn.q_proj.weight\n",
      "decoder.layers.1.encoder_attn.q_proj.bias\n",
      "decoder.layers.1.encoder_attn.out_proj.weight\n",
      "decoder.layers.1.encoder_attn.out_proj.bias\n",
      "decoder.layers.1.encoder_attn_layer_norm.weight\n",
      "decoder.layers.1.encoder_attn_layer_norm.bias\n",
      "decoder.layers.1.fc1.weight\n",
      "decoder.layers.1.fc1.bias\n",
      "decoder.layers.1.fc2.weight\n",
      "decoder.layers.1.fc2.bias\n",
      "decoder.layers.1.final_layer_norm.weight\n",
      "decoder.layers.1.final_layer_norm.bias\n",
      "decoder.layers.2.self_attn.k_proj.weight\n",
      "decoder.layers.2.self_attn.k_proj.bias\n",
      "decoder.layers.2.self_attn.v_proj.weight\n",
      "decoder.layers.2.self_attn.v_proj.bias\n",
      "decoder.layers.2.self_attn.q_proj.weight\n",
      "decoder.layers.2.self_attn.q_proj.bias\n",
      "decoder.layers.2.self_attn.out_proj.weight\n",
      "decoder.layers.2.self_attn.out_proj.bias\n",
      "decoder.layers.2.self_attn_layer_norm.weight\n",
      "decoder.layers.2.self_attn_layer_norm.bias\n",
      "decoder.layers.2.encoder_attn.k_proj.weight\n",
      "decoder.layers.2.encoder_attn.k_proj.bias\n",
      "decoder.layers.2.encoder_attn.v_proj.weight\n",
      "decoder.layers.2.encoder_attn.v_proj.bias\n",
      "decoder.layers.2.encoder_attn.q_proj.weight\n",
      "decoder.layers.2.encoder_attn.q_proj.bias\n",
      "decoder.layers.2.encoder_attn.out_proj.weight\n",
      "decoder.layers.2.encoder_attn.out_proj.bias\n",
      "decoder.layers.2.encoder_attn_layer_norm.weight\n",
      "decoder.layers.2.encoder_attn_layer_norm.bias\n",
      "decoder.layers.2.fc1.weight\n",
      "decoder.layers.2.fc1.bias\n",
      "decoder.layers.2.fc2.weight\n",
      "decoder.layers.2.fc2.bias\n",
      "decoder.layers.2.final_layer_norm.weight\n",
      "decoder.layers.2.final_layer_norm.bias\n",
      "decoder.layers.3.self_attn.k_proj.weight\n",
      "decoder.layers.3.self_attn.k_proj.bias\n",
      "decoder.layers.3.self_attn.v_proj.weight\n",
      "decoder.layers.3.self_attn.v_proj.bias\n",
      "decoder.layers.3.self_attn.q_proj.weight\n",
      "decoder.layers.3.self_attn.q_proj.bias\n",
      "decoder.layers.3.self_attn.out_proj.weight\n",
      "decoder.layers.3.self_attn.out_proj.bias\n",
      "decoder.layers.3.self_attn_layer_norm.weight\n",
      "decoder.layers.3.self_attn_layer_norm.bias\n",
      "decoder.layers.3.encoder_attn.k_proj.weight\n",
      "decoder.layers.3.encoder_attn.k_proj.bias\n",
      "decoder.layers.3.encoder_attn.v_proj.weight\n",
      "decoder.layers.3.encoder_attn.v_proj.bias\n",
      "decoder.layers.3.encoder_attn.q_proj.weight\n",
      "decoder.layers.3.encoder_attn.q_proj.bias\n",
      "decoder.layers.3.encoder_attn.out_proj.weight\n",
      "decoder.layers.3.encoder_attn.out_proj.bias\n",
      "decoder.layers.3.encoder_attn_layer_norm.weight\n",
      "decoder.layers.3.encoder_attn_layer_norm.bias\n",
      "decoder.layers.3.fc1.weight\n",
      "decoder.layers.3.fc1.bias\n",
      "decoder.layers.3.fc2.weight\n",
      "decoder.layers.3.fc2.bias\n",
      "decoder.layers.3.final_layer_norm.weight\n",
      "decoder.layers.3.final_layer_norm.bias\n",
      "decoder.layers.4.self_attn.k_proj.weight\n",
      "decoder.layers.4.self_attn.k_proj.bias\n",
      "decoder.layers.4.self_attn.v_proj.weight\n",
      "decoder.layers.4.self_attn.v_proj.bias\n",
      "decoder.layers.4.self_attn.q_proj.weight\n",
      "decoder.layers.4.self_attn.q_proj.bias\n",
      "decoder.layers.4.self_attn.out_proj.weight\n",
      "decoder.layers.4.self_attn.out_proj.bias\n",
      "decoder.layers.4.self_attn_layer_norm.weight\n",
      "decoder.layers.4.self_attn_layer_norm.bias\n",
      "decoder.layers.4.encoder_attn.k_proj.weight\n",
      "decoder.layers.4.encoder_attn.k_proj.bias\n",
      "decoder.layers.4.encoder_attn.v_proj.weight\n",
      "decoder.layers.4.encoder_attn.v_proj.bias\n",
      "decoder.layers.4.encoder_attn.q_proj.weight\n",
      "decoder.layers.4.encoder_attn.q_proj.bias\n",
      "decoder.layers.4.encoder_attn.out_proj.weight\n",
      "decoder.layers.4.encoder_attn.out_proj.bias\n",
      "decoder.layers.4.encoder_attn_layer_norm.weight\n",
      "decoder.layers.4.encoder_attn_layer_norm.bias\n",
      "decoder.layers.4.fc1.weight\n",
      "decoder.layers.4.fc1.bias\n",
      "decoder.layers.4.fc2.weight\n",
      "decoder.layers.4.fc2.bias\n",
      "decoder.layers.4.final_layer_norm.weight\n",
      "decoder.layers.4.final_layer_norm.bias\n",
      "decoder.layers.5.self_attn.k_proj.weight\n",
      "decoder.layers.5.self_attn.k_proj.bias\n",
      "decoder.layers.5.self_attn.v_proj.weight\n",
      "decoder.layers.5.self_attn.v_proj.bias\n",
      "decoder.layers.5.self_attn.q_proj.weight\n",
      "decoder.layers.5.self_attn.q_proj.bias\n",
      "decoder.layers.5.self_attn.out_proj.weight\n",
      "decoder.layers.5.self_attn.out_proj.bias\n",
      "decoder.layers.5.self_attn_layer_norm.weight\n",
      "decoder.layers.5.self_attn_layer_norm.bias\n",
      "decoder.layers.5.encoder_attn.k_proj.weight\n",
      "decoder.layers.5.encoder_attn.k_proj.bias\n",
      "decoder.layers.5.encoder_attn.v_proj.weight\n",
      "decoder.layers.5.encoder_attn.v_proj.bias\n",
      "decoder.layers.5.encoder_attn.q_proj.weight\n",
      "decoder.layers.5.encoder_attn.q_proj.bias\n",
      "decoder.layers.5.encoder_attn.out_proj.weight\n",
      "decoder.layers.5.encoder_attn.out_proj.bias\n",
      "decoder.layers.5.encoder_attn_layer_norm.weight\n",
      "decoder.layers.5.encoder_attn_layer_norm.bias\n",
      "decoder.layers.5.fc1.weight\n",
      "decoder.layers.5.fc1.bias\n",
      "decoder.layers.5.fc2.weight\n",
      "decoder.layers.5.fc2.bias\n",
      "decoder.layers.5.final_layer_norm.weight\n",
      "decoder.layers.5.final_layer_norm.bias\n",
      "decoder.output_projection.weight\n"
     ]
    }
   ],
   "source": [
    "for key in model.state_dict().keys():\n",
    "    print(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_transformer_checkpoint = \"../../checkpoints/transformer_iwslt16/checkpoint_best.pt\"\n",
    "pretrain_state = checkpoint_utils.load_checkpoint_to_cpu(pretrained_transformer_checkpoint)\n",
    "pretrain_state_dict = pretrain_state[\"model\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.all(model.state_dict()['encoder.transformer_encoder.layers.2.self_attn_layer_norm.weight'] == pretrain_state_dict['encoder.layers.2.self_attn_layer_norm.weight'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'decoder.embed_positions._float_tensor',\n",
       " 'decoder.version',\n",
       " 'encoder.fc1.bias',\n",
       " 'encoder.fc1.weight',\n",
       " 'encoder.fc2.bias',\n",
       " 'encoder.fc2.weight',\n",
       " 'encoder.final_layer_norm.bias',\n",
       " 'encoder.final_layer_norm.weight',\n",
       " 'encoder.layer_norm_query_sent.bias',\n",
       " 'encoder.layer_norm_query_sent.weight',\n",
       " 'encoder.layer_norm_query_word.bias',\n",
       " 'encoder.layer_norm_query_word.weight',\n",
       " 'encoder.layer_norm_sentence_level.bias',\n",
       " 'encoder.layer_norm_sentence_level.weight',\n",
       " 'encoder.layer_norm_word_level.bias',\n",
       " 'encoder.layer_norm_word_level.weight',\n",
       " 'encoder.linear.bias',\n",
       " 'encoder.linear.weight',\n",
       " 'encoder.sent_attn.k_proj.bias',\n",
       " 'encoder.sent_attn.k_proj.weight',\n",
       " 'encoder.sent_attn.out_proj.bias',\n",
       " 'encoder.sent_attn.out_proj.weight',\n",
       " 'encoder.sent_attn.q_proj.bias',\n",
       " 'encoder.sent_attn.q_proj.weight',\n",
       " 'encoder.sent_attn.v_proj.bias',\n",
       " 'encoder.sent_attn.v_proj.weight',\n",
       " 'encoder.transformer_encoder.embed_positions._float_tensor',\n",
       " 'encoder.transformer_encoder.version',\n",
       " 'encoder.word_attn.k_proj.bias',\n",
       " 'encoder.word_attn.k_proj.weight',\n",
       " 'encoder.word_attn.out_proj.bias',\n",
       " 'encoder.word_attn.out_proj.weight',\n",
       " 'encoder.word_attn.q_proj.bias',\n",
       " 'encoder.word_attn.q_proj.weight',\n",
       " 'encoder.word_attn.v_proj.bias',\n",
       " 'encoder.word_attn.v_proj.weight'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state_dict = model.state_dict()\n",
    "matching = set()\n",
    "for key in pretrain_state_dict.keys():\n",
    "                    for search_key in [\n",
    "                        \"embed_tokens\", \"embed_positions\", \"layers\",\n",
    "                        \"output_projection\"\n",
    "                    ]:\n",
    "                        if search_key in key:\n",
    "                            if \"encoder.\" in key:\n",
    "                                state_dict_key = \"encoder.transformer_encoder.\" + key[key.find(search_key):]\n",
    "                            else:\n",
    "                                state_dict_key = key\n",
    "                            assert state_dict_key in state_dict, \"{} not in state_dict\".format(state_dict_key)\n",
    "                            if torch.all(state_dict[state_dict_key]==pretrain_state_dict[key]):\n",
    "                                matching.add(state_dict_key)\n",
    "non_matching_keys = set(state_dict) - matching\n",
    "non_matching_keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Which modules are freezed?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HanEncoder\n",
      "\tHiddenStatesCache\n",
      "\tTransformerEncoder\n",
      "\tSigmoid\n",
      "TransformerDecoder\n",
      "\tFairseqDropout\n",
      "\tEmbedding\n",
      "\tSinusoidalPositionalEmbedding\n",
      "\tModuleList\n",
      "\tLinear\n"
     ]
    }
   ],
   "source": [
    "for module in model.children():\n",
    "    print(module._get_name())\n",
    "    for mod in module.children():\n",
    "        requires_grad = False\n",
    "        for p in mod.parameters():\n",
    "            if p.requires_grad == True:\n",
    "                requires_grad = True\n",
    "        if requires_grad == False:\n",
    "            print('\\t'+mod._get_name())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fc2 weight before train:\n",
      "Parameter containing:\n",
      "tensor([[-0.1340, -0.3676, -0.0647, -0.2280, -0.2112],\n",
      "        [ 0.1760, -0.4117, -0.3137, -0.1145,  0.0494],\n",
      "        [-0.0822,  0.3440, -0.2333,  0.4418, -0.1833],\n",
      "        [ 0.1637, -0.1574,  0.2041, -0.0125, -0.0325],\n",
      "        [-0.0159, -0.1092,  0.1287,  0.3985,  0.1439]], requires_grad=True)\n",
      "fc2 weight after train:\n",
      "Parameter containing:\n",
      "tensor([[-0.1343, -0.2993, -0.1109, -0.1309, -0.2369],\n",
      "        [ 0.1500, -0.3714, -0.3045, -0.0124,  0.0821],\n",
      "        [-0.1014,  0.2075, -0.1208,  0.3245, -0.0842],\n",
      "        [ 0.1652, -0.1190,  0.1817,  0.0053, -0.0595],\n",
      "        [-0.0374, -0.1891,  0.2104,  0.3621,  0.2304]], requires_grad=True)\n",
      "fc2 pretrained weight (same as the one above):\n",
      "Parameter containing:\n",
      "tensor([[-0.1343, -0.2993, -0.1109, -0.1309, -0.2369],\n",
      "        [ 0.1500, -0.3714, -0.3045, -0.0124,  0.0821],\n",
      "        [-0.1014,  0.2075, -0.1208,  0.3245, -0.0842],\n",
      "        [ 0.1652, -0.1190,  0.1817,  0.0053, -0.0595],\n",
      "        [-0.0374, -0.1891,  0.2104,  0.3621,  0.2304]], requires_grad=True)\n",
      "fc2 weight (frozen) after retrain:\n",
      "Parameter containing:\n",
      "tensor([[-0.1343, -0.2993, -0.1109, -0.1309, -0.2369],\n",
      "        [ 0.1500, -0.3714, -0.3045, -0.0124,  0.0821],\n",
      "        [-0.1014,  0.2075, -0.1208,  0.3245, -0.0842],\n",
      "        [ 0.1652, -0.1190,  0.1817,  0.0053, -0.0595],\n",
      "        [-0.0374, -0.1891,  0.2104,  0.3621,  0.2304]])\n",
      "fc2 weight (unfrozen) after re-retrain:\n",
      "Parameter containing:\n",
      "tensor([[-0.1499, -0.1969, -0.2100, -0.0055, -0.2143],\n",
      "        [ 0.2764, -0.4546, -0.0803, -0.2216,  0.0596],\n",
      "        [-0.0729,  0.0741,  0.1254,  0.0798, -0.0043],\n",
      "        [ 0.0944,  0.0594,  0.0307,  0.1376,  0.0827],\n",
      "        [-0.0438, -0.0450, -0.0550,  0.5629,  0.2057]], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "# toy feed-forward net\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(10, 5)\n",
    "        self.fc2 = nn.Linear(5, 5)\n",
    "        self.fc3 = nn.Linear(5, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# define random data\n",
    "random_input = Variable(torch.randn(10,))\n",
    "random_target = Variable(torch.randn(1,))\n",
    "\n",
    "# define net\n",
    "net = Net()\n",
    "\n",
    "# print fc2 weight\n",
    "print('fc2 weight before train:')\n",
    "print(net.fc2.weight)\n",
    "\n",
    "# train the net\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.1)\n",
    "for i in range(100):\n",
    "    net.zero_grad()\n",
    "    output = net(random_input)\n",
    "    loss = criterion(output, random_target)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# print the trained fc2 weight\n",
    "print('fc2 weight after train:')\n",
    "print(net.fc2.weight)\n",
    "\n",
    "# save the net\n",
    "torch.save(net.state_dict(), 'model')\n",
    "\n",
    "# delete and redefine the net\n",
    "del net\n",
    "net = Net()\n",
    "\n",
    "# load the weight\n",
    "net.load_state_dict(torch.load('model'))\n",
    "\n",
    "# print the pre-trained fc2 weight\n",
    "print('fc2 pretrained weight (same as the one above):')\n",
    "print(net.fc2.weight)\n",
    "\n",
    "# define new random data\n",
    "random_input = Variable(torch.randn(10,))\n",
    "random_target = Variable(torch.randn(1,))\n",
    "\n",
    "# we want to freeze the fc2 layer this time: only train fc1 and fc3\n",
    "net.fc2.weight.requires_grad = False\n",
    "net.fc2.bias.requires_grad = False\n",
    "\n",
    "# train again\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# NOTE: pytorch optimizer explicitly accepts parameter that requires grad\n",
    "# see https://github.com/pytorch/pytorch/issues/679\n",
    "optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)\n",
    "# this raises ValueError: optimizing a parameter that doesn't require gradients\n",
    "#optimizer = optim.Adam(net.parameters(), lr=0.1)\n",
    "\n",
    "for i in range(100):\n",
    "    net.zero_grad()\n",
    "    output = net(random_input)\n",
    "    loss = criterion(output, random_target)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# print the retrained fc2 weight\n",
    "# note that the weight is same as the one before retraining: only fc1 & fc3 changed\n",
    "print('fc2 weight (frozen) after retrain:')\n",
    "print(net.fc2.weight)\n",
    "\n",
    "# let's unfreeze the fc2 layer this time for extra tuning\n",
    "net.fc2.weight.requires_grad = True\n",
    "net.fc2.bias.requires_grad = True\n",
    "\n",
    "# add the unfrozen fc2 weight to the current optimizer\n",
    "optimizer.add_param_group({'params': net.fc2.parameters()})\n",
    "\n",
    "# re-retrain\n",
    "for i in range(100):\n",
    "    net.zero_grad()\n",
    "    output = net(random_input)\n",
    "    loss = criterion(output, random_target)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# print the re-retrained fc2 weight\n",
    "# note that this time the fc2 weight also changed\n",
    "print('fc2 weight (unfrozen) after re-retrain:')\n",
    "print(net.fc2.weight)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
