<kaggle_start><data_title>BMS-train-full<data_name>bmstrainpart1
<code>import numpy as np  # linear algebra
import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory
import os

inputDir = "/kaggle/input"


def _walkdir(dirPath):
    for dirname, _, filenames in os.walk(dirPath):
        for filename in filenames:
            print(os.path.join(dirname, filename))


# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
import gc

# gc.collect()


# # Refference
# 1. frog dudes masterpieces https://www.kaggle.com/c/bms-molecular-translation/discussion/231190
# ## 待研究
# - 青蛙做了什么优化
# 1. cnn（tnt） + tokenizer + transformer decoder 为什么比 cnn + lstm 好 ？
# ans: Ranger
# 2. 青蛙的tnt结构可以直接迁移到tpu上吗？ 似乎是可以的，结构中并没有啥np的结构
# 3. 什么样的情况下模型可以反复训练
# 4. 似乎结构上只有预处理 padding mask 上有优化
# - 为什么内存会超
# 1. dataloader 里有什么增加了内存
# 2. transformer 训练过程中有什么变量一直没释放
# 3. dataloader 太大，可以分批小范围load
# 4. 训练过程中做do_valid 不合适 非常占资源
# - 如果
# 1. 如果不做patch会省内存嘛？
# 2. 我能不能搭建个vit
# 3. 如果不处理内存爆掉的问题会如何
# - 结论
# 1. 想来想去，这个版本实在是太慢了，而且有包内存的问题，即使不做do validate 也只能持续20k个batch左右
# 2. 后面还是根据这个博客 手动搭建一下transformer https://www.kaggle.com/drzhuzhe/training-on-gpu-bms/
import math
import gc
import torch
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.sampler import *
import torch.nn.functional as F
from torch.nn.parallel.data_parallel import data_parallel
import random
import pickle
import time
from timeit import default_timer as timer
import cv2

# %pylab inline
import matplotlib.pyplot as plt

# import matplotlib.image as mpimg
import collections
from collections import defaultdict


# # Set up Torch
def seed_py(seed):
    random.seed(seed)
    np.random.seed(seed)
    return seed


def seed_torch(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    return seed


seed = int(time.time())  # 335202   #5202  #123  #
seed_py(seed)
seed_torch(seed)
torch.backends.cudnn.benchmark = False  ##uses the inbuilt cudnn auto-tuner to find the fastest convolution algorithms. -
torch.backends.cudnn.enabled = True
torch.backends.cudnn.deterministic = True
COMMON_STRING = "\tpytorch\n"
COMMON_STRING += "\t\tseed = %d\n" % seed
COMMON_STRING += "\t\ttorch.__version__              = %s\n" % torch.__version__
COMMON_STRING += "\t\ttorch.version.cuda             = %s\n" % torch.version.cuda
COMMON_STRING += (
    "\t\ttorch.backends.cudnn.version() = %s\n" % torch.backends.cudnn.version()
)
try:
    COMMON_STRING += (
        "\t\tos['CUDA_VISIBLE_DEVICES']     = %s\n" % os.environ["CUDA_VISIBLE_DEVICES"]
    )
    NUM_CUDA_DEVICES = len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))
except Exception:
    COMMON_STRING += "\t\tos['CUDA_VISIBLE_DEVICES']     = None\n"
    NUM_CUDA_DEVICES = 1
COMMON_STRING += "\t\ttorch.cuda.device_count()      = %d\n" % torch.cuda.device_count()
COMMON_STRING += (
    "\t\ttorch.cuda.get_device_properties() = %s\n"
    % str(torch.cuda.get_device_properties(0))[21:]
)
COMMON_STRING += "\n"


# # Configure
origin_data_dir = "../input/bms-molecular-translation"
data_dir = "../input/bmstrainpart1/bms-moleular-translation"
csv_data_dir = "../input/bmd-mate-csv/csv"
STOI = {
    "<sos>": 190,
    "<eos>": 191,
    "<pad>": 192,
}
patch_size = 16
pixel_pad = 3
pixel_scale = 0.8  # 1.0  #0.62=36/58 #1.0
max_length = 300  # 278 #275
pixel_stride = 4
pixel_dim = 24
patch_dim = 384
_walkdir(csv_data_dir)


# # Utils
def read_pickle_from_file(pickle_file):
    with open(pickle_file, "rb") as f:
        x = pickle.load(f)
    return x


def uncompress_array(compressed_k):
    compressed_k.seek(0)
    k = np.load(compressed_k, allow_pickle=True)["arr_0"]
    return k


# draw -----------------------------------
def image_show(name, image, resize=0.1):
    H, W = image.shape[0:2]
    plt.figure(figsize=(round(resize * W), round(resize * H)), dpi=20)
    # fig = plt.figure(figsize = (500,500), dpi=20)
    # ax = fig.add_subplot(1)
    plt.imshow(image, cmap=plt.cm.gray)
    # ax.imshow(image)
    # image = cv2.resize(image, (round(resize*W), round(resize*H)), interpolation = cv2.INTER_AREA)
    # plt.imshow(image)


def resize_image(image, scale=1):
    if scale == 1:
        f = pixel_scale * 58 / 36  # 1.2414 #0.80555
        b = int(round(36 * 0.5))
    if scale == 2:
        f = pixel_scale * 1
        b = int(round(58 * 0.5))
    image = image[b:-b, b:-b]  # remove border
    if not np.isclose(1, f, rtol=1e-02, atol=1e-02):
        h, w = image.shape
        fw = int(round(f * w))
        fh = int(round(f * h))
        image = cv2.resize(image, dsize=(fw, fh), interpolation=cv2.INTER_AREA)
    return image


def repad_image(image, multiplier=16):
    h, w = image.shape
    fh = int(np.ceil(h / multiplier)) * multiplier
    fw = int(np.ceil(w / multiplier)) * multiplier
    m = np.full((fh, fw), 255, np.uint8)
    m[0:h, 0:w] = image
    return m


# it seems to fill patches to Images
def patch_to_image(patch, coord, width, height):
    image = np.full((height, width), 255, np.uint8)
    p = pixel_pad
    patch = patch[:, p:-p, p:-p]
    num_patch = len(patch)
    for i in range(num_patch):
        y, x = coord[i]
        x = x * patch_size
        y = y * patch_size
        image[y : y + patch_size, x : x + patch_size] = patch[i]
        cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)
    return image


def image_to_patch(image, patch_size, pixel_pad, threshold=0):
    p = pixel_pad
    h, w = image.shape
    x, y = np.meshgrid(np.arange(w // patch_size), np.arange(h // patch_size))
    yx = np.stack([y, x], 2).reshape(-1, 2)
    s = patch_size + 2 * p
    m = torch.from_numpy(image).reshape(1, 1, h, w).float()
    k = F.unfold(m, kernel_size=s, stride=patch_size, padding=p)
    k = k.permute(0, 2, 1).reshape(-1, s * s)
    k = k.data.cpu().numpy().reshape(-1, s, s)
    # print(k.shape)
    sum = (1 - k[:, p:-p, p:-p] / 255).reshape(len(k), -1).sum(-1)
    # only store patches
    i = np.where(sum > threshold)
    # print(sum)
    patch = k[i]
    coord = yx[i]
    return patch, coord


"""
# check folder data
#df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')
df = read_pickle_from_file(csv_data_dir+'/df_train.more.csv.pickle')
df_fold = pd.read_csv(csv_data_dir+'/df_fold.fine.csv')
df_meta = pd.read_csv(csv_data_dir+'/df_train_image_meta.csv')
df = df.merge(df_fold, on='image_id')
df = df.merge(df_meta, on='image_id')
print(len(df), len(df_fold), len(df_meta))
print(len(df['fold'].unique()))
df.head(10)
"""


# divide dataset into training and validate
def make_fold(mode="train-1"):
    if "train" in mode:
        df = read_pickle_from_file(csv_data_dir + "/df_train.more.csv.pickle")
        # df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')
        # df_fold = pd.read_csv(data_dir+'/df_fold.csv')
        df_fold = pd.read_csv(csv_data_dir + "/df_fold.fine.csv")
        df_meta = pd.read_csv(csv_data_dir + "/df_train_image_meta.csv")
        df = df.merge(df_fold, on="image_id")
        df = df.merge(df_meta, on="image_id")
        df.loc[:, "path"] = "train_patch16_s0.800"
        df["fold"] = df["fold"].astype(int)
        # print(df.groupby(['fold']).size()) #404_031
        # print(df.columns)
        fold = int(mode[-1]) * 10
        # print(fold)
        df_train = df[df.fold != fold].reset_index(drop=True)
        df_valid = df[df.fold == fold].reset_index(drop=True)
        return df_train, df_valid
    # Index(['image_id', 'InChI'], dtype='object')
    if "test" in mode:
        # df = pd.read_csv(data_dir+'/sample_submission.csv')
        df = pd.read_csv(data_dir + "/submit_lb3.80.csv")
        df_meta = pd.read_csv(data_dir + "/df_test_image_meta.csv")
        df = df.merge(df_meta, on="image_id")
        df.loc[:, "path"] = "test"
        # df.loc[:, 'InChI'] = '0'
        df.loc[:, "formula"] = "0"
        df.loc[:, "text"] = "0"
        df.loc[:, "sequence"] = pd.Series([[0]] * len(df))
        df.loc[:, "length"] = df.InChI.str.len()
        df_test = df
        return df_test


# df_train, df_valid = make_fold()
# print(len(df_train), len(df_valid))
# -----------------------------------------------------------------------
# tokenization, padding, ...
def pad_sequence_to_max_length(sequence, max_length, padding_value):
    batch_size = len(sequence)
    pad_sequence = np.full((batch_size, max_length), padding_value, np.int32)
    for b, s in enumerate(sequence):
        L = len(s)
        pad_sequence[b, :L, ...] = s
    return pad_sequence


def load_tokenizer():
    tokenizer = YNakamaTokenizer(is_load=True)
    print("len(tokenizer) : vocab_size", len(tokenizer))
    for k, v in STOI.items():
        assert tokenizer.stoi[k] == v
    return tokenizer


def null_augment(r):
    return r


# getDataset
class BmsDataset(Dataset):
    def __init__(self, df, tokenizer, augment=null_augment):
        super().__init__()
        self.tokenizer = tokenizer
        self.df = df
        self.augment = augment
        self.length = len(self.df)

    def __str__(self):
        string = ""
        string += "\tlen = %d\n" % len(self)
        string += "\tdf  = %s\n" % str(self.df.shape)
        g = self.df["length"].values.astype(np.int32) // 20
        g = np.bincount(g, minlength=14)
        string += "\tlength distribution\n"
        for n in range(14):
            string += "\t\t %3d = %8d (%0.4f)\n" % ((n + 1) * 20, g[n], g[n] / g.sum())
        return string

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        d = self.df.iloc[index]
        token = d.sequence
        patch_file = data_dir + "/%s/%s/%s/%s/%s.pickle" % (
            d.path,
            d.image_id[0],
            d.image_id[1],
            d.image_id[2],
            d.image_id,
        )
        k = read_pickle_from_file(patch_file)
        patch = uncompress_array(k["patch"])
        patch = np.concatenate(
            [
                np.zeros(
                    (1, patch_size + 2 * pixel_pad, patch_size + 2 * pixel_pad),
                    np.uint8,
                ),
                patch,
            ],
            0,
        )  # cls token
        coord = k["coord"]
        w = k["width"]
        h = k["height"]
        h = h // patch_size - 1
        w = w // patch_size - 1
        coord = np.insert(coord, 0, [h, w], 0)  # cls token
        # debug
        # image = patch_to_image(patch, coord, k['width' ], k['height'])
        # image_show('image', image, resize=1)
        # cv2.waitKey(0)
        # image = cv2.imread(image_file,cv2.IMREAD_GRAYSCALE) #
        r = {
            "index": index,
            "image_id": d.image_id,
            "InChI": d.InChI,
            "d": d,
            "token": token,
            #'image' : image,
            "patch": patch,
            "coord": coord,
        }
        if self.augment is not None:
            r = self.augment(r)
        # Warrning runtime keep this variable shall cause Out of memery crash
        del patch_file, patch, coord, k
        # gc.collect()
        return r


# set up tokenizer
class YNakamaTokenizer(object):
    def __init__(self, is_load=True):
        self.stoi = {}
        self.itos = {}
        if is_load:
            self.stoi = _TOKENIZER_
            self.itos = {k: v for v, k in self.stoi.items()}

    def __len__(self):
        return len(self.stoi)

    def build_vocab(self, text):
        vocab = set()
        for t in text:
            vocab.update(t.split(" "))
        vocab = sorted(vocab)
        vocab.append("<sos>")
        vocab.append("<eos>")
        vocab.append("<pad>")
        for i, s in enumerate(vocab):
            self.stoi[s] = i
        self.itos = {k: v for v, k in self.stoi.items()}

    def one_text_to_sequence(self, text):
        sequence = []
        sequence.append(self.stoi["<sos>"])
        for s in text.split(" "):
            sequence.append(self.stoi[s])
        sequence.append(self.stoi["<eos>"])
        return sequence

    def one_sequence_to_text(self, sequence):
        return "".join(list(map(lambda i: self.itos[i], sequence)))

    def one_predict_to_inchi(self, predict):
        inchi = "InChI=1S/"
        for p in predict:
            if p == self.stoi["<eos>"] or p == self.stoi["<pad>"]:
                break
            inchi += self.itos[p]
        return inchi

    # ---
    def text_to_sequence(self, text):
        sequence = [self.one_text_to_sequence(t) for t in text]
        return sequence

    def sequence_to_text(self, sequence):
        text = [self.one_sequence_to_text(s) for s in sequence]
        return text

    def predict_to_inchi(self, predict):
        inchi = [self.one_predict_to_inchi(p) for p in predict]
        return inchi


def null_collate(batch, is_sort_decreasing_length=True):
    collate = defaultdict(list)
    if is_sort_decreasing_length:  # sort by decreasing length
        sort = np.argsort([-len(r["token"]) for r in batch])
        batch = [batch[s] for s in sort]
    for r in batch:
        for k, v in r.items():
            collate[k].append(v)
    # ----
    batch_size = len(batch)
    collate["length"] = [len(l) for l in collate["token"]]
    token = [np.array(t, np.int32) for t in collate["token"]]
    token = pad_sequence_to_max_length(
        token, max_length=max_length, padding_value=STOI["<pad>"]
    )
    collate["token"] = torch.from_numpy(token).long()
    max_of_length = max(collate["length"])
    token_pad_mask = np.zeros((batch_size, max_of_length, max_of_length))
    for b in range(batch_size):
        L = collate["length"][b]
        token_pad_mask[b, :L, :L] = 1  # +1 for cls_token
    collate["token_pad_mask"] = torch.from_numpy(token_pad_mask).byte()
    # -----
    # image = np.stack(collate['image'])
    # image = image.astype(np.float32) / 255
    # collate['image'] = torch.from_numpy(image).unsqueeze(1).repeat(1,3,1,1)
    # -----
    collate["num_patch"] = [len(l) for l in collate["patch"]]
    max_of_num_patch = max(collate["num_patch"])
    patch_pad_mask = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))
    patch = np.full(
        (
            batch_size,
            max_of_num_patch,
            patch_size + 2 * pixel_pad,
            patch_size + 2 * pixel_pad,
        ),
        255,
    )  # pad as 255
    coord = np.zeros((batch_size, max_of_num_patch, 2))
    for b in range(batch_size):
        N = collate["num_patch"][b]
        patch[b, :N] = collate["patch"][b]
        coord[b, :N] = collate["coord"][b]
        patch_pad_mask[b, :N, :N] = 1  # +1 for cls_token
    collate["patch"] = torch.from_numpy(patch).half() / 255
    collate["coord"] = torch.from_numpy(coord).long()
    collate["patch_pad_mask"] = torch.from_numpy(patch_pad_mask).byte()
    del patch_pad_mask, patch, coord, token_pad_mask, batch
    return collate


def run_check_dataset():
    tokenizer = load_tokenizer()
    df_train, df_valid = make_fold("train-1")
    # df_train = make_fold('test') #1616107
    # dataset = BmsDataset(df_train, tokenizer, remote_augment)
    # dataset = BmsDataset(df_valid, tokenizer)
    # print(dataset)
    train_dataset = BmsDataset(df_train, tokenizer)
    # for i in range(len(dataset)):
    """
    for i in range(5):
        #i = np.random.choice(len(dataset))
        r = dataset[i]
        print(r['index'])
        print(r['image_id'])
        #print(r['formula'])
        print(r['InChI'])
        print(r['token'])
        print('image : ')
        #print('\t', r['image'].shape)
        print('')
        #---
        image = patch_to_image(r['patch'], r['coord'], width=1024, height=1024)
        #image_show('image', image, resize=.1)
        del image
        gc.collect()
        #cv2.waitKey(0)
    """
    # exit(0)
    # """
    loader = DataLoader(
        train_dataset,
        # sampler = RandomSampler(train_dataset),
        batch_size=8,
        drop_last=True,
        num_workers=8,
        pin_memory=True,
        collate_fn=null_collate,
    )
    for t, batch in enumerate(loader):
        # if t>30: break
        print(t, "-----------")
        # print('index : ', batch['index'])
        """
        print('image : ')
        #print('\t', batch['image'].shape, batch['image'].is_contiguous())
        print('\t', batch['patch'].shape, batch['patch'].is_contiguous())
        print('\t', batch['coord'].shape, batch['coord'].is_contiguous())
        #print('\t', batch['mask'].shape, batch['mask'].is_contiguous())
        print('length  : ')
        print('\t',len( batch['length']))
        print('\t', batch['length'])
        print('token  : ')
        print('\t', batch['token'].shape, batch['token'].is_contiguous())
        """
        # print('\t', batch['token'])
        # print('')
        del t, batch
    del loader, train_dataset, df_train, df_valid, tokenizer
    # """


# 数据load到4000 条 左右就会挂 ， 讲道理 load数据应该不会不断加内存吧
# run_check_dataset()


# # Net Modeules
# ## 1D Encoder and Decoder
from typing import Tuple, Dict
import torch.nn as nn
from fairseq import utils
from fairseq.models import *
from fairseq.modules import *


# ------------------------------------------------------
# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need
class PositionEncode1D(nn.Module):
    def __init__(self, dim, max_length):
        super().__init__()
        assert dim % 2 == 0
        self.max_length = max_length
        d = torch.exp(torch.arange(0.0, dim, 2) * (-math.log(10000.0) / dim))
        position = torch.arange(0.0, max_length).unsqueeze(1)
        pos = torch.zeros(1, max_length, dim)
        pos[0, :, 0::2] = torch.sin(position * d)
        pos[0, :, 1::2] = torch.cos(position * d)
        self.register_buffer("pos", pos)

    def forward(self, x):
        batch_size, T, dim = x.shape
        x = x + self.pos[:, :T]
        return x


"""
# https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py
#https://github.com/pytorch/fairseq/issues/568
# fairseq/fairseq/models/fairseq_encoder.py
# https://github.com/pytorch/fairseq/blob/master/fairseq/modules/transformer_layer.py
class TransformerEncode(FairseqEncoder):
    def __init__(self, dim, ff_dim, num_head, num_layer):
        super().__init__({})
        #print('my TransformerEncode()')
        self.layer = nn.ModuleList([
            TransformerEncoderLayer(Namespace({
                'encoder_embed_dim': dim,
                'encoder_attention_heads': num_head,
                'attention_dropout': 0.1,
                'dropout': 0.1,
                'encoder_normalize_before': True,
                'encoder_ffn_embed_dim': ff_dim,
            })) for i in range(num_layer)
        ])
        self.layer_norm = nn.LayerNorm(dim)
    def forward(self, x):# T x B x C
        #print('my TransformerEncode forward()')
        for layer in self.layer:
            x = layer(x)
        x = self.layer_norm(x)
        return x
"""


# https://mt.cs.upc.edu/2020/12/21/the-transformer-fairseq-edition/
# for debug
# class TransformerDecode(FairseqDecoder):
#     def __init__(self, dim, ff_dim, num_head, num_layer):
#         super().__init__({})
#         print('my TransformerDecode()')
#
#         self.layer = nn.ModuleList([
#             TransformerDecoderLayer(Namespace({
#                 'decoder_embed_dim': dim,
#                 'decoder_attention_heads': num_head,
#                 'attention_dropout': 0.1,
#                 'dropout': 0.1,
#                 'decoder_normalize_before': True,
#                 'decoder_ffn_embed_dim': ff_dim,
#             })) for i in range(num_layer)
#         ])
#         self.layer_norm = nn.LayerNorm(dim)
#
#
#     def forward(self, x, mem, x_mask):# T x B x C
#         print('my TransformerDecode forward()')
#         for layer in self.layer:
#             x = layer(x, mem, self_attn_mask=x_mask)[0]
#         x = self.layer_norm(x)
#         return x  # T x B x C
# https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute
class Namespace(object):
    def __init__(self, adict):
        self.__dict__.update(adict)


# https://fairseq.readthedocs.io/en/latest/tutorial_simple_lstm.html
# see https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py
class TransformerDecode(FairseqIncrementalDecoder):
    def __init__(self, dim, ff_dim, num_head, num_layer):
        super().__init__({})
        # print('my TransformerDecode()')
        self.layer = nn.ModuleList(
            [
                TransformerDecoderLayer(
                    Namespace(
                        {
                            "decoder_embed_dim": dim,
                            "decoder_attention_heads": num_head,
                            "attention_dropout": 0.1,
                            "dropout": 0.1,
                            "decoder_normalize_before": True,
                            "decoder_ffn_embed_dim": ff_dim,
                        }
                    )
                )
                for i in range(num_layer)
            ]
        )
        self.layer_norm = nn.LayerNorm(dim)

    def forward(self, x, mem, x_mask, x_pad_mask, mem_pad_mask):
        # print('my TransformerDecode forward()')
        for layer in self.layer:
            x = layer(
                x,
                mem,
                self_attn_mask=x_mask,
                self_attn_padding_mask=x_pad_mask,
                encoder_padding_mask=mem_pad_mask,
            )[0]
        x = self.layer_norm(x)
        return x  # T x B x C

    # def forward_one(self, x, mem, incremental_state):
    def forward_one(
        self,
        x: Tensor,
        mem: Tensor,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
    ) -> Tensor:
        x = x[-1:]
        for layer in self.layer:
            x = layer(x, mem, incremental_state=incremental_state)[0]
        x = self.layer_norm(x)
        return x


"""
https://fairseq.readthedocs.io/en/latest/_modules/fairseq/modules/transformer_layer.html
class TransformerDecoderLayer(nn.Module):
def forward(
        self,
        x,
        encoder_out: Optional[torch.Tensor] = None,
        encoder_padding_mask: Optional[torch.Tensor] = None,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
        prev_self_attn_state: Optional[List[torch.Tensor]] = None,
        prev_attn_state: Optional[List[torch.Tensor]] = None,
        self_attn_mask: Optional[torch.Tensor] = None,
        self_attn_padding_mask: Optional[torch.Tensor] = None,
        need_attn: bool = False,
        need_head_weights: bool = False,
    ):
https://github.com/pytorch/fairseq/blob/05b86005bcca0155319fa9b81abfd69f63c06906/examples/simultaneous_translation/utils/data_utils.py#L31      
"""


# ## TNT
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence

# https://github.com/pytorch/pytorch/issues/1788
# https://stackoverflow.com/questions/51030782/why-do-we-pack-the-sequences-in-pytorch
from timm.models.vision_transformer import Mlp
from timm.models.layers import DropPath, trunc_normal_


# from timm.models.tnt import *
# 这个多头注意力会造成内存不断增加吗？
class Attention(nn.Module):
    """Multi-Head Attention"""

    def __init__(
        self, dim, hidden_dim, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        head_dim = hidden_dim // num_heads
        self.head_dim = head_dim
        self.scale = head_dim**-0.5
        self.qk = nn.Linear(dim, hidden_dim * 2, bias=qkv_bias)
        self.v = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop, inplace=True)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop, inplace=True)

    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        B, N, C = x.shape
        qk = (
            self.qk(x)
            .reshape(B, N, 2, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k = qk[0], qk[1]  # make torchscript happy (cannot use tensor as tuple)
        v = self.v(x).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        # ---
        attn = (q @ k.transpose(-2, -1)) * self.scale  # B x self.num_heads x NxN
        if mask is not None:
            # mask = mask.unsqueeze(1).repeat(1,self.num_heads,1,1)
            mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
            attn = attn.masked_fill(mask == 0, -6e4)
            # attn = attn.masked_fill(mask == 0, -half('inf'))
            # https://github.com/NVIDIA/apex/issues/93
            # How to use fp16 training with masked operations
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
        x = self.proj(x)
        x = self.proj_drop(x)
        del mask, attn
        return x


class Block(nn.Module):
    def __init__(
        self,
        dim,
        in_dim,
        num_pixel,
        num_heads=12,
        in_num_head=4,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        # Inner transformer
        self.norm_in = norm_layer(in_dim)
        self.attn_in = Attention(
            in_dim,
            in_dim,
            num_heads=in_num_head,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.norm_mlp_in = norm_layer(in_dim)
        self.mlp_in = Mlp(
            in_features=in_dim,
            hidden_features=int(in_dim * 4),
            out_features=in_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm1_proj = norm_layer(in_dim)
        self.proj = nn.Linear(in_dim * num_pixel, dim, bias=True)
        # Outer transformer
        self.norm_out = norm_layer(dim)
        self.attn_out = Attention(
            dim,
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm_mlp = norm_layer(dim)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),
            out_features=dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, pixel_embed, patch_embed, mask):
        # inner
        pixel_embed = pixel_embed + self.drop_path(
            self.attn_in(self.norm_in(pixel_embed))
        )
        pixel_embed = pixel_embed + self.drop_path(
            self.mlp_in(self.norm_mlp_in(pixel_embed))
        )
        # outer
        B, N, C = patch_embed.size()
        patch_embed[:, 1:] = (
            patch_embed[:, 1:]
            + self.proj(self.norm1_proj(pixel_embed).reshape(B, N, -1))[:, 1:]
        )
        patch_embed = patch_embed + self.drop_path(
            self.attn_out(self.norm_out(patch_embed), mask)
        )
        patch_embed = patch_embed + self.drop_path(self.mlp(self.norm_mlp(patch_embed)))
        return pixel_embed, patch_embed


# ---------------------------------
class PixelEmbed(nn.Module):
    def __init__(self, patch_size=16, in_dim=48, stride=4):
        super().__init__()
        self.in_dim = in_dim
        self.proj = nn.Conv2d(3, self.in_dim, kernel_size=7, padding=0, stride=stride)

    def forward(self, patch, pixel_pos):
        BN = len(patch)
        x = patch
        x = self.proj(x)
        # x = x.transpose(1, 2).reshape(B * self.num_patches, self.in_dim, self.new_patch_size, self.new_patch_size)
        x = x + pixel_pos
        x = x.reshape(BN, self.in_dim, -1).transpose(1, 2)
        return x


# ---------------------------------
class TNT(nn.Module):
    """Transformer in Transformer - https://arxiv.org/abs/2103.00112"""

    def __init__(
        self,
        patch_size=patch_size,
        embed_dim=patch_dim,
        in_dim=pixel_dim,
        depth=12,
        num_heads=6,
        in_num_head=4,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        norm_layer=nn.LayerNorm,
        first_stride=pixel_stride,
    ):
        super().__init__()
        self.num_features = (
            self.embed_dim
        ) = embed_dim  # num_features for consistency with other models
        self.pixel_embed = PixelEmbed(
            patch_size=patch_size, in_dim=in_dim, stride=first_stride
        )
        # num_patches = self.pixel_embed.num_patches
        # self.num_patches = num_patches
        new_patch_size = 4  # self.pixel_embed.new_patch_size
        num_pixel = new_patch_size**2
        self.norm1_proj = norm_layer(num_pixel * in_dim)
        self.proj = nn.Linear(num_pixel * in_dim, embed_dim)
        self.norm2_proj = norm_layer(embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.patch_pos = nn.Embedding(
            100 * 100, embed_dim
        )  # nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pixel_pos = nn.Parameter(
            torch.zeros(1, in_dim, new_patch_size, new_patch_size)
        )
        self.pos_drop = nn.Dropout(p=drop_rate)
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, depth)
        ]  # stochastic depth decay rule
        blocks = []
        for i in range(depth):
            blocks.append(
                Block(
                    dim=embed_dim,
                    in_dim=in_dim,
                    num_pixel=num_pixel,
                    num_heads=num_heads,
                    in_num_head=in_num_head,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[i],
                    norm_layer=norm_layer,
                )
            )
        self.blocks = nn.ModuleList(blocks)
        self.norm = norm_layer(embed_dim)
        trunc_normal_(self.cls_token, std=0.02)
        # trunc_normal_(self.patch_pos, std=.02)
        trunc_normal_(self.pixel_pos, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {"patch_pos", "pixel_pos", "cls_token"}

    def forward(self, patch, coord, mask):
        B = len(patch)
        batch_size, max_of_num_patch, s, s = patch.shape
        patch = patch.reshape(batch_size * max_of_num_patch, 1, s, s).repeat(1, 3, 1, 1)
        pixel_embed = self.pixel_embed(patch, self.pixel_pos)
        patch_embed = self.norm2_proj(
            self.proj(self.norm1_proj(pixel_embed.reshape(B, max_of_num_patch, -1)))
        )
        # patch_embed = torch.cat((self.cls_token.expand(B, -1, -1), patch_embed), dim=1)
        # patch_embed = patch_embed + self.patch_pos
        # patch_embed[:, 1:] = patch_embed[:, 1:] + self.patch_pos(coord[:, :, 0] * 100 + coord[:, :, 1])
        patch_embed[:, :1] = self.cls_token.expand(B, -1, -1)
        patch_embed = patch_embed + self.patch_pos(
            coord[:, :, 0] * 100 + coord[:, :, 1]
        )
        patch_embed = self.pos_drop(patch_embed)
        for blk in self.blocks:
            pixel_embed, patch_embed = blk(pixel_embed, patch_embed, mask)
        patch_embed = self.norm(patch_embed)
        del patch, mask
        return patch_embed


def make_dummy_data():
    # make dummy data
    # image_id,width,height,scale,orientation
    meta = [
        [
            "000011a64c74",
            325,
            229,
            2,
            0,
        ],
        [
            "000019cc0cd2",
            288,
            148,
            1,
            0,
        ],
        [
            "0000252b6d2b",
            509,
            335,
            2,
            0,
        ],
        [
            "000026b49b7e",
            243,
            177,
            1,
            0,
        ],
        [
            "000026fc6c36",
            294,
            112,
            1,
            0,
        ],
        [
            "000028818203",
            402,
            328,
            2,
            0,
        ],
        [
            "000029a61c01",
            395,
            294,
            2,
            0,
        ],
        [
            "000035624718",
            309,
            145,
            1,
            0,
        ],
    ]
    batch_size = 8
    # <todo> check border for padding
    # <todo> pepper noise
    batch = {
        "num_patch": [],
        "patch": [],
        "coord": [],
    }
    for b in range(batch_size):
        image_id = meta[b][0]
        scale = meta[b][3]
        image_file = origin_data_dir + "/%s/%s/%s/%s/%s.png" % (
            "train",
            image_id[0],
            image_id[1],
            image_id[2],
            image_id,
        )
        # print(image_file)
        image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)
        image = resize_image(image, scale)
        image = repad_image(image, patch_size)  # remove border and repad
        # print(image.shape)
        k, yx = image_to_patch(image, patch_size, pixel_pad, threshold=0)
        for y, x in yx:
            # cv2.circle(image,(x,y),8,128,1)
            x = x * patch_size
            y = y * patch_size
            cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)
        # image_show('image-%d' % b, image, resize=1)
        del image
        # cv2.waitKey(1)
        batch["patch"].append(k)
        batch["coord"].append(yx)
        batch["num_patch"].append(len(k))
    # ----
    max_of_num_patch = max(batch["num_patch"])
    mask = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))
    patch = np.zeros(
        (
            batch_size,
            max_of_num_patch,
            patch_size + 2 * pixel_pad,
            patch_size + 2 * pixel_pad,
        )
    )
    coord = np.zeros((batch_size, max_of_num_patch, 2))
    for b in range(batch_size):
        N = batch["num_patch"][b]
        patch[b, :N] = batch["patch"][b]
        coord[b, :N] = batch["coord"][b]
        mask[b, :N, :N] = 1
    num_patch = batch["num_patch"]
    patch = torch.from_numpy(patch).float()
    coord = torch.from_numpy(coord).long()
    mask = torch.from_numpy(mask).byte()
    return patch, coord, num_patch, mask


def run_check_tnt_patch():
    patch, coord, num_patch, mask = make_dummy_data()
    tnt = TNT()
    patch_embed = tnt(patch, coord, mask)
    print(patch_embed.shape)


# run_check_tnt_patch()


# ## TNT + PositionEncode1D + Transformer Decoder
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence

vocab_size = 193
text_dim = 384
decoder_dim = 384
num_layer = 3
num_head = 8
ff_dim = 1024


class Net(nn.Module):
    def __init__(
        self,
    ):
        super(Net, self).__init__()
        self.cnn = TNT()
        self.image_encode = nn.Identity()
        # ---
        self.text_pos = PositionEncode1D(text_dim, max_length)
        self.token_embed = nn.Embedding(vocab_size, text_dim)
        self.text_decode = TransformerDecode(decoder_dim, ff_dim, num_head, num_layer)
        # ---
        self.logit = nn.Linear(decoder_dim, vocab_size)
        self.dropout = nn.Dropout(p=0.5)
        # ----
        # initialization
        self.token_embed.weight.data.uniform_(-0.1, 0.1)
        self.logit.bias.data.fill_(0)
        self.logit.weight.data.uniform_(-0.1, 0.1)

    @torch.jit.unused
    def forward(self, patch, coord, token, patch_pad_mask, token_pad_mask):
        device = patch.device
        batch_size = len(patch)
        # ---
        patch = patch * 2 - 1
        image_embed = self.cnn(patch, coord, patch_pad_mask)
        image_embed = self.image_encode(image_embed).permute(1, 0, 2).contiguous()
        text_embed = self.token_embed(token)
        text_embed = self.text_pos(text_embed).permute(1, 0, 2).contiguous()
        max_of_length = token_pad_mask.shape[-1]
        text_mask = np.triu(np.ones((max_of_length, max_of_length)), k=1).astype(
            np.uint8
        )
        text_mask = torch.autograd.Variable(torch.from_numpy(text_mask) == 1).to(device)
        # ----
        # <todo> perturb mask as aug
        text_pad_mask = token_pad_mask[:, :, 0] == 0
        image_pad_mask = patch_pad_mask[:, :, 0] == 0
        x = self.text_decode(
            text_embed[:max_of_length],
            image_embed,
            text_mask,
            text_pad_mask,
            image_pad_mask,
        )
        x = x.permute(1, 0, 2).contiguous()
        l = self.logit(x)
        logit = torch.zeros((batch_size, max_length, vocab_size), device=device)
        logit[:, :max_of_length] = l
        del image_embed, text_mask, text_pad_mask, image_pad_mask
        return logit

    # submit function has not been coded. i will leave it as an exercise for the kaggler
    # @torch.jit.export
    # def forward_argmax_decode(self, patch, coord, mask):
    #
    #     image_dim   = 384
    #     text_dim    = 384
    #     decoder_dim = 384
    #     num_layer = 3
    #     num_head  = 8
    #     ff_dim    = 1024
    #
    #     STOI = {
    #         '<sos>': 190,
    #         '<eos>': 191,
    #         '<pad>': 192,
    #     }
    #     max_length = 278 # 275
    #
    #
    #     #---------------------------------
    #     device = patch.device
    #     batch_size = len(patch)
    #
    #     patch = patch*2-1
    #     image_embed = self.cnn(patch, coord, mask)
    #     image_embed = self.image_encode(image_embed).permute(1,0,2).contiguous()
    #
    #     token = torch.full((batch_size, max_length), STOI['<pad>'],dtype=torch.long, device=device)
    #     text_pos = self.text_pos.pos
    #     token[:,0] = STOI['<sos>']
    #
    #
    #     #-------------------------------------
    #     eos = STOI['<eos>']
    #     pad = STOI['<pad>']
    #     # https://github.com/alexmt-scale/causal-transformer-decoder/blob/master/tests/test_consistency.py
    #     # slow version
    #     # if 0:
    #     #     for t in range(max_length-1):
    #     #         last_token = token [:,:(t+1)]
    #     #         text_embed = self.token_embed(last_token)
    #     #         text_embed = self.text_pos(text_embed).permute(1,0,2).contiguous() #text_embed + text_pos[:,:(t+1)] #
    #     #
    #     #         text_mask = np.triu(np.ones((t+1, t+1)), k=1).astype(np.uint8)
    #     #         text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)==1).to(device)
    #     #
    #     #         x = self.text_decode(text_embed, image_embed, text_mask)
    #     #         x = x.permute(1,0,2).contiguous()
    #     #
    #     #         l = self.logit(x[:,-1])
    #     #         k = torch.argmax(l, -1)  # predict max
    #     #         token[:, t+1] = k
    #     #         if ((k == eos) | (k == pad)).all():  break
    #
    #     # fast version
    #     if 1:
    #         #incremental_state = {}
    #         incremental_state = torch.jit.annotate(
    #             Dict[str, Dict[str, Optional[Tensor]]],
    #             torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {}),
    #         )
    #         for t in range(max_length-1):
    #             #last_token = token [:,:(t+1)]
    #             #text_embed = self.token_embed(last_token)
    #             #text_embed = self.text_pos(text_embed) #text_embed + text_pos[:,:(t+1)] #
    #
    #             last_token = token[:, t]
    #             text_embed = self.token_embed(last_token)
    #             text_embed = text_embed + text_pos[:,t] #
    #             text_embed = text_embed.reshape(1,batch_size,text_dim)
    #
    #             x = self.text_decode.forward_one(text_embed, image_embed, incremental_state)
    #             x = x.reshape(batch_size,decoder_dim)
    #             #print(incremental_state.keys())
    #
    #             l = self.logit(x)
    #             k = torch.argmax(l, -1)  # predict max
    #             token[:, t+1] = k
    #             if ((k == eos) | (k == pad)).all():  break
    #
    #     predict = token[:, 1:]
    #     return predict


# loss #################################################################
def seq_cross_entropy_loss(logit, token, length):
    truth = token[:, 1:]
    L = [l - 1 for l in length]
    logit = pack_padded_sequence(logit, L, batch_first=True).data
    truth = pack_padded_sequence(truth, L, batch_first=True).data
    loss = F.cross_entropy(logit, truth, ignore_index=STOI["<pad>"])
    return loss


# https://www.aclweb.org/anthology/2020.findings-emnlp.276.pdf
def seq_focal_cross_entropy_loss(logit, token, length):
    gamma = 0.5  # {0.5,1.0}
    # label_smooth = 0.90
    # ---
    truth = token[:, 1:]
    L = [l - 1 for l in length]
    logit = pack_padded_sequence(logit, L, batch_first=True).data
    truth = pack_padded_sequence(truth, L, batch_first=True).data
    # loss = F.cross_entropy(logit, truth, ignore_index=STOI['<pad>'])
    # non_pad = torch.where(truth != STOI['<pad>'])[0]  # & (t!=STOI['<sos>'])
    # ---
    # p = F.softmax(logit,-1)
    # logp = - torch.log(torch.clamp(p, 1e-4, 1 - 1e-4))
    logp = F.log_softmax(logit, -1)
    logp = logp.gather(1, truth.reshape(-1, 1)).reshape(-1)
    p = logp.exp()
    loss = -((1 - p) ** gamma) * logp  # focal
    # loss = - ((1 + p) ** gamma)*logp  #anti-focal
    loss = loss.mean()
    return loss


def np_loss_cross_entropy(probability, truth):
    batch_size = len(probability)
    truth = truth.reshape(-1)
    p = probability[np.arange(batch_size), truth]
    loss = -np.log(np.clip(p, 1e-6, 1))
    loss = loss.mean()
    return loss


# check #################################################################
def run_check_net():
    patch, coord, num_patch, patch_pad_mask = make_dummy_data()
    batch_size = len(patch)
    token = np.full((batch_size, max_length), STOI["<pad>"], np.int64)  # token
    length = np.random.randint(5, max_length - 2, batch_size)
    length = np.sort(length)[::-1].copy()
    max_of_length = max(length)
    token_pad_mask = np.zeros((batch_size, max_of_length, max_of_length))
    for b in range(batch_size):
        l = length[b]
        t = np.random.choice(vocab_size, l)
        t = np.insert(t, 0, STOI["<sos>"])
        t = np.insert(t, len(t), STOI["<eos>"])
        L = len(t)
        token[b, :L] = t
        token_pad_mask[b, :L, :L] = 1
    token = torch.from_numpy(token).long()
    token_pad_mask = torch.from_numpy(token_pad_mask).byte()
    # ---
    net = Net()
    net.train()
    logit = net(patch, coord, token, patch_pad_mask, token_pad_mask)
    loss = seq_cross_entropy_loss(logit, token, length)
    print("vocab_size", vocab_size)
    print("max_length", max_length)
    print("")
    print(length)
    print(length.shape)
    print(token.shape)
    print("---")
    print(logit.shape)
    print(loss)
    print("---")
    del net
    # ---
    # print('torch.jit.script(net)')
    # net.eval()
    # net = torch.jit.script(net)
    #
    # predict = net.forward_argmax_decode(patch, coord, mask)
    # print(predict.shape)


# run_check_net()
"""
import sys
local_vars = list(locals().items())
for var, obj in local_vars:
    print(var, sys.getsizeof(obj))
"""
# ----------------
is_mixed_precision = True  # False  #
###################################################################################################
import torch.cuda.amp as amp

if is_mixed_precision:

    class AmpNet(Net):
        @torch.cuda.amp.autocast()
        def forward(self, *args):
            return super(AmpNet, self).forward(*args)

else:
    AmpNet = Net
import sys
from datetime import datetime
import Levenshtein

IDENTIFIER = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")


# http://stackoverflow.com/questions/34950201/pycharm-print-end-r-statement-not-working
class Logger(object):
    def __init__(self):
        self.terminal = sys.stdout  # stdout
        self.file = None

    def open(self, file, mode=None):
        if mode is None:
            mode = "w"
        self.file = open(file, mode)

    def write(self, message, is_terminal=1, is_file=1):
        if "\r" in message:
            is_file = 0
        if is_terminal == 1:
            self.terminal.write(message)
            self.terminal.flush()
            # time.sleep(1)
        if is_file == 1:
            self.file.write(message)
            self.file.flush()

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass


# etc ------------------------------------
def time_to_str(t, mode="min"):
    if mode == "min":
        t = int(t) / 60
        hr = t // 60
        min = t % 60
        return "%2d hr %02d min" % (hr, min)
    elif mode == "sec":
        t = int(t)
        min = t // 60
        sec = t % 60
        return "%2d min %02d sec" % (min, sec)
    else:
        raise NotImplementedError


# ## Optimizer
from torch.optim.optimizer import Optimizer
import itertools as it


def get_learning_rate(optimizer):
    lr = []
    for param_group in optimizer.param_groups:
        lr += [param_group["lr"]]
    assert len(lr) == 1  # we support only one param_group
    lr = lr[0]
    return lr


class RAdam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        self.buffer = [[None, None, None] for ind in range(10)]
        super(RAdam, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(RAdam, self).__setstate__(state)

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError("RAdam does not support sparse gradients")
                p_data_fp32 = p.data.float()
                state = self.state[p]
                if len(state) == 0:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
                else:
                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                beta1, beta2 = group["betas"]
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                state["step"] += 1
                buffered = self.buffer[int(state["step"] % 10)]
                if state["step"] == buffered[0]:
                    N_sma, step_size = buffered[1], buffered[2]
                else:
                    buffered[0] = state["step"]
                    beta2_t = beta2 ** state["step"]
                    N_sma_max = 2 / (1 - beta2) - 1
                    N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
                    buffered[1] = N_sma
                    # more conservative since it's an approximated value
                    if N_sma >= 5:
                        step_size = math.sqrt(
                            (1 - beta2_t)
                            * (N_sma - 4)
                            / (N_sma_max - 4)
                            * (N_sma - 2)
                            / N_sma
                            * N_sma_max
                            / (N_sma_max - 2)
                        ) / (1 - beta1 ** state["step"])
                    else:
                        step_size = 1.0 / (1 - beta1 ** state["step"])
                    buffered[2] = step_size
                if group["weight_decay"] != 0:
                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
                # more conservative since it's an approximated value
                if N_sma >= 5:
                    denom = exp_avg_sq.sqrt().add_(group["eps"])
                    p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size * group["lr"])
                else:
                    p_data_fp32.add_(exp_avg, alpha=-step_size * group["lr"])
                p.data.copy_(p_data_fp32)
        return loss


class Lookahead(Optimizer):
    def __init__(self, optimizer, alpha=0.5, k=6):
        if not 0.0 <= alpha <= 1.0:
            raise ValueError(f"Invalid slow update rate: {alpha}")
        if not 1 <= k:
            raise ValueError(f"Invalid lookahead steps: {k}")
        self.optimizer = optimizer
        self.param_groups = self.optimizer.param_groups
        self.alpha = alpha
        self.k = k
        for group in self.param_groups:
            group["step_counter"] = 0
        self.slow_weights = [
            [p.clone().detach() for p in group["params"]] for group in self.param_groups
        ]
        for w in it.chain(*self.slow_weights):
            w.requires_grad = False
        self.state = optimizer.state

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
        loss = self.optimizer.step()
        for group, slow_weights in zip(self.param_groups, self.slow_weights):
            group["step_counter"] += 1
            if group["step_counter"] % self.k != 0:
                continue
            for p, q in zip(group["params"], slow_weights):
                if p.grad is None:
                    continue
                q.data.add_(p.data - q.data, alpha=self.alpha)
                p.data.copy_(q.data)
        return loss


# ## Training
def do_valid(net, tokenizer, valid_loader):
    valid_probability = []
    valid_truth = []
    valid_length = []
    valid_num = 0
    net.eval()
    start_timer = timer()
    for t, batch in enumerate(valid_loader):
        batch_size = len(batch["index"])
        length = batch["length"]
        token = batch["token"].cuda()
        token_pad_mask = batch["token_pad_mask"].cuda()
        # image  = batch['image' ].cuda()
        num_patch = batch["num_patch"]
        patch = batch["patch"].cuda()
        coord = batch["coord"].cuda()
        patch_pad_mask = batch["patch_pad_mask"].cuda()
        with torch.no_grad():
            logit = data_parallel(
                net, (patch, coord, token, patch_pad_mask, token_pad_mask)
            )  # net(image, token, length)
            probability = F.softmax(logit, -1)
        valid_num += batch_size
        valid_probability.append(probability.data.cpu().numpy())
        valid_truth.append(token.data.cpu().numpy())
        valid_length.extend(length)
        print(
            "\r %8d / %d  %s"
            % (
                valid_num,
                len(valid_loader.sampler),
                time_to_str(timer() - start_timer, "sec"),
            ),
            end="",
            flush=True,
        )
        del t, batch
    assert valid_num == len(valid_loader.sampler)  # len(valid_loader.dataset))
    # print('')
    # ----------------------
    probability = np.concatenate(valid_probability)
    predict = probability.argmax(-1)
    truth = np.concatenate(valid_truth)
    length = valid_length
    # ----
    p = probability[:, :-1].reshape(-1, vocab_size)
    t = truth[:, 1:].reshape(-1)
    non_pad = np.where(t != STOI["<pad>"])[0]  # & (t!=STOI['<sos>'])
    p = p[non_pad]
    t = t[non_pad]
    loss = np_loss_cross_entropy(p, t)
    # ----
    lb_score = 0
    if 1:
        score = []
        for i, (p, t) in enumerate(zip(predict, truth)):
            t = truth[i][1 : length[i] - 1]
            p = predict[i][1 : length[i] - 1]
            t = tokenizer.one_predict_to_inchi(t)
            p = tokenizer.one_predict_to_inchi(p)
            s = Levenshtein.distance(p, t)
            score.append(s)
        lb_score = np.mean(score)
    """
    if 1:
        score = []
        for i, (p, t) in enumerate(zip(predict, truth)):
            t = truth[i][1:length[i]-1]     # in the buggy version, i have used 1 instead of i
            p = predict[i][1:length[i]-1]
            t = tokenizer.one_predict_to_inchi(t)
            p = tokenizer.one_predict_to_inchi(p)
            s = Levenshtein.distance(p, t)
            score.append(s)
        lb_score = np.mean(score)
    """
    # lb_score = compute_lb_score(k, t)
    del (
        valid_loader,
        net,
        predict,
        truth,
        valid_probability,
        valid_truth,
        valid_length,
        valid_num,
    )
    return [loss, lb_score]


def run_train():
    fold = 3
    out_dir = "./tnt-patch1-s0.8/fold%d" % fold
    initial_checkpoint = None
    # initial_checkpoint = \
    # out_dir + '/checkpoint/00755000_model.pth'#None #
    #'/root/share1/kaggle/2021/bms-moleular-translation/result/try22/tnt-patch1/fold3/checkpoint/00697000_model.pth'
    debug = 0
    start_lr = 0.00001  # 1
    batch_size = 32  # 24
    ## setup  ----------------------------------------
    for f in ["checkpoint", "train", "valid", "backup"]:
        os.makedirs(out_dir + "/" + f, exist_ok=True)
    # backup_project_as_zip(PROJECT_PATH, out_dir +'/backup/code.train.%s.zip'%IDENTIFIER)
    log = Logger()
    log.open(out_dir + "/log.train.txt", mode="a")
    log.write("\n--- [START %s] %s\n\n" % (IDENTIFIER, "-" * 64))
    log.write("\t%s\n" % COMMON_STRING)
    # log.write('\t__file__ = %s\n' % __file__)
    log.write("\tout_dir  = %s\n" % out_dir)
    log.write("\n")
    ## dataset ------------------------------------
    df_train, df_valid = make_fold("train-%d" % fold)
    df_valid = df_valid.iloc[:5_000]
    tokenizer = load_tokenizer()
    train_dataset = BmsDataset(df_train, tokenizer)
    valid_dataset = BmsDataset(df_valid, tokenizer)
    train_loader = DataLoader(
        train_dataset,
        sampler=RandomSampler(train_dataset),
        # sampler=UniformLengthSampler(train_dataset, is_shuffle=True), #200_000
        batch_size=batch_size,
        drop_last=True,
        num_workers=8,
        pin_memory=True,
        worker_init_fn=lambda id: np.random.seed(torch.initial_seed() // 2**32 + id),
        collate_fn=null_collate,
    )
    valid_loader = DataLoader(
        valid_dataset,
        # sampler=UniformLengthSampler(valid_dataset, 5_000),
        sampler=SequentialSampler(valid_dataset),
        batch_size=32,
        drop_last=False,
        num_workers=8,
        pin_memory=True,
        collate_fn=null_collate,
    )
    log.write("train_dataset : \n%s\n" % (train_dataset))
    log.write("valid_dataset : \n%s\n" % (valid_dataset))
    log.write("\n")
    ## net ----------------------------------------
    log.write("** net setting **\n")
    scaler = amp.GradScaler()
    net = AmpNet().cuda()
    if initial_checkpoint is not None:
        f = torch.load(initial_checkpoint, map_location=lambda storage, loc: storage)
        start_iteration = f["iteration"]
        start_epoch = f["epoch"]
        state_dict = f["state_dict"]
        # ---
        # state_dict = {k.replace('cnn.e.','cnn.'):v for k,v in state_dict.items()}
        # del state_dict['text_pos.pos']
        # del state_dict['cnn.head.weight']
        # del state_dict['cnn.head.bias']
        # net.load_state_dict(state_dict, strict=False)
        # ---
        net.load_state_dict(state_dict, strict=True)  # True
    else:
        start_iteration = 0
        start_epoch = 0
    log.write("\tinitial_checkpoint = %s\n" % initial_checkpoint)
    log.write("\n")
    # -----------------------------------------------
    if 0:  ##freeze
        for p in net.encoder.parameters():
            p.requires_grad = False
    optimizer = Lookahead(
        RAdam(filter(lambda p: p.requires_grad, net.parameters()), lr=start_lr),
        alpha=0.5,
        k=5,
    )
    # optimizer = RAdam(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr)
    num_iteration = 10 * 1000
    # num_iteration = 10
    iter_log = 1000
    iter_valid = 1000
    iter_save = list(range(0, num_iteration, 1000))  # 1*1000
    # iter_save = [0]
    log.write("optimizer\n  %s\n" % (optimizer))
    log.write("\n")
    ## start training here! ##############################################
    log.write("** start training here! **\n")
    log.write("   is_mixed_precision = %s \n" % str(is_mixed_precision))
    log.write("   batch_size = %d\n" % (batch_size))
    # log.write('   experiment = %s\n' % str(__file__.split('/')[-2:]))
    log.write(
        "                      |----- VALID ---|---- TRAIN/BATCH --------------\n"
    )
    log.write(
        "rate     iter   epoch | loss  lb(lev) | loss0  loss1  | time          \n"
    )
    log.write(
        "----------------------------------------------------------------------\n"
    )

    # 0.00000   0.00* 0.00  | 0.000  0.000  | 0.000  0.000  |  0 hr 00 min
    def message(mode="print"):
        if mode == ("print"):
            asterisk = " "
            loss = batch_loss
        if mode == ("log"):
            asterisk = "*" if iteration in iter_save else " "
            loss = train_loss
        #'%4.3f  %5.2f  | ' % (*valid_loss,) + \
        text = (
            "%0.5f  %5.4f%s %4.2f  | "
            % (
                rate,
                iteration / 10000,
                asterisk,
                epoch,
            )
            + "%4.3f  %4.3f  %4.3f  | " % (*loss,)
            + "%s" % (time_to_str(timer() - start_timer, "min"))
        )
        return text

    # ----
    valid_loss = np.zeros(2, np.float32)
    train_loss = np.zeros(3, np.float32)
    batch_loss = np.zeros_like(train_loss)
    sum_train_loss = np.zeros_like(train_loss)
    sum_train = 0
    loss0 = torch.FloatTensor([0]).cuda().sum()
    loss1 = torch.FloatTensor([0]).cuda().sum()
    loss2 = torch.FloatTensor([0]).cuda().sum()
    start_timer = timer()
    iteration = start_iteration
    epoch = start_epoch
    rate = 0
    # while iteration < num_iteration:
    for t, batch in enumerate(train_loader):
        if iteration in iter_save:
            if iteration != start_iteration:
                torch.save(
                    {
                        "state_dict": net.state_dict(),
                        "iteration": iteration,
                        "epoch": epoch,
                    },
                    out_dir + "/checkpoint/%08d_model.pth" % (iteration),
                )
                pass
        """
        if (iteration % iter_valid == 0):
            if iteration != start_iteration:
                valid_loss = do_valid(net, tokenizer, valid_loader)  #
                pass
        """
        if iteration % iter_log == 0:
            print("\r", end="", flush=True)
            log.write(message(mode="log") + "\n")
        # learning rate schduler ------------
        rate = get_learning_rate(optimizer)
        # one iteration update  -------------
        batch_size = len(batch["index"])
        length = batch["length"]
        token = batch["token"].cuda()
        token_pad_mask = batch["token_pad_mask"].cuda()
        # image  = batch['image' ].cuda()
        num_patch = batch["num_patch"]
        patch = batch["patch"].cuda()
        coord = batch["coord"].cuda()
        patch_pad_mask = batch["patch_pad_mask"].cuda()
        # ----
        net.train()
        optimizer.zero_grad()
        if is_mixed_precision:
            with amp.autocast():
                # assert(False)
                logit = data_parallel(
                    net, (patch, coord, token, patch_pad_mask, token_pad_mask)
                )  # net(image, token, length)
                loss0 = seq_cross_entropy_loss(logit, token, length)
                # loss0 = seq_anti_focal_cross_entropy_loss(logit, token, length)
            scaler.scale(loss0).backward()
            # scaler.unscale_(optimizer)
            # torch.nn.utils.clip_grad_norm_(net.parameters(), 2)
            scaler.step(optimizer)
            scaler.update()
        else:
            assert False
            # print('fp32')
            # image_embed = encoder(image)
            logit, weight = decoder(image_embed, token, length)
            (loss0).backward()
            optimizer.step()
        # print statistics  --------
        epoch += 1 / len(train_loader)
        iteration += 1
        batch_loss = np.array([loss0.item(), loss1.item(), loss2.item()])
        sum_train_loss += batch_loss
        sum_train += 1
        if iteration % 100 == 0:
            train_loss = sum_train_loss / (sum_train + 1e-12)
            sum_train_loss[...] = 0
            sum_train = 0
        print("\r", end="", flush=True)
        print(message(mode="print"), end="", flush=True)
        # debug--------------------------
        if debug:
            pass
        # delete per batch
        del (
            batch_size,
            length,
            token,
            token_pad_mask,
            num_patch,
            patch,
            coord,
            patch_pad_mask,
            t,
            batch,
        )
        if iteration > num_iteration:
            break
    log.write("\n")


run_train()
#!ls ./tnt-patch1-s0.8/fold3/checkpoint/
# list(range(0, 3000, 1000))
"""
import sys
# These are the usual ipython objects, including this one you are creating
ipython_vars = ['In', 'Out', 'exit', 'quit', 'get_ipython', 'ipython_vars']
# Get a sorted list of the objects and their sizes
sorted([(x, sys.getsizeof(globals().get(x))) for x in dir() if not x.startswith('_') and x not in sys.modules and x not in ipython_vars], key=lambda x: x[1], reverse=True)
"""
<|endoftext|>
[{"Id": 63434456, "ScriptId": 17031148, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 4321793, "CreationDate": "05/20/2021 12:54:39", "VersionNumber": 7.0, "Title": "Training on GPU - BMS", "EvaluationDate": "05/20/2021", "IsChange": true, "TotalLines": 1818.0, "LinesInsertedFromPrevious": 147.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 1671.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 2}]
[{"Id": 83325210, "KernelVersionId": 63434456, "SourceDatasetVersionId": 2239678}, {"Id": 83325211, "KernelVersionId": 63434456, "SourceDatasetVersionId": 2240275}]
[{"Id": 2239678, "DatasetId": 1343985, "DatasourceVersionId": 2281534, "CreatorUserId": 4321793, "LicenseName": "Unknown", "CreationDate": "05/17/2021 08:17:11", "VersionNumber": 3.0, "Title": "BMS-train-full", "Slug": "bmstrainpart1", "Subtitle": NaN, "Description": NaN, "VersionNotes": "tar.gz.file", "TotalCompressedBytes": 0.0, "TotalUncompressedBytes": 0.0}]
[{"Id": 1343985, "CreatorUserId": 4321793, "OwnerUserId": 4321793.0, "OwnerOrganizationId": NaN, "CurrentDatasetVersionId": 2239678.0, "CurrentDatasourceVersionId": 2281534.0, "ForumId": 1362981, "Type": 2, "CreationDate": "05/16/2021 11:02:38", "LastActivityDate": "05/16/2021", "TotalViews": 898, "TotalDownloads": 1, "TotalVotes": 1, "TotalKernels": 1}]
import numpy as np  # linear algebra
import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory
import os

inputDir = "/kaggle/input"


def _walkdir(dirPath):
    for dirname, _, filenames in os.walk(dirPath):
        for filename in filenames:
            print(os.path.join(dirname, filename))


# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
import gc

# gc.collect()


# # Refference
# 1. frog dudes masterpieces https://www.kaggle.com/c/bms-molecular-translation/discussion/231190
# ## 待研究
# - 青蛙做了什么优化
# 1. cnn（tnt） + tokenizer + transformer decoder 为什么比 cnn + lstm 好 ？
# ans: Ranger
# 2. 青蛙的tnt结构可以直接迁移到tpu上吗？ 似乎是可以的，结构中并没有啥np的结构
# 3. 什么样的情况下模型可以反复训练
# 4. 似乎结构上只有预处理 padding mask 上有优化
# - 为什么内存会超
# 1. dataloader 里有什么增加了内存
# 2. transformer 训练过程中有什么变量一直没释放
# 3. dataloader 太大，可以分批小范围load
# 4. 训练过程中做do_valid 不合适 非常占资源
# - 如果
# 1. 如果不做patch会省内存嘛？
# 2. 我能不能搭建个vit
# 3. 如果不处理内存爆掉的问题会如何
# - 结论
# 1. 想来想去，这个版本实在是太慢了，而且有包内存的问题，即使不做do validate 也只能持续20k个batch左右
# 2. 后面还是根据这个博客 手动搭建一下transformer https://www.kaggle.com/drzhuzhe/training-on-gpu-bms/
import math
import gc
import torch
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.sampler import *
import torch.nn.functional as F
from torch.nn.parallel.data_parallel import data_parallel
import random
import pickle
import time
from timeit import default_timer as timer
import cv2

# %pylab inline
import matplotlib.pyplot as plt

# import matplotlib.image as mpimg
import collections
from collections import defaultdict


# # Set up Torch
def seed_py(seed):
    random.seed(seed)
    np.random.seed(seed)
    return seed


def seed_torch(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    return seed


seed = int(time.time())  # 335202   #5202  #123  #
seed_py(seed)
seed_torch(seed)
torch.backends.cudnn.benchmark = False  ##uses the inbuilt cudnn auto-tuner to find the fastest convolution algorithms. -
torch.backends.cudnn.enabled = True
torch.backends.cudnn.deterministic = True
COMMON_STRING = "\tpytorch\n"
COMMON_STRING += "\t\tseed = %d\n" % seed
COMMON_STRING += "\t\ttorch.__version__              = %s\n" % torch.__version__
COMMON_STRING += "\t\ttorch.version.cuda             = %s\n" % torch.version.cuda
COMMON_STRING += (
    "\t\ttorch.backends.cudnn.version() = %s\n" % torch.backends.cudnn.version()
)
try:
    COMMON_STRING += (
        "\t\tos['CUDA_VISIBLE_DEVICES']     = %s\n" % os.environ["CUDA_VISIBLE_DEVICES"]
    )
    NUM_CUDA_DEVICES = len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))
except Exception:
    COMMON_STRING += "\t\tos['CUDA_VISIBLE_DEVICES']     = None\n"
    NUM_CUDA_DEVICES = 1
COMMON_STRING += "\t\ttorch.cuda.device_count()      = %d\n" % torch.cuda.device_count()
COMMON_STRING += (
    "\t\ttorch.cuda.get_device_properties() = %s\n"
    % str(torch.cuda.get_device_properties(0))[21:]
)
COMMON_STRING += "\n"


# # Configure
origin_data_dir = "../input/bms-molecular-translation"
data_dir = "../input/bmstrainpart1/bms-moleular-translation"
csv_data_dir = "../input/bmd-mate-csv/csv"
STOI = {
    "<sos>": 190,
    "<eos>": 191,
    "<pad>": 192,
}
patch_size = 16
pixel_pad = 3
pixel_scale = 0.8  # 1.0  #0.62=36/58 #1.0
max_length = 300  # 278 #275
pixel_stride = 4
pixel_dim = 24
patch_dim = 384
_walkdir(csv_data_dir)


# # Utils
def read_pickle_from_file(pickle_file):
    with open(pickle_file, "rb") as f:
        x = pickle.load(f)
    return x


def uncompress_array(compressed_k):
    compressed_k.seek(0)
    k = np.load(compressed_k, allow_pickle=True)["arr_0"]
    return k


# draw -----------------------------------
def image_show(name, image, resize=0.1):
    H, W = image.shape[0:2]
    plt.figure(figsize=(round(resize * W), round(resize * H)), dpi=20)
    # fig = plt.figure(figsize = (500,500), dpi=20)
    # ax = fig.add_subplot(1)
    plt.imshow(image, cmap=plt.cm.gray)
    # ax.imshow(image)
    # image = cv2.resize(image, (round(resize*W), round(resize*H)), interpolation = cv2.INTER_AREA)
    # plt.imshow(image)


def resize_image(image, scale=1):
    if scale == 1:
        f = pixel_scale * 58 / 36  # 1.2414 #0.80555
        b = int(round(36 * 0.5))
    if scale == 2:
        f = pixel_scale * 1
        b = int(round(58 * 0.5))
    image = image[b:-b, b:-b]  # remove border
    if not np.isclose(1, f, rtol=1e-02, atol=1e-02):
        h, w = image.shape
        fw = int(round(f * w))
        fh = int(round(f * h))
        image = cv2.resize(image, dsize=(fw, fh), interpolation=cv2.INTER_AREA)
    return image


def repad_image(image, multiplier=16):
    h, w = image.shape
    fh = int(np.ceil(h / multiplier)) * multiplier
    fw = int(np.ceil(w / multiplier)) * multiplier
    m = np.full((fh, fw), 255, np.uint8)
    m[0:h, 0:w] = image
    return m


# it seems to fill patches to Images
def patch_to_image(patch, coord, width, height):
    image = np.full((height, width), 255, np.uint8)
    p = pixel_pad
    patch = patch[:, p:-p, p:-p]
    num_patch = len(patch)
    for i in range(num_patch):
        y, x = coord[i]
        x = x * patch_size
        y = y * patch_size
        image[y : y + patch_size, x : x + patch_size] = patch[i]
        cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)
    return image


def image_to_patch(image, patch_size, pixel_pad, threshold=0):
    p = pixel_pad
    h, w = image.shape
    x, y = np.meshgrid(np.arange(w // patch_size), np.arange(h // patch_size))
    yx = np.stack([y, x], 2).reshape(-1, 2)
    s = patch_size + 2 * p
    m = torch.from_numpy(image).reshape(1, 1, h, w).float()
    k = F.unfold(m, kernel_size=s, stride=patch_size, padding=p)
    k = k.permute(0, 2, 1).reshape(-1, s * s)
    k = k.data.cpu().numpy().reshape(-1, s, s)
    # print(k.shape)
    sum = (1 - k[:, p:-p, p:-p] / 255).reshape(len(k), -1).sum(-1)
    # only store patches
    i = np.where(sum > threshold)
    # print(sum)
    patch = k[i]
    coord = yx[i]
    return patch, coord


"""
# check folder data
#df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')
df = read_pickle_from_file(csv_data_dir+'/df_train.more.csv.pickle')
df_fold = pd.read_csv(csv_data_dir+'/df_fold.fine.csv')
df_meta = pd.read_csv(csv_data_dir+'/df_train_image_meta.csv')
df = df.merge(df_fold, on='image_id')
df = df.merge(df_meta, on='image_id')
print(len(df), len(df_fold), len(df_meta))
print(len(df['fold'].unique()))
df.head(10)
"""


# divide dataset into training and validate
def make_fold(mode="train-1"):
    if "train" in mode:
        df = read_pickle_from_file(csv_data_dir + "/df_train.more.csv.pickle")
        # df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')
        # df_fold = pd.read_csv(data_dir+'/df_fold.csv')
        df_fold = pd.read_csv(csv_data_dir + "/df_fold.fine.csv")
        df_meta = pd.read_csv(csv_data_dir + "/df_train_image_meta.csv")
        df = df.merge(df_fold, on="image_id")
        df = df.merge(df_meta, on="image_id")
        df.loc[:, "path"] = "train_patch16_s0.800"
        df["fold"] = df["fold"].astype(int)
        # print(df.groupby(['fold']).size()) #404_031
        # print(df.columns)
        fold = int(mode[-1]) * 10
        # print(fold)
        df_train = df[df.fold != fold].reset_index(drop=True)
        df_valid = df[df.fold == fold].reset_index(drop=True)
        return df_train, df_valid
    # Index(['image_id', 'InChI'], dtype='object')
    if "test" in mode:
        # df = pd.read_csv(data_dir+'/sample_submission.csv')
        df = pd.read_csv(data_dir + "/submit_lb3.80.csv")
        df_meta = pd.read_csv(data_dir + "/df_test_image_meta.csv")
        df = df.merge(df_meta, on="image_id")
        df.loc[:, "path"] = "test"
        # df.loc[:, 'InChI'] = '0'
        df.loc[:, "formula"] = "0"
        df.loc[:, "text"] = "0"
        df.loc[:, "sequence"] = pd.Series([[0]] * len(df))
        df.loc[:, "length"] = df.InChI.str.len()
        df_test = df
        return df_test


# df_train, df_valid = make_fold()
# print(len(df_train), len(df_valid))
# -----------------------------------------------------------------------
# tokenization, padding, ...
def pad_sequence_to_max_length(sequence, max_length, padding_value):
    batch_size = len(sequence)
    pad_sequence = np.full((batch_size, max_length), padding_value, np.int32)
    for b, s in enumerate(sequence):
        L = len(s)
        pad_sequence[b, :L, ...] = s
    return pad_sequence


def load_tokenizer():
    tokenizer = YNakamaTokenizer(is_load=True)
    print("len(tokenizer) : vocab_size", len(tokenizer))
    for k, v in STOI.items():
        assert tokenizer.stoi[k] == v
    return tokenizer


def null_augment(r):
    return r


# getDataset
class BmsDataset(Dataset):
    def __init__(self, df, tokenizer, augment=null_augment):
        super().__init__()
        self.tokenizer = tokenizer
        self.df = df
        self.augment = augment
        self.length = len(self.df)

    def __str__(self):
        string = ""
        string += "\tlen = %d\n" % len(self)
        string += "\tdf  = %s\n" % str(self.df.shape)
        g = self.df["length"].values.astype(np.int32) // 20
        g = np.bincount(g, minlength=14)
        string += "\tlength distribution\n"
        for n in range(14):
            string += "\t\t %3d = %8d (%0.4f)\n" % ((n + 1) * 20, g[n], g[n] / g.sum())
        return string

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        d = self.df.iloc[index]
        token = d.sequence
        patch_file = data_dir + "/%s/%s/%s/%s/%s.pickle" % (
            d.path,
            d.image_id[0],
            d.image_id[1],
            d.image_id[2],
            d.image_id,
        )
        k = read_pickle_from_file(patch_file)
        patch = uncompress_array(k["patch"])
        patch = np.concatenate(
            [
                np.zeros(
                    (1, patch_size + 2 * pixel_pad, patch_size + 2 * pixel_pad),
                    np.uint8,
                ),
                patch,
            ],
            0,
        )  # cls token
        coord = k["coord"]
        w = k["width"]
        h = k["height"]
        h = h // patch_size - 1
        w = w // patch_size - 1
        coord = np.insert(coord, 0, [h, w], 0)  # cls token
        # debug
        # image = patch_to_image(patch, coord, k['width' ], k['height'])
        # image_show('image', image, resize=1)
        # cv2.waitKey(0)
        # image = cv2.imread(image_file,cv2.IMREAD_GRAYSCALE) #
        r = {
            "index": index,
            "image_id": d.image_id,
            "InChI": d.InChI,
            "d": d,
            "token": token,
            #'image' : image,
            "patch": patch,
            "coord": coord,
        }
        if self.augment is not None:
            r = self.augment(r)
        # Warrning runtime keep this variable shall cause Out of memery crash
        del patch_file, patch, coord, k
        # gc.collect()
        return r


# set up tokenizer
class YNakamaTokenizer(object):
    def __init__(self, is_load=True):
        self.stoi = {}
        self.itos = {}
        if is_load:
            self.stoi = _TOKENIZER_
            self.itos = {k: v for v, k in self.stoi.items()}

    def __len__(self):
        return len(self.stoi)

    def build_vocab(self, text):
        vocab = set()
        for t in text:
            vocab.update(t.split(" "))
        vocab = sorted(vocab)
        vocab.append("<sos>")
        vocab.append("<eos>")
        vocab.append("<pad>")
        for i, s in enumerate(vocab):
            self.stoi[s] = i
        self.itos = {k: v for v, k in self.stoi.items()}

    def one_text_to_sequence(self, text):
        sequence = []
        sequence.append(self.stoi["<sos>"])
        for s in text.split(" "):
            sequence.append(self.stoi[s])
        sequence.append(self.stoi["<eos>"])
        return sequence

    def one_sequence_to_text(self, sequence):
        return "".join(list(map(lambda i: self.itos[i], sequence)))

    def one_predict_to_inchi(self, predict):
        inchi = "InChI=1S/"
        for p in predict:
            if p == self.stoi["<eos>"] or p == self.stoi["<pad>"]:
                break
            inchi += self.itos[p]
        return inchi

    # ---
    def text_to_sequence(self, text):
        sequence = [self.one_text_to_sequence(t) for t in text]
        return sequence

    def sequence_to_text(self, sequence):
        text = [self.one_sequence_to_text(s) for s in sequence]
        return text

    def predict_to_inchi(self, predict):
        inchi = [self.one_predict_to_inchi(p) for p in predict]
        return inchi


def null_collate(batch, is_sort_decreasing_length=True):
    collate = defaultdict(list)
    if is_sort_decreasing_length:  # sort by decreasing length
        sort = np.argsort([-len(r["token"]) for r in batch])
        batch = [batch[s] for s in sort]
    for r in batch:
        for k, v in r.items():
            collate[k].append(v)
    # ----
    batch_size = len(batch)
    collate["length"] = [len(l) for l in collate["token"]]
    token = [np.array(t, np.int32) for t in collate["token"]]
    token = pad_sequence_to_max_length(
        token, max_length=max_length, padding_value=STOI["<pad>"]
    )
    collate["token"] = torch.from_numpy(token).long()
    max_of_length = max(collate["length"])
    token_pad_mask = np.zeros((batch_size, max_of_length, max_of_length))
    for b in range(batch_size):
        L = collate["length"][b]
        token_pad_mask[b, :L, :L] = 1  # +1 for cls_token
    collate["token_pad_mask"] = torch.from_numpy(token_pad_mask).byte()
    # -----
    # image = np.stack(collate['image'])
    # image = image.astype(np.float32) / 255
    # collate['image'] = torch.from_numpy(image).unsqueeze(1).repeat(1,3,1,1)
    # -----
    collate["num_patch"] = [len(l) for l in collate["patch"]]
    max_of_num_patch = max(collate["num_patch"])
    patch_pad_mask = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))
    patch = np.full(
        (
            batch_size,
            max_of_num_patch,
            patch_size + 2 * pixel_pad,
            patch_size + 2 * pixel_pad,
        ),
        255,
    )  # pad as 255
    coord = np.zeros((batch_size, max_of_num_patch, 2))
    for b in range(batch_size):
        N = collate["num_patch"][b]
        patch[b, :N] = collate["patch"][b]
        coord[b, :N] = collate["coord"][b]
        patch_pad_mask[b, :N, :N] = 1  # +1 for cls_token
    collate["patch"] = torch.from_numpy(patch).half() / 255
    collate["coord"] = torch.from_numpy(coord).long()
    collate["patch_pad_mask"] = torch.from_numpy(patch_pad_mask).byte()
    del patch_pad_mask, patch, coord, token_pad_mask, batch
    return collate


def run_check_dataset():
    tokenizer = load_tokenizer()
    df_train, df_valid = make_fold("train-1")
    # df_train = make_fold('test') #1616107
    # dataset = BmsDataset(df_train, tokenizer, remote_augment)
    # dataset = BmsDataset(df_valid, tokenizer)
    # print(dataset)
    train_dataset = BmsDataset(df_train, tokenizer)
    # for i in range(len(dataset)):
    """
    for i in range(5):
        #i = np.random.choice(len(dataset))
        r = dataset[i]
        print(r['index'])
        print(r['image_id'])
        #print(r['formula'])
        print(r['InChI'])
        print(r['token'])
        print('image : ')
        #print('\t', r['image'].shape)
        print('')
        #---
        image = patch_to_image(r['patch'], r['coord'], width=1024, height=1024)
        #image_show('image', image, resize=.1)
        del image
        gc.collect()
        #cv2.waitKey(0)
    """
    # exit(0)
    # """
    loader = DataLoader(
        train_dataset,
        # sampler = RandomSampler(train_dataset),
        batch_size=8,
        drop_last=True,
        num_workers=8,
        pin_memory=True,
        collate_fn=null_collate,
    )
    for t, batch in enumerate(loader):
        # if t>30: break
        print(t, "-----------")
        # print('index : ', batch['index'])
        """
        print('image : ')
        #print('\t', batch['image'].shape, batch['image'].is_contiguous())
        print('\t', batch['patch'].shape, batch['patch'].is_contiguous())
        print('\t', batch['coord'].shape, batch['coord'].is_contiguous())
        #print('\t', batch['mask'].shape, batch['mask'].is_contiguous())
        print('length  : ')
        print('\t',len( batch['length']))
        print('\t', batch['length'])
        print('token  : ')
        print('\t', batch['token'].shape, batch['token'].is_contiguous())
        """
        # print('\t', batch['token'])
        # print('')
        del t, batch
    del loader, train_dataset, df_train, df_valid, tokenizer
    # """


# 数据load到4000 条 左右就会挂 ， 讲道理 load数据应该不会不断加内存吧
# run_check_dataset()


# # Net Modeules
# ## 1D Encoder and Decoder
from typing import Tuple, Dict
import torch.nn as nn
from fairseq import utils
from fairseq.models import *
from fairseq.modules import *


# ------------------------------------------------------
# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need
class PositionEncode1D(nn.Module):
    def __init__(self, dim, max_length):
        super().__init__()
        assert dim % 2 == 0
        self.max_length = max_length
        d = torch.exp(torch.arange(0.0, dim, 2) * (-math.log(10000.0) / dim))
        position = torch.arange(0.0, max_length).unsqueeze(1)
        pos = torch.zeros(1, max_length, dim)
        pos[0, :, 0::2] = torch.sin(position * d)
        pos[0, :, 1::2] = torch.cos(position * d)
        self.register_buffer("pos", pos)

    def forward(self, x):
        batch_size, T, dim = x.shape
        x = x + self.pos[:, :T]
        return x


"""
# https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py
#https://github.com/pytorch/fairseq/issues/568
# fairseq/fairseq/models/fairseq_encoder.py
# https://github.com/pytorch/fairseq/blob/master/fairseq/modules/transformer_layer.py
class TransformerEncode(FairseqEncoder):
    def __init__(self, dim, ff_dim, num_head, num_layer):
        super().__init__({})
        #print('my TransformerEncode()')
        self.layer = nn.ModuleList([
            TransformerEncoderLayer(Namespace({
                'encoder_embed_dim': dim,
                'encoder_attention_heads': num_head,
                'attention_dropout': 0.1,
                'dropout': 0.1,
                'encoder_normalize_before': True,
                'encoder_ffn_embed_dim': ff_dim,
            })) for i in range(num_layer)
        ])
        self.layer_norm = nn.LayerNorm(dim)
    def forward(self, x):# T x B x C
        #print('my TransformerEncode forward()')
        for layer in self.layer:
            x = layer(x)
        x = self.layer_norm(x)
        return x
"""


# https://mt.cs.upc.edu/2020/12/21/the-transformer-fairseq-edition/
# for debug
# class TransformerDecode(FairseqDecoder):
#     def __init__(self, dim, ff_dim, num_head, num_layer):
#         super().__init__({})
#         print('my TransformerDecode()')
#
#         self.layer = nn.ModuleList([
#             TransformerDecoderLayer(Namespace({
#                 'decoder_embed_dim': dim,
#                 'decoder_attention_heads': num_head,
#                 'attention_dropout': 0.1,
#                 'dropout': 0.1,
#                 'decoder_normalize_before': True,
#                 'decoder_ffn_embed_dim': ff_dim,
#             })) for i in range(num_layer)
#         ])
#         self.layer_norm = nn.LayerNorm(dim)
#
#
#     def forward(self, x, mem, x_mask):# T x B x C
#         print('my TransformerDecode forward()')
#         for layer in self.layer:
#             x = layer(x, mem, self_attn_mask=x_mask)[0]
#         x = self.layer_norm(x)
#         return x  # T x B x C
# https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute
class Namespace(object):
    def __init__(self, adict):
        self.__dict__.update(adict)


# https://fairseq.readthedocs.io/en/latest/tutorial_simple_lstm.html
# see https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py
class TransformerDecode(FairseqIncrementalDecoder):
    def __init__(self, dim, ff_dim, num_head, num_layer):
        super().__init__({})
        # print('my TransformerDecode()')
        self.layer = nn.ModuleList(
            [
                TransformerDecoderLayer(
                    Namespace(
                        {
                            "decoder_embed_dim": dim,
                            "decoder_attention_heads": num_head,
                            "attention_dropout": 0.1,
                            "dropout": 0.1,
                            "decoder_normalize_before": True,
                            "decoder_ffn_embed_dim": ff_dim,
                        }
                    )
                )
                for i in range(num_layer)
            ]
        )
        self.layer_norm = nn.LayerNorm(dim)

    def forward(self, x, mem, x_mask, x_pad_mask, mem_pad_mask):
        # print('my TransformerDecode forward()')
        for layer in self.layer:
            x = layer(
                x,
                mem,
                self_attn_mask=x_mask,
                self_attn_padding_mask=x_pad_mask,
                encoder_padding_mask=mem_pad_mask,
            )[0]
        x = self.layer_norm(x)
        return x  # T x B x C

    # def forward_one(self, x, mem, incremental_state):
    def forward_one(
        self,
        x: Tensor,
        mem: Tensor,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
    ) -> Tensor:
        x = x[-1:]
        for layer in self.layer:
            x = layer(x, mem, incremental_state=incremental_state)[0]
        x = self.layer_norm(x)
        return x


"""
https://fairseq.readthedocs.io/en/latest/_modules/fairseq/modules/transformer_layer.html
class TransformerDecoderLayer(nn.Module):
def forward(
        self,
        x,
        encoder_out: Optional[torch.Tensor] = None,
        encoder_padding_mask: Optional[torch.Tensor] = None,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
        prev_self_attn_state: Optional[List[torch.Tensor]] = None,
        prev_attn_state: Optional[List[torch.Tensor]] = None,
        self_attn_mask: Optional[torch.Tensor] = None,
        self_attn_padding_mask: Optional[torch.Tensor] = None,
        need_attn: bool = False,
        need_head_weights: bool = False,
    ):
https://github.com/pytorch/fairseq/blob/05b86005bcca0155319fa9b81abfd69f63c06906/examples/simultaneous_translation/utils/data_utils.py#L31      
"""


# ## TNT
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence

# https://github.com/pytorch/pytorch/issues/1788
# https://stackoverflow.com/questions/51030782/why-do-we-pack-the-sequences-in-pytorch
from timm.models.vision_transformer import Mlp
from timm.models.layers import DropPath, trunc_normal_


# from timm.models.tnt import *
# 这个多头注意力会造成内存不断增加吗？
class Attention(nn.Module):
    """Multi-Head Attention"""

    def __init__(
        self, dim, hidden_dim, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        head_dim = hidden_dim // num_heads
        self.head_dim = head_dim
        self.scale = head_dim**-0.5
        self.qk = nn.Linear(dim, hidden_dim * 2, bias=qkv_bias)
        self.v = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop, inplace=True)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop, inplace=True)

    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        B, N, C = x.shape
        qk = (
            self.qk(x)
            .reshape(B, N, 2, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k = qk[0], qk[1]  # make torchscript happy (cannot use tensor as tuple)
        v = self.v(x).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        # ---
        attn = (q @ k.transpose(-2, -1)) * self.scale  # B x self.num_heads x NxN
        if mask is not None:
            # mask = mask.unsqueeze(1).repeat(1,self.num_heads,1,1)
            mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
            attn = attn.masked_fill(mask == 0, -6e4)
            # attn = attn.masked_fill(mask == 0, -half('inf'))
            # https://github.com/NVIDIA/apex/issues/93
            # How to use fp16 training with masked operations
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
        x = self.proj(x)
        x = self.proj_drop(x)
        del mask, attn
        return x


class Block(nn.Module):
    def __init__(
        self,
        dim,
        in_dim,
        num_pixel,
        num_heads=12,
        in_num_head=4,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        # Inner transformer
        self.norm_in = norm_layer(in_dim)
        self.attn_in = Attention(
            in_dim,
            in_dim,
            num_heads=in_num_head,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.norm_mlp_in = norm_layer(in_dim)
        self.mlp_in = Mlp(
            in_features=in_dim,
            hidden_features=int(in_dim * 4),
            out_features=in_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm1_proj = norm_layer(in_dim)
        self.proj = nn.Linear(in_dim * num_pixel, dim, bias=True)
        # Outer transformer
        self.norm_out = norm_layer(dim)
        self.attn_out = Attention(
            dim,
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm_mlp = norm_layer(dim)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),
            out_features=dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, pixel_embed, patch_embed, mask):
        # inner
        pixel_embed = pixel_embed + self.drop_path(
            self.attn_in(self.norm_in(pixel_embed))
        )
        pixel_embed = pixel_embed + self.drop_path(
            self.mlp_in(self.norm_mlp_in(pixel_embed))
        )
        # outer
        B, N, C = patch_embed.size()
        patch_embed[:, 1:] = (
            patch_embed[:, 1:]
            + self.proj(self.norm1_proj(pixel_embed).reshape(B, N, -1))[:, 1:]
        )
        patch_embed = patch_embed + self.drop_path(
            self.attn_out(self.norm_out(patch_embed), mask)
        )
        patch_embed = patch_embed + self.drop_path(self.mlp(self.norm_mlp(patch_embed)))
        return pixel_embed, patch_embed


# ---------------------------------
class PixelEmbed(nn.Module):
    def __init__(self, patch_size=16, in_dim=48, stride=4):
        super().__init__()
        self.in_dim = in_dim
        self.proj = nn.Conv2d(3, self.in_dim, kernel_size=7, padding=0, stride=stride)

    def forward(self, patch, pixel_pos):
        BN = len(patch)
        x = patch
        x = self.proj(x)
        # x = x.transpose(1, 2).reshape(B * self.num_patches, self.in_dim, self.new_patch_size, self.new_patch_size)
        x = x + pixel_pos
        x = x.reshape(BN, self.in_dim, -1).transpose(1, 2)
        return x


# ---------------------------------
class TNT(nn.Module):
    """Transformer in Transformer - https://arxiv.org/abs/2103.00112"""

    def __init__(
        self,
        patch_size=patch_size,
        embed_dim=patch_dim,
        in_dim=pixel_dim,
        depth=12,
        num_heads=6,
        in_num_head=4,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        norm_layer=nn.LayerNorm,
        first_stride=pixel_stride,
    ):
        super().__init__()
        self.num_features = (
            self.embed_dim
        ) = embed_dim  # num_features for consistency with other models
        self.pixel_embed = PixelEmbed(
            patch_size=patch_size, in_dim=in_dim, stride=first_stride
        )
        # num_patches = self.pixel_embed.num_patches
        # self.num_patches = num_patches
        new_patch_size = 4  # self.pixel_embed.new_patch_size
        num_pixel = new_patch_size**2
        self.norm1_proj = norm_layer(num_pixel * in_dim)
        self.proj = nn.Linear(num_pixel * in_dim, embed_dim)
        self.norm2_proj = norm_layer(embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.patch_pos = nn.Embedding(
            100 * 100, embed_dim
        )  # nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pixel_pos = nn.Parameter(
            torch.zeros(1, in_dim, new_patch_size, new_patch_size)
        )
        self.pos_drop = nn.Dropout(p=drop_rate)
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, depth)
        ]  # stochastic depth decay rule
        blocks = []
        for i in range(depth):
            blocks.append(
                Block(
                    dim=embed_dim,
                    in_dim=in_dim,
                    num_pixel=num_pixel,
                    num_heads=num_heads,
                    in_num_head=in_num_head,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[i],
                    norm_layer=norm_layer,
                )
            )
        self.blocks = nn.ModuleList(blocks)
        self.norm = norm_layer(embed_dim)
        trunc_normal_(self.cls_token, std=0.02)
        # trunc_normal_(self.patch_pos, std=.02)
        trunc_normal_(self.pixel_pos, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {"patch_pos", "pixel_pos", "cls_token"}

    def forward(self, patch, coord, mask):
        B = len(patch)
        batch_size, max_of_num_patch, s, s = patch.shape
        patch = patch.reshape(batch_size * max_of_num_patch, 1, s, s).repeat(1, 3, 1, 1)
        pixel_embed = self.pixel_embed(patch, self.pixel_pos)
        patch_embed = self.norm2_proj(
            self.proj(self.norm1_proj(pixel_embed.reshape(B, max_of_num_patch, -1)))
        )
        # patch_embed = torch.cat((self.cls_token.expand(B, -1, -1), patch_embed), dim=1)
        # patch_embed = patch_embed + self.patch_pos
        # patch_embed[:, 1:] = patch_embed[:, 1:] + self.patch_pos(coord[:, :, 0] * 100 + coord[:, :, 1])
        patch_embed[:, :1] = self.cls_token.expand(B, -1, -1)
        patch_embed = patch_embed + self.patch_pos(
            coord[:, :, 0] * 100 + coord[:, :, 1]
        )
        patch_embed = self.pos_drop(patch_embed)
        for blk in self.blocks:
            pixel_embed, patch_embed = blk(pixel_embed, patch_embed, mask)
        patch_embed = self.norm(patch_embed)
        del patch, mask
        return patch_embed


def make_dummy_data():
    # make dummy data
    # image_id,width,height,scale,orientation
    meta = [
        [
            "000011a64c74",
            325,
            229,
            2,
            0,
        ],
        [
            "000019cc0cd2",
            288,
            148,
            1,
            0,
        ],
        [
            "0000252b6d2b",
            509,
            335,
            2,
            0,
        ],
        [
            "000026b49b7e",
            243,
            177,
            1,
            0,
        ],
        [
            "000026fc6c36",
            294,
            112,
            1,
            0,
        ],
        [
            "000028818203",
            402,
            328,
            2,
            0,
        ],
        [
            "000029a61c01",
            395,
            294,
            2,
            0,
        ],
        [
            "000035624718",
            309,
            145,
            1,
            0,
        ],
    ]
    batch_size = 8
    # <todo> check border for padding
    # <todo> pepper noise
    batch = {
        "num_patch": [],
        "patch": [],
        "coord": [],
    }
    for b in range(batch_size):
        image_id = meta[b][0]
        scale = meta[b][3]
        image_file = origin_data_dir + "/%s/%s/%s/%s/%s.png" % (
            "train",
            image_id[0],
            image_id[1],
            image_id[2],
            image_id,
        )
        # print(image_file)
        image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)
        image = resize_image(image, scale)
        image = repad_image(image, patch_size)  # remove border and repad
        # print(image.shape)
        k, yx = image_to_patch(image, patch_size, pixel_pad, threshold=0)
        for y, x in yx:
            # cv2.circle(image,(x,y),8,128,1)
            x = x * patch_size
            y = y * patch_size
            cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)
        # image_show('image-%d' % b, image, resize=1)
        del image
        # cv2.waitKey(1)
        batch["patch"].append(k)
        batch["coord"].append(yx)
        batch["num_patch"].append(len(k))
    # ----
    max_of_num_patch = max(batch["num_patch"])
    mask = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))
    patch = np.zeros(
        (
            batch_size,
            max_of_num_patch,
            patch_size + 2 * pixel_pad,
            patch_size + 2 * pixel_pad,
        )
    )
    coord = np.zeros((batch_size, max_of_num_patch, 2))
    for b in range(batch_size):
        N = batch["num_patch"][b]
        patch[b, :N] = batch["patch"][b]
        coord[b, :N] = batch["coord"][b]
        mask[b, :N, :N] = 1
    num_patch = batch["num_patch"]
    patch = torch.from_numpy(patch).float()
    coord = torch.from_numpy(coord).long()
    mask = torch.from_numpy(mask).byte()
    return patch, coord, num_patch, mask


def run_check_tnt_patch():
    patch, coord, num_patch, mask = make_dummy_data()
    tnt = TNT()
    patch_embed = tnt(patch, coord, mask)
    print(patch_embed.shape)


# run_check_tnt_patch()


# ## TNT + PositionEncode1D + Transformer Decoder
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence

vocab_size = 193
text_dim = 384
decoder_dim = 384
num_layer = 3
num_head = 8
ff_dim = 1024


class Net(nn.Module):
    def __init__(
        self,
    ):
        super(Net, self).__init__()
        self.cnn = TNT()
        self.image_encode = nn.Identity()
        # ---
        self.text_pos = PositionEncode1D(text_dim, max_length)
        self.token_embed = nn.Embedding(vocab_size, text_dim)
        self.text_decode = TransformerDecode(decoder_dim, ff_dim, num_head, num_layer)
        # ---
        self.logit = nn.Linear(decoder_dim, vocab_size)
        self.dropout = nn.Dropout(p=0.5)
        # ----
        # initialization
        self.token_embed.weight.data.uniform_(-0.1, 0.1)
        self.logit.bias.data.fill_(0)
        self.logit.weight.data.uniform_(-0.1, 0.1)

    @torch.jit.unused
    def forward(self, patch, coord, token, patch_pad_mask, token_pad_mask):
        device = patch.device
        batch_size = len(patch)
        # ---
        patch = patch * 2 - 1
        image_embed = self.cnn(patch, coord, patch_pad_mask)
        image_embed = self.image_encode(image_embed).permute(1, 0, 2).contiguous()
        text_embed = self.token_embed(token)
        text_embed = self.text_pos(text_embed).permute(1, 0, 2).contiguous()
        max_of_length = token_pad_mask.shape[-1]
        text_mask = np.triu(np.ones((max_of_length, max_of_length)), k=1).astype(
            np.uint8
        )
        text_mask = torch.autograd.Variable(torch.from_numpy(text_mask) == 1).to(device)
        # ----
        # <todo> perturb mask as aug
        text_pad_mask = token_pad_mask[:, :, 0] == 0
        image_pad_mask = patch_pad_mask[:, :, 0] == 0
        x = self.text_decode(
            text_embed[:max_of_length],
            image_embed,
            text_mask,
            text_pad_mask,
            image_pad_mask,
        )
        x = x.permute(1, 0, 2).contiguous()
        l = self.logit(x)
        logit = torch.zeros((batch_size, max_length, vocab_size), device=device)
        logit[:, :max_of_length] = l
        del image_embed, text_mask, text_pad_mask, image_pad_mask
        return logit

    # submit function has not been coded. i will leave it as an exercise for the kaggler
    # @torch.jit.export
    # def forward_argmax_decode(self, patch, coord, mask):
    #
    #     image_dim   = 384
    #     text_dim    = 384
    #     decoder_dim = 384
    #     num_layer = 3
    #     num_head  = 8
    #     ff_dim    = 1024
    #
    #     STOI = {
    #         '<sos>': 190,
    #         '<eos>': 191,
    #         '<pad>': 192,
    #     }
    #     max_length = 278 # 275
    #
    #
    #     #---------------------------------
    #     device = patch.device
    #     batch_size = len(patch)
    #
    #     patch = patch*2-1
    #     image_embed = self.cnn(patch, coord, mask)
    #     image_embed = self.image_encode(image_embed).permute(1,0,2).contiguous()
    #
    #     token = torch.full((batch_size, max_length), STOI['<pad>'],dtype=torch.long, device=device)
    #     text_pos = self.text_pos.pos
    #     token[:,0] = STOI['<sos>']
    #
    #
    #     #-------------------------------------
    #     eos = STOI['<eos>']
    #     pad = STOI['<pad>']
    #     # https://github.com/alexmt-scale/causal-transformer-decoder/blob/master/tests/test_consistency.py
    #     # slow version
    #     # if 0:
    #     #     for t in range(max_length-1):
    #     #         last_token = token [:,:(t+1)]
    #     #         text_embed = self.token_embed(last_token)
    #     #         text_embed = self.text_pos(text_embed).permute(1,0,2).contiguous() #text_embed + text_pos[:,:(t+1)] #
    #     #
    #     #         text_mask = np.triu(np.ones((t+1, t+1)), k=1).astype(np.uint8)
    #     #         text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)==1).to(device)
    #     #
    #     #         x = self.text_decode(text_embed, image_embed, text_mask)
    #     #         x = x.permute(1,0,2).contiguous()
    #     #
    #     #         l = self.logit(x[:,-1])
    #     #         k = torch.argmax(l, -1)  # predict max
    #     #         token[:, t+1] = k
    #     #         if ((k == eos) | (k == pad)).all():  break
    #
    #     # fast version
    #     if 1:
    #         #incremental_state = {}
    #         incremental_state = torch.jit.annotate(
    #             Dict[str, Dict[str, Optional[Tensor]]],
    #             torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {}),
    #         )
    #         for t in range(max_length-1):
    #             #last_token = token [:,:(t+1)]
    #             #text_embed = self.token_embed(last_token)
    #             #text_embed = self.text_pos(text_embed) #text_embed + text_pos[:,:(t+1)] #
    #
    #             last_token = token[:, t]
    #             text_embed = self.token_embed(last_token)
    #             text_embed = text_embed + text_pos[:,t] #
    #             text_embed = text_embed.reshape(1,batch_size,text_dim)
    #
    #             x = self.text_decode.forward_one(text_embed, image_embed, incremental_state)
    #             x = x.reshape(batch_size,decoder_dim)
    #             #print(incremental_state.keys())
    #
    #             l = self.logit(x)
    #             k = torch.argmax(l, -1)  # predict max
    #             token[:, t+1] = k
    #             if ((k == eos) | (k == pad)).all():  break
    #
    #     predict = token[:, 1:]
    #     return predict


# loss #################################################################
def seq_cross_entropy_loss(logit, token, length):
    truth = token[:, 1:]
    L = [l - 1 for l in length]
    logit = pack_padded_sequence(logit, L, batch_first=True).data
    truth = pack_padded_sequence(truth, L, batch_first=True).data
    loss = F.cross_entropy(logit, truth, ignore_index=STOI["<pad>"])
    return loss


# https://www.aclweb.org/anthology/2020.findings-emnlp.276.pdf
def seq_focal_cross_entropy_loss(logit, token, length):
    gamma = 0.5  # {0.5,1.0}
    # label_smooth = 0.90
    # ---
    truth = token[:, 1:]
    L = [l - 1 for l in length]
    logit = pack_padded_sequence(logit, L, batch_first=True).data
    truth = pack_padded_sequence(truth, L, batch_first=True).data
    # loss = F.cross_entropy(logit, truth, ignore_index=STOI['<pad>'])
    # non_pad = torch.where(truth != STOI['<pad>'])[0]  # & (t!=STOI['<sos>'])
    # ---
    # p = F.softmax(logit,-1)
    # logp = - torch.log(torch.clamp(p, 1e-4, 1 - 1e-4))
    logp = F.log_softmax(logit, -1)
    logp = logp.gather(1, truth.reshape(-1, 1)).reshape(-1)
    p = logp.exp()
    loss = -((1 - p) ** gamma) * logp  # focal
    # loss = - ((1 + p) ** gamma)*logp  #anti-focal
    loss = loss.mean()
    return loss


def np_loss_cross_entropy(probability, truth):
    batch_size = len(probability)
    truth = truth.reshape(-1)
    p = probability[np.arange(batch_size), truth]
    loss = -np.log(np.clip(p, 1e-6, 1))
    loss = loss.mean()
    return loss


# check #################################################################
def run_check_net():
    patch, coord, num_patch, patch_pad_mask = make_dummy_data()
    batch_size = len(patch)
    token = np.full((batch_size, max_length), STOI["<pad>"], np.int64)  # token
    length = np.random.randint(5, max_length - 2, batch_size)
    length = np.sort(length)[::-1].copy()
    max_of_length = max(length)
    token_pad_mask = np.zeros((batch_size, max_of_length, max_of_length))
    for b in range(batch_size):
        l = length[b]
        t = np.random.choice(vocab_size, l)
        t = np.insert(t, 0, STOI["<sos>"])
        t = np.insert(t, len(t), STOI["<eos>"])
        L = len(t)
        token[b, :L] = t
        token_pad_mask[b, :L, :L] = 1
    token = torch.from_numpy(token).long()
    token_pad_mask = torch.from_numpy(token_pad_mask).byte()
    # ---
    net = Net()
    net.train()
    logit = net(patch, coord, token, patch_pad_mask, token_pad_mask)
    loss = seq_cross_entropy_loss(logit, token, length)
    print("vocab_size", vocab_size)
    print("max_length", max_length)
    print("")
    print(length)
    print(length.shape)
    print(token.shape)
    print("---")
    print(logit.shape)
    print(loss)
    print("---")
    del net
    # ---
    # print('torch.jit.script(net)')
    # net.eval()
    # net = torch.jit.script(net)
    #
    # predict = net.forward_argmax_decode(patch, coord, mask)
    # print(predict.shape)


# run_check_net()
"""
import sys
local_vars = list(locals().items())
for var, obj in local_vars:
    print(var, sys.getsizeof(obj))
"""
# ----------------
is_mixed_precision = True  # False  #
###################################################################################################
import torch.cuda.amp as amp

if is_mixed_precision:

    class AmpNet(Net):
        @torch.cuda.amp.autocast()
        def forward(self, *args):
            return super(AmpNet, self).forward(*args)

else:
    AmpNet = Net
import sys
from datetime import datetime
import Levenshtein

IDENTIFIER = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")


# http://stackoverflow.com/questions/34950201/pycharm-print-end-r-statement-not-working
class Logger(object):
    def __init__(self):
        self.terminal = sys.stdout  # stdout
        self.file = None

    def open(self, file, mode=None):
        if mode is None:
            mode = "w"
        self.file = open(file, mode)

    def write(self, message, is_terminal=1, is_file=1):
        if "\r" in message:
            is_file = 0
        if is_terminal == 1:
            self.terminal.write(message)
            self.terminal.flush()
            # time.sleep(1)
        if is_file == 1:
            self.file.write(message)
            self.file.flush()

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass


# etc ------------------------------------
def time_to_str(t, mode="min"):
    if mode == "min":
        t = int(t) / 60
        hr = t // 60
        min = t % 60
        return "%2d hr %02d min" % (hr, min)
    elif mode == "sec":
        t = int(t)
        min = t // 60
        sec = t % 60
        return "%2d min %02d sec" % (min, sec)
    else:
        raise NotImplementedError


# ## Optimizer
from torch.optim.optimizer import Optimizer
import itertools as it


def get_learning_rate(optimizer):
    lr = []
    for param_group in optimizer.param_groups:
        lr += [param_group["lr"]]
    assert len(lr) == 1  # we support only one param_group
    lr = lr[0]
    return lr


class RAdam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        self.buffer = [[None, None, None] for ind in range(10)]
        super(RAdam, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(RAdam, self).__setstate__(state)

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError("RAdam does not support sparse gradients")
                p_data_fp32 = p.data.float()
                state = self.state[p]
                if len(state) == 0:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
                else:
                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                beta1, beta2 = group["betas"]
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                state["step"] += 1
                buffered = self.buffer[int(state["step"] % 10)]
                if state["step"] == buffered[0]:
                    N_sma, step_size = buffered[1], buffered[2]
                else:
                    buffered[0] = state["step"]
                    beta2_t = beta2 ** state["step"]
                    N_sma_max = 2 / (1 - beta2) - 1
                    N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
                    buffered[1] = N_sma
                    # more conservative since it's an approximated value
                    if N_sma >= 5:
                        step_size = math.sqrt(
                            (1 - beta2_t)
                            * (N_sma - 4)
                            / (N_sma_max - 4)
                            * (N_sma - 2)
                            / N_sma
                            * N_sma_max
                            / (N_sma_max - 2)
                        ) / (1 - beta1 ** state["step"])
                    else:
                        step_size = 1.0 / (1 - beta1 ** state["step"])
                    buffered[2] = step_size
                if group["weight_decay"] != 0:
                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
                # more conservative since it's an approximated value
                if N_sma >= 5:
                    denom = exp_avg_sq.sqrt().add_(group["eps"])
                    p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size * group["lr"])
                else:
                    p_data_fp32.add_(exp_avg, alpha=-step_size * group["lr"])
                p.data.copy_(p_data_fp32)
        return loss


class Lookahead(Optimizer):
    def __init__(self, optimizer, alpha=0.5, k=6):
        if not 0.0 <= alpha <= 1.0:
            raise ValueError(f"Invalid slow update rate: {alpha}")
        if not 1 <= k:
            raise ValueError(f"Invalid lookahead steps: {k}")
        self.optimizer = optimizer
        self.param_groups = self.optimizer.param_groups
        self.alpha = alpha
        self.k = k
        for group in self.param_groups:
            group["step_counter"] = 0
        self.slow_weights = [
            [p.clone().detach() for p in group["params"]] for group in self.param_groups
        ]
        for w in it.chain(*self.slow_weights):
            w.requires_grad = False
        self.state = optimizer.state

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
        loss = self.optimizer.step()
        for group, slow_weights in zip(self.param_groups, self.slow_weights):
            group["step_counter"] += 1
            if group["step_counter"] % self.k != 0:
                continue
            for p, q in zip(group["params"], slow_weights):
                if p.grad is None:
                    continue
                q.data.add_(p.data - q.data, alpha=self.alpha)
                p.data.copy_(q.data)
        return loss


# ## Training
def do_valid(net, tokenizer, valid_loader):
    valid_probability = []
    valid_truth = []
    valid_length = []
    valid_num = 0
    net.eval()
    start_timer = timer()
    for t, batch in enumerate(valid_loader):
        batch_size = len(batch["index"])
        length = batch["length"]
        token = batch["token"].cuda()
        token_pad_mask = batch["token_pad_mask"].cuda()
        # image  = batch['image' ].cuda()
        num_patch = batch["num_patch"]
        patch = batch["patch"].cuda()
        coord = batch["coord"].cuda()
        patch_pad_mask = batch["patch_pad_mask"].cuda()
        with torch.no_grad():
            logit = data_parallel(
                net, (patch, coord, token, patch_pad_mask, token_pad_mask)
            )  # net(image, token, length)
            probability = F.softmax(logit, -1)
        valid_num += batch_size
        valid_probability.append(probability.data.cpu().numpy())
        valid_truth.append(token.data.cpu().numpy())
        valid_length.extend(length)
        print(
            "\r %8d / %d  %s"
            % (
                valid_num,
                len(valid_loader.sampler),
                time_to_str(timer() - start_timer, "sec"),
            ),
            end="",
            flush=True,
        )
        del t, batch
    assert valid_num == len(valid_loader.sampler)  # len(valid_loader.dataset))
    # print('')
    # ----------------------
    probability = np.concatenate(valid_probability)
    predict = probability.argmax(-1)
    truth = np.concatenate(valid_truth)
    length = valid_length
    # ----
    p = probability[:, :-1].reshape(-1, vocab_size)
    t = truth[:, 1:].reshape(-1)
    non_pad = np.where(t != STOI["<pad>"])[0]  # & (t!=STOI['<sos>'])
    p = p[non_pad]
    t = t[non_pad]
    loss = np_loss_cross_entropy(p, t)
    # ----
    lb_score = 0
    if 1:
        score = []
        for i, (p, t) in enumerate(zip(predict, truth)):
            t = truth[i][1 : length[i] - 1]
            p = predict[i][1 : length[i] - 1]
            t = tokenizer.one_predict_to_inchi(t)
            p = tokenizer.one_predict_to_inchi(p)
            s = Levenshtein.distance(p, t)
            score.append(s)
        lb_score = np.mean(score)
    """
    if 1:
        score = []
        for i, (p, t) in enumerate(zip(predict, truth)):
            t = truth[i][1:length[i]-1]     # in the buggy version, i have used 1 instead of i
            p = predict[i][1:length[i]-1]
            t = tokenizer.one_predict_to_inchi(t)
            p = tokenizer.one_predict_to_inchi(p)
            s = Levenshtein.distance(p, t)
            score.append(s)
        lb_score = np.mean(score)
    """
    # lb_score = compute_lb_score(k, t)
    del (
        valid_loader,
        net,
        predict,
        truth,
        valid_probability,
        valid_truth,
        valid_length,
        valid_num,
    )
    return [loss, lb_score]


def run_train():
    fold = 3
    out_dir = "./tnt-patch1-s0.8/fold%d" % fold
    initial_checkpoint = None
    # initial_checkpoint = \
    # out_dir + '/checkpoint/00755000_model.pth'#None #
    #'/root/share1/kaggle/2021/bms-moleular-translation/result/try22/tnt-patch1/fold3/checkpoint/00697000_model.pth'
    debug = 0
    start_lr = 0.00001  # 1
    batch_size = 32  # 24
    ## setup  ----------------------------------------
    for f in ["checkpoint", "train", "valid", "backup"]:
        os.makedirs(out_dir + "/" + f, exist_ok=True)
    # backup_project_as_zip(PROJECT_PATH, out_dir +'/backup/code.train.%s.zip'%IDENTIFIER)
    log = Logger()
    log.open(out_dir + "/log.train.txt", mode="a")
    log.write("\n--- [START %s] %s\n\n" % (IDENTIFIER, "-" * 64))
    log.write("\t%s\n" % COMMON_STRING)
    # log.write('\t__file__ = %s\n' % __file__)
    log.write("\tout_dir  = %s\n" % out_dir)
    log.write("\n")
    ## dataset ------------------------------------
    df_train, df_valid = make_fold("train-%d" % fold)
    df_valid = df_valid.iloc[:5_000]
    tokenizer = load_tokenizer()
    train_dataset = BmsDataset(df_train, tokenizer)
    valid_dataset = BmsDataset(df_valid, tokenizer)
    train_loader = DataLoader(
        train_dataset,
        sampler=RandomSampler(train_dataset),
        # sampler=UniformLengthSampler(train_dataset, is_shuffle=True), #200_000
        batch_size=batch_size,
        drop_last=True,
        num_workers=8,
        pin_memory=True,
        worker_init_fn=lambda id: np.random.seed(torch.initial_seed() // 2**32 + id),
        collate_fn=null_collate,
    )
    valid_loader = DataLoader(
        valid_dataset,
        # sampler=UniformLengthSampler(valid_dataset, 5_000),
        sampler=SequentialSampler(valid_dataset),
        batch_size=32,
        drop_last=False,
        num_workers=8,
        pin_memory=True,
        collate_fn=null_collate,
    )
    log.write("train_dataset : \n%s\n" % (train_dataset))
    log.write("valid_dataset : \n%s\n" % (valid_dataset))
    log.write("\n")
    ## net ----------------------------------------
    log.write("** net setting **\n")
    scaler = amp.GradScaler()
    net = AmpNet().cuda()
    if initial_checkpoint is not None:
        f = torch.load(initial_checkpoint, map_location=lambda storage, loc: storage)
        start_iteration = f["iteration"]
        start_epoch = f["epoch"]
        state_dict = f["state_dict"]
        # ---
        # state_dict = {k.replace('cnn.e.','cnn.'):v for k,v in state_dict.items()}
        # del state_dict['text_pos.pos']
        # del state_dict['cnn.head.weight']
        # del state_dict['cnn.head.bias']
        # net.load_state_dict(state_dict, strict=False)
        # ---
        net.load_state_dict(state_dict, strict=True)  # True
    else:
        start_iteration = 0
        start_epoch = 0
    log.write("\tinitial_checkpoint = %s\n" % initial_checkpoint)
    log.write("\n")
    # -----------------------------------------------
    if 0:  ##freeze
        for p in net.encoder.parameters():
            p.requires_grad = False
    optimizer = Lookahead(
        RAdam(filter(lambda p: p.requires_grad, net.parameters()), lr=start_lr),
        alpha=0.5,
        k=5,
    )
    # optimizer = RAdam(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr)
    num_iteration = 10 * 1000
    # num_iteration = 10
    iter_log = 1000
    iter_valid = 1000
    iter_save = list(range(0, num_iteration, 1000))  # 1*1000
    # iter_save = [0]
    log.write("optimizer\n  %s\n" % (optimizer))
    log.write("\n")
    ## start training here! ##############################################
    log.write("** start training here! **\n")
    log.write("   is_mixed_precision = %s \n" % str(is_mixed_precision))
    log.write("   batch_size = %d\n" % (batch_size))
    # log.write('   experiment = %s\n' % str(__file__.split('/')[-2:]))
    log.write(
        "                      |----- VALID ---|---- TRAIN/BATCH --------------\n"
    )
    log.write(
        "rate     iter   epoch | loss  lb(lev) | loss0  loss1  | time          \n"
    )
    log.write(
        "----------------------------------------------------------------------\n"
    )

    # 0.00000   0.00* 0.00  | 0.000  0.000  | 0.000  0.000  |  0 hr 00 min
    def message(mode="print"):
        if mode == ("print"):
            asterisk = " "
            loss = batch_loss
        if mode == ("log"):
            asterisk = "*" if iteration in iter_save else " "
            loss = train_loss
        #'%4.3f  %5.2f  | ' % (*valid_loss,) + \
        text = (
            "%0.5f  %5.4f%s %4.2f  | "
            % (
                rate,
                iteration / 10000,
                asterisk,
                epoch,
            )
            + "%4.3f  %4.3f  %4.3f  | " % (*loss,)
            + "%s" % (time_to_str(timer() - start_timer, "min"))
        )
        return text

    # ----
    valid_loss = np.zeros(2, np.float32)
    train_loss = np.zeros(3, np.float32)
    batch_loss = np.zeros_like(train_loss)
    sum_train_loss = np.zeros_like(train_loss)
    sum_train = 0
    loss0 = torch.FloatTensor([0]).cuda().sum()
    loss1 = torch.FloatTensor([0]).cuda().sum()
    loss2 = torch.FloatTensor([0]).cuda().sum()
    start_timer = timer()
    iteration = start_iteration
    epoch = start_epoch
    rate = 0
    # while iteration < num_iteration:
    for t, batch in enumerate(train_loader):
        if iteration in iter_save:
            if iteration != start_iteration:
                torch.save(
                    {
                        "state_dict": net.state_dict(),
                        "iteration": iteration,
                        "epoch": epoch,
                    },
                    out_dir + "/checkpoint/%08d_model.pth" % (iteration),
                )
                pass
        """
        if (iteration % iter_valid == 0):
            if iteration != start_iteration:
                valid_loss = do_valid(net, tokenizer, valid_loader)  #
                pass
        """
        if iteration % iter_log == 0:
            print("\r", end="", flush=True)
            log.write(message(mode="log") + "\n")
        # learning rate schduler ------------
        rate = get_learning_rate(optimizer)
        # one iteration update  -------------
        batch_size = len(batch["index"])
        length = batch["length"]
        token = batch["token"].cuda()
        token_pad_mask = batch["token_pad_mask"].cuda()
        # image  = batch['image' ].cuda()
        num_patch = batch["num_patch"]
        patch = batch["patch"].cuda()
        coord = batch["coord"].cuda()
        patch_pad_mask = batch["patch_pad_mask"].cuda()
        # ----
        net.train()
        optimizer.zero_grad()
        if is_mixed_precision:
            with amp.autocast():
                # assert(False)
                logit = data_parallel(
                    net, (patch, coord, token, patch_pad_mask, token_pad_mask)
                )  # net(image, token, length)
                loss0 = seq_cross_entropy_loss(logit, token, length)
                # loss0 = seq_anti_focal_cross_entropy_loss(logit, token, length)
            scaler.scale(loss0).backward()
            # scaler.unscale_(optimizer)
            # torch.nn.utils.clip_grad_norm_(net.parameters(), 2)
            scaler.step(optimizer)
            scaler.update()
        else:
            assert False
            # print('fp32')
            # image_embed = encoder(image)
            logit, weight = decoder(image_embed, token, length)
            (loss0).backward()
            optimizer.step()
        # print statistics  --------
        epoch += 1 / len(train_loader)
        iteration += 1
        batch_loss = np.array([loss0.item(), loss1.item(), loss2.item()])
        sum_train_loss += batch_loss
        sum_train += 1
        if iteration % 100 == 0:
            train_loss = sum_train_loss / (sum_train + 1e-12)
            sum_train_loss[...] = 0
            sum_train = 0
        print("\r", end="", flush=True)
        print(message(mode="print"), end="", flush=True)
        # debug--------------------------
        if debug:
            pass
        # delete per batch
        del (
            batch_size,
            length,
            token,
            token_pad_mask,
            num_patch,
            patch,
            coord,
            patch_pad_mask,
            t,
            batch,
        )
        if iteration > num_iteration:
            break
    log.write("\n")


run_train()
#!ls ./tnt-patch1-s0.8/fold3/checkpoint/
# list(range(0, 3000, 1000))
"""
import sys
# These are the usual ipython objects, including this one you are creating
ipython_vars = ['In', 'Out', 'exit', 'quit', 'get_ipython', 'ipython_vars']
# Get a sorted list of the objects and their sizes
sorted([(x, sys.getsizeof(globals().get(x))) for x in dir() if not x.startswith('_') and x not in sys.modules and x not in ipython_vars], key=lambda x: x[1], reverse=True)
"""
