{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n\ninputDir = '/kaggle/input'\n\ndef _walkdir(dirPath): \n    for dirname, _, filenames in os.walk(dirPath):\n        for filename in filenames:\n            print(os.path.join(dirname, filename))\n\n# 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\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-05-20T07:22:52.13085Z","iopub.execute_input":"2021-05-20T07:22:52.131428Z","iopub.status.idle":"2021-05-20T07:22:52.138403Z","shell.execute_reply.started":"2021-05-20T07:22:52.131361Z","shell.execute_reply":"2021-05-20T07:22:52.13753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\n#gc.collect()","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.139855Z","iopub.execute_input":"2021-05-20T07:22:52.14031Z","iopub.status.idle":"2021-05-20T07:22:52.149723Z","shell.execute_reply.started":"2021-05-20T07:22:52.140272Z","shell.execute_reply":"2021-05-20T07:22:52.148932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Refference\n\n1. frog dudes masterpieces https://www.kaggle.com/c/bms-molecular-translation/discussion/231190","metadata":{}},{"cell_type":"markdown","source":"## 待研究\n- 青蛙做了什么优化 \n1. cnn（tnt） + tokenizer + transformer decoder 为什么比 cnn + lstm 好 ？\n    ans: Ranger \n2. 青蛙的tnt结构可以直接迁移到tpu上吗？ 似乎是可以的，结构中并没有啥np的结构\n3. 什么样的情况下模型可以反复训练\n4. 似乎结构上只有预处理 padding mask 上有优化\n\n- 为什么内存会超\n1. dataloader 里有什么增加了内存\n2. transformer 训练过程中有什么变量一直没释放\n3. dataloader 太大，可以分批小范围load\n4. 训练过程中做do_valid 不合适 非常占资源\n\n- 如果\n1. 如果不做patch会省内存嘛？\n2. 我能不能搭建个vit\n3. 如果不处理内存爆掉的问题会如何\n\n- 结论\n1. 想来想去，这个版本实在是太慢了，而且有包内存的问题，即使不做do validate 也只能持续20k个batch左右\n2. 后面还是根据这个博客 手动搭建一下transformer https://www.kaggle.com/drzhuzhe/training-on-gpu-bms/","metadata":{}},{"cell_type":"code","source":"import math\nimport gc\n\nimport torch\nfrom torch.utils.data.dataset import Dataset\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.sampler import *\nimport torch.nn.functional as F\nfrom torch.nn.parallel.data_parallel import data_parallel\n\nimport random\nimport pickle\nimport time\nfrom timeit import default_timer as timer\n\nimport cv2\n\n#%pylab inline\nimport matplotlib.pyplot as plt\n#import matplotlib.image as mpimg\n\nimport collections\nfrom collections import defaultdict","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.151824Z","iopub.execute_input":"2021-05-20T07:22:52.152446Z","iopub.status.idle":"2021-05-20T07:22:52.159837Z","shell.execute_reply.started":"2021-05-20T07:22:52.152408Z","shell.execute_reply":"2021-05-20T07:22:52.159047Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Set up Torch","metadata":{}},{"cell_type":"code","source":"def seed_py(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    return seed\n\ndef seed_torch(seed):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    return seed","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.162884Z","iopub.execute_input":"2021-05-20T07:22:52.163145Z","iopub.status.idle":"2021-05-20T07:22:52.173234Z","shell.execute_reply.started":"2021-05-20T07:22:52.163115Z","shell.execute_reply":"2021-05-20T07:22:52.172089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed = int(time.time())#335202   #5202  #123  #\nseed_py(seed)\nseed_torch(seed)\n\ntorch.backends.cudnn.benchmark     = False  ##uses the inbuilt cudnn auto-tuner to find the fastest convolution algorithms. -\ntorch.backends.cudnn.enabled       = True\ntorch.backends.cudnn.deterministic = True\n\nCOMMON_STRING = '\\tpytorch\\n'\nCOMMON_STRING += '\\t\\tseed = %d\\n'%seed\nCOMMON_STRING += '\\t\\ttorch.__version__              = %s\\n'%torch.__version__\nCOMMON_STRING += '\\t\\ttorch.version.cuda             = %s\\n'%torch.version.cuda\nCOMMON_STRING += '\\t\\ttorch.backends.cudnn.version() = %s\\n'%torch.backends.cudnn.version()\ntry:\n    COMMON_STRING += '\\t\\tos[\\'CUDA_VISIBLE_DEVICES\\']     = %s\\n'%os.environ['CUDA_VISIBLE_DEVICES']\n    NUM_CUDA_DEVICES = len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))\nexcept Exception:\n    COMMON_STRING += '\\t\\tos[\\'CUDA_VISIBLE_DEVICES\\']     = None\\n'\n    NUM_CUDA_DEVICES = 1\n\nCOMMON_STRING += '\\t\\ttorch.cuda.device_count()      = %d\\n'%torch.cuda.device_count()\nCOMMON_STRING += '\\t\\ttorch.cuda.get_device_properties() = %s\\n' % str(torch.cuda.get_device_properties(0))[21:]\n\nCOMMON_STRING += '\\n'\n","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.176044Z","iopub.execute_input":"2021-05-20T07:22:52.176325Z","iopub.status.idle":"2021-05-20T07:22:52.184265Z","shell.execute_reply.started":"2021-05-20T07:22:52.176299Z","shell.execute_reply":"2021-05-20T07:22:52.183328Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configure","metadata":{}},{"cell_type":"code","source":"origin_data_dir = '../input/bms-molecular-translation'\ndata_dir = '../input/bmstrainpart1/bms-moleular-translation'\ncsv_data_dir = '../input/bmd-mate-csv/csv'\nSTOI = {\n    '<sos>': 190,\n    '<eos>': 191,\n    '<pad>': 192,\n}\n\n\npatch_size   = 16\npixel_pad    = 3\npixel_scale  = 0.8  #1.0  #0.62=36/58 #1.0\nmax_length = 300 #278 #275\n\npixel_stride = 4\npixel_dim  = 24\npatch_dim  = 384\n\n_walkdir(csv_data_dir)","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.187031Z","iopub.execute_input":"2021-05-20T07:22:52.187279Z","iopub.status.idle":"2021-05-20T07:22:52.200315Z","shell.execute_reply.started":"2021-05-20T07:22:52.187247Z","shell.execute_reply":"2021-05-20T07:22:52.199551Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"def read_pickle_from_file(pickle_file):\n    with open(pickle_file,'rb') as f:\n        x = pickle.load(f)\n    return x\n\ndef uncompress_array(compressed_k):\n    compressed_k.seek(0)\n    k  = np.load(compressed_k,allow_pickle=True)['arr_0']\n    return k","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.201727Z","iopub.execute_input":"2021-05-20T07:22:52.20223Z","iopub.status.idle":"2021-05-20T07:22:52.208386Z","shell.execute_reply.started":"2021-05-20T07:22:52.202167Z","shell.execute_reply":"2021-05-20T07:22:52.20743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# draw -----------------------------------\ndef image_show(name, image, resize=.1):\n    H,W = image.shape[0:2]\n    plt.figure(figsize = (round(resize*W), round(resize*H)), dpi=20)\n    #fig = plt.figure(figsize = (500,500), dpi=20)\n    #ax = fig.add_subplot(1)                        \n    plt.imshow(image, cmap = plt.cm.gray)\n    #ax.imshow(image)\n    #image = cv2.resize(image, (round(resize*W), round(resize*H)), interpolation = cv2.INTER_AREA)\n    #plt.imshow(image)\n    \ndef resize_image(image, scale=1):\n    if scale==1 :\n        f = pixel_scale * 58/36  #1.2414 #0.80555\n        b = int(round(36*0.5))\n\n    if scale==2 :\n        f = pixel_scale * 1\n        b = int(round(58*0.5))\n\n    image = image[b:-b,b:-b] #remove border\n    if not np.isclose(1,f, rtol=1e-02, atol=1e-02):\n        h, w = image.shape\n        fw = int(round(f*w))\n        fh = int(round(f*h))\n        image = cv2.resize(image, dsize=(fw, fh), interpolation=cv2.INTER_AREA)\n    return image\n\ndef repad_image(image, multiplier=16):\n    h, w = image.shape\n    fh = int(np.ceil(h/multiplier))*multiplier\n    fw = int(np.ceil(w/multiplier))*multiplier\n    m  = np.full((fh, fw), 255, np.uint8)\n    m[0:h, 0:w] = image\n    return m","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.209882Z","iopub.execute_input":"2021-05-20T07:22:52.210293Z","iopub.status.idle":"2021-05-20T07:22:52.222038Z","shell.execute_reply.started":"2021-05-20T07:22:52.210257Z","shell.execute_reply":"2021-05-20T07:22:52.221157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# it seems to fill patches to Images\ndef patch_to_image(patch, coord, width, height):\n    image = np.full((height,width), 255, np.uint8)\n    p = pixel_pad\n    patch = patch[:, p:-p, p:-p]\n    num_patch = len(patch)\n\n    for i in range(num_patch):\n        y,x = coord[i]\n        x = x * patch_size\n        y = y * patch_size\n        image[y:y+patch_size,x:x+patch_size] = patch[i]\n        cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)\n    return image\n\ndef image_to_patch(image, patch_size, pixel_pad, threshold=0):\n    p = pixel_pad\n    h, w = image.shape\n\n    x, y = np.meshgrid(np.arange(w // patch_size), np.arange(h // patch_size))\n    yx = np.stack([y, x], 2).reshape(-1, 2)\n\n    s = patch_size + 2*p\n    m = torch.from_numpy(image).reshape(1, 1, h, w).float()\n    k = F.unfold(m, kernel_size=s, stride=patch_size, padding=p)\n    k = k.permute(0, 2, 1).reshape(-1, s * s)\n    k = k.data.cpu().numpy().reshape(-1, s, s)\n    #print(k.shape)\n\n    sum = (1 - k[:, p:-p, p:-p]/255).reshape(len(k), -1).sum(-1)\n    # only store patches\n    i = np.where(sum > threshold)\n    #print(sum)\n    patch = k[i]\n    coord = yx[i]\n    return  patch, coord","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.31779Z","iopub.execute_input":"2021-05-20T07:22:52.318131Z","iopub.status.idle":"2021-05-20T07:22:52.330468Z","shell.execute_reply.started":"2021-05-20T07:22:52.31809Z","shell.execute_reply":"2021-05-20T07:22:52.329545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\n# check folder data\n#df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')\ndf = read_pickle_from_file(csv_data_dir+'/df_train.more.csv.pickle')\ndf_fold = pd.read_csv(csv_data_dir+'/df_fold.fine.csv')\ndf_meta = pd.read_csv(csv_data_dir+'/df_train_image_meta.csv')\ndf = df.merge(df_fold, on='image_id')\ndf = df.merge(df_meta, on='image_id')\nprint(len(df), len(df_fold), len(df_meta))\nprint(len(df['fold'].unique()))\ndf.head(10)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.331841Z","iopub.execute_input":"2021-05-20T07:22:52.332399Z","iopub.status.idle":"2021-05-20T07:22:52.34759Z","shell.execute_reply.started":"2021-05-20T07:22:52.332361Z","shell.execute_reply":"2021-05-20T07:22:52.346786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# divide dataset into training and validate\ndef make_fold(mode='train-1'):\n    if 'train' in mode:\n        df = read_pickle_from_file(csv_data_dir+'/df_train.more.csv.pickle')\n        #df = pd.read_csv('../input/tarfile/df_train_patch_s0.800.csv')\n        #df_fold = pd.read_csv(data_dir+'/df_fold.csv')\n        df_fold = pd.read_csv(csv_data_dir+'/df_fold.fine.csv')\n        df_meta = pd.read_csv(csv_data_dir+'/df_train_image_meta.csv')\n        df = df.merge(df_fold, on='image_id')\n        df = df.merge(df_meta, on='image_id')\n        df.loc[:,'path']='train_patch16_s0.800'\n\n        df['fold'] = df['fold'].astype(int)\n        #print(df.groupby(['fold']).size()) #404_031\n        #print(df.columns)\n\n        fold = int(mode[-1])*10\n        #print(fold)\n        df_train = df[df.fold != fold].reset_index(drop=True)\n        df_valid = df[df.fold == fold].reset_index(drop=True)\n        return df_train, df_valid\n\n    # Index(['image_id', 'InChI'], dtype='object')\n    if 'test' in mode:\n        #df = pd.read_csv(data_dir+'/sample_submission.csv')\n        df = pd.read_csv(data_dir+'/submit_lb3.80.csv')\n        df_meta = pd.read_csv(data_dir+'/df_test_image_meta.csv')\n        df = df.merge(df_meta, on='image_id')\n\n        df.loc[:, 'path'] = 'test'\n        #df.loc[:, 'InChI'] = '0'\n        df.loc[:, 'formula'] = '0'\n        df.loc[:, 'text'] =  '0'\n        df.loc[:, 'sequence'] = pd.Series([[0]] * len(df))\n        df.loc[:, 'length'] = df.InChI.str.len()\n\n        df_test = df\n        return df_test","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.350779Z","iopub.execute_input":"2021-05-20T07:22:52.351069Z","iopub.status.idle":"2021-05-20T07:22:52.360763Z","shell.execute_reply.started":"2021-05-20T07:22:52.351037Z","shell.execute_reply":"2021-05-20T07:22:52.359241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#df_train, df_valid = make_fold()\n#print(len(df_train), len(df_valid))","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.362601Z","iopub.execute_input":"2021-05-20T07:22:52.36304Z","iopub.status.idle":"2021-05-20T07:22:52.372092Z","shell.execute_reply.started":"2021-05-20T07:22:52.363004Z","shell.execute_reply":"2021-05-20T07:22:52.371314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#-----------------------------------------------------------------------\n# tokenization, padding, ...\ndef pad_sequence_to_max_length(sequence, max_length, padding_value):\n    batch_size =len(sequence)\n    pad_sequence = np.full((batch_size,max_length), padding_value, np.int32)\n    for b, s in enumerate(sequence):\n        L = len(s)\n        pad_sequence[b, :L, ...] = s\n    return pad_sequence\n\ndef load_tokenizer():\n    tokenizer = YNakamaTokenizer(is_load=True)\n    print('len(tokenizer) : vocab_size', len(tokenizer))\n    for k,v in STOI.items():\n        assert  tokenizer.stoi[k]==v\n    return tokenizer\n\ndef null_augment(r):\n    return r","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.373477Z","iopub.execute_input":"2021-05-20T07:22:52.373901Z","iopub.status.idle":"2021-05-20T07:22:52.381952Z","shell.execute_reply.started":"2021-05-20T07:22:52.373845Z","shell.execute_reply":"2021-05-20T07:22:52.381176Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_TOKENIZER_ = {'(': 0, ')': 1, '+': 2, ',': 3, '-': 4, '/b': 5, '/c': 6, '/h': 7, '/i': 8, '/m': 9, '/s': 10, '/t': 11, '0': 12, '1': 13, '10': 14, '100': 15, '101': 16, '102': 17, '103': 18, '104': 19, '105': 20, '106': 21, '107': 22, '108': 23, '109': 24, '11': 25, '110': 26, '111': 27, '112': 28, '113': 29, '114': 30, '115': 31, '116': 32, '117': 33, '118': 34, '119': 35, '12': 36, '120': 37, '121': 38, '122': 39, '123': 40, '124': 41, '125': 42, '126': 43, '127': 44, '128': 45, '129': 46, '13': 47, '130': 48, '131': 49, '132': 50, '133': 51, '134': 52, '135': 53, '136': 54, '137': 55, '138': 56, '139': 57, '14': 58, '140': 59, '141': 60, '142': 61, '143': 62, '144': 63, '145': 64, '146': 65, '147': 66, '148': 67, '149': 68, '15': 69, '150': 70, '151': 71, '152': 72, '153': 73, '154': 74, '155': 75, '156': 76, '157': 77, '158': 78, '159': 79, '16': 80, '161': 81, '163': 82, '165': 83, '167': 84, '17': 85, '18': 86, '19': 87, '2': 88, '20': 89, '21': 90, '22': 91, '23': 92, '24': 93, '25': 94, '26': 95, '27': 96, '28': 97, '29': 98, '3': 99, '30': 100, '31': 101, '32': 102, '33': 103, '34': 104, '35': 105, '36': 106, '37': 107, '38': 108, '39': 109, '4': 110, '40': 111, '41': 112, '42': 113, '43': 114, '44': 115, '45': 116, '46': 117, '47': 118, '48': 119, '49': 120, '5': 121, '50': 122, '51': 123, '52': 124, '53': 125, '54': 126, '55': 127, '56': 128, '57': 129, '58': 130, '59': 131, '6': 132, '60': 133, '61': 134, '62': 135, '63': 136, '64': 137, '65': 138, '66': 139, '67': 140, '68': 141, '69': 142, '7': 143, '70': 144, '71': 145, '72': 146, '73': 147, '74': 148, '75': 149, '76': 150, '77': 151, '78': 152, '79': 153, '8': 154, '80': 155, '81': 156, '82': 157, '83': 158, '84': 159, '85': 160, '86': 161, '87': 162, '88': 163, '89': 164, '9': 165, '90': 166, '91': 167, '92': 168, '93': 169, '94': 170, '95': 171, '96': 172, '97': 173, '98': 174, '99': 175, 'B': 176, 'Br': 177, 'C': 178, 'Cl': 179, 'D': 180, 'F': 181, 'H': 182, 'I': 183, 'N': 184, 'O': 185, 'P': 186, 'S': 187, 'Si': 188, 'T': 189, '<sos>': 190, '<eos>': 191, '<pad>': 192}","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.383202Z","iopub.execute_input":"2021-05-20T07:22:52.383563Z","iopub.status.idle":"2021-05-20T07:22:52.402379Z","shell.execute_reply.started":"2021-05-20T07:22:52.383526Z","shell.execute_reply":"2021-05-20T07:22:52.401571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# getDataset\nclass BmsDataset(Dataset):\n    def __init__(self, df, tokenizer, augment=null_augment):\n        super().__init__()\n        self.tokenizer = tokenizer\n        self.df = df\n        self.augment = augment\n        self.length = len(self.df)\n\n    def __str__(self):\n        string  = ''\n        string += '\\tlen = %d\\n'%len(self)\n        string += '\\tdf  = %s\\n'%str(self.df.shape)\n\n        g = self.df['length'].values.astype(np.int32)//20\n        g = np.bincount(g,minlength=14)\n        string += '\\tlength distribution\\n'\n        for n in range(14):\n            string += '\\t\\t %3d = %8d (%0.4f)\\n'%((n+1)*20,g[n], g[n]/g.sum() )\n        return string\n\n    def __len__(self):\n        return self.length\n\n    def __getitem__(self, index):\n        d = self.df.iloc[index]\n        token = d.sequence\n\n        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)\n        k = read_pickle_from_file(patch_file)    \n        \n        patch = uncompress_array(k['patch'])\n        patch = np.concatenate([\n            np.zeros((1, patch_size+2*pixel_pad, patch_size+2*pixel_pad), np.uint8),\n            patch],0) #cls token\n\n        coord  = k['coord']\n        w = k['width' ]\n        h = k['height']\n\n        h = h // patch_size -1\n        w = w // patch_size -1\n        coord = np.insert(coord, 0, [h, w], 0) #cls token\n\n        #debug\n        # image = patch_to_image(patch, coord, k['width' ], k['height'])\n        # image_show('image', image, resize=1)\n        # cv2.waitKey(0)\n        #image = cv2.imread(image_file,cv2.IMREAD_GRAYSCALE) #\n\n        r = {\n            'index'    : index,\n            'image_id' : d.image_id,\n            'InChI'    : d.InChI,\n            'd' : d,\n            'token' : token,\n            #'image' : image,\n            'patch' : patch,\n            'coord' : coord,\n        }\n        if self.augment is not None: r = self.augment(r)\n        \n        # Warrning runtime keep this variable shall cause Out of memery crash \n        del patch_file, patch, coord, k\n        #gc.collect()\n        \n        return r","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.40373Z","iopub.execute_input":"2021-05-20T07:22:52.404123Z","iopub.status.idle":"2021-05-20T07:22:52.418549Z","shell.execute_reply.started":"2021-05-20T07:22:52.404087Z","shell.execute_reply":"2021-05-20T07:22:52.417321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# set up tokenizer\nclass YNakamaTokenizer(object):\n\n    def __init__(self, is_load=True):\n        self.stoi = {}\n        self.itos = {}\n\n        if is_load:\n            self.stoi = _TOKENIZER_\n            self.itos = {k: v for v, k in self.stoi.items()}\n\n    def __len__(self):\n        return len(self.stoi)\n\n    def build_vocab(self, text):\n        vocab = set()\n        for t in text:\n            vocab.update(t.split(' '))\n        vocab = sorted(vocab)\n        vocab.append('<sos>')\n        vocab.append('<eos>')\n        vocab.append('<pad>')\n        for i, s in enumerate(vocab):\n            self.stoi[s] = i\n        self.itos = {k: v for v, k in self.stoi.items()}\n\n    def one_text_to_sequence(self, text):\n        sequence = []\n        sequence.append(self.stoi['<sos>'])\n        for s in text.split(' '):\n            sequence.append(self.stoi[s])\n        sequence.append(self.stoi['<eos>'])\n        return sequence\n\n    def one_sequence_to_text(self, sequence):\n        return ''.join(list(map(lambda i: self.itos[i], sequence)))\n\n    def one_predict_to_inchi(self, predict):\n        inchi = 'InChI=1S/'\n        for p in predict:\n            if p == self.stoi['<eos>'] or p == self.stoi['<pad>']:\n                break\n            inchi += self.itos[p]\n        return inchi\n\n    # ---\n    def text_to_sequence(self, text):\n        sequence = [\n            self.one_text_to_sequence(t)\n            for t in text\n        ]\n        return sequence\n\n    def sequence_to_text(self, sequence):\n        text = [\n            self.one_sequence_to_text(s)\n            for s in sequence\n        ]\n        return text\n\n    def predict_to_inchi(self, predict):\n        inchi = [\n            self.one_predict_to_inchi(p)\n            for p in predict\n        ]\n        return inchi\n    \ndef null_collate(batch, is_sort_decreasing_length=True):\n    collate = defaultdict(list)\n\n    if is_sort_decreasing_length: #sort by decreasing length\n        sort  = np.argsort([-len(r['token']) for r in batch])\n        batch = [batch[s] for s in sort]\n\n    for r in batch:\n        for k, v in r.items():\n            collate[k].append(v)\n    #----\n\n    batch_size = len(batch)\n    collate['length'] = [len(l) for l in collate['token']]\n\n    token  = [np.array(t,np.int32) for t in collate['token']]\n    token  = pad_sequence_to_max_length(token, max_length=max_length, padding_value=STOI['<pad>'])\n    collate['token'] = torch.from_numpy(token).long()\n\n    max_of_length = max(collate['length'])\n    token_pad_mask  = np.zeros((batch_size, max_of_length, max_of_length))\n    for b in range(batch_size):\n        L = collate['length'][b]\n        token_pad_mask [b, :L, :L] = 1 #+1 for cls_token\n\n    collate['token_pad_mask'] = torch.from_numpy(token_pad_mask).byte()\n    #-----\n    # image = np.stack(collate['image'])\n    # image = image.astype(np.float32) / 255\n    # collate['image'] = torch.from_numpy(image).unsqueeze(1).repeat(1,3,1,1)\n\n    #-----\n\n    collate['num_patch'] = [len(l) for l in collate['patch']]\n\n    max_of_num_patch = max(collate['num_patch'])\n    patch_pad_mask  = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))\n    patch = np.full((batch_size, max_of_num_patch, patch_size+2*pixel_pad, patch_size+2*pixel_pad),255) #pad as 255\n    coord = np.zeros((batch_size, max_of_num_patch, 2))\n    for b in range(batch_size):\n        N = collate['num_patch'][b]\n        patch[b, :N] = collate['patch'][b]\n        coord[b, :N] = collate['coord'][b]\n        patch_pad_mask [b, :N, :N] = 1 #+1 for cls_token\n\n    collate['patch'] = torch.from_numpy(patch).half() / 255\n    collate['coord'] = torch.from_numpy(coord).long()\n    collate['patch_pad_mask' ] = torch.from_numpy(patch_pad_mask).byte()\n    \n    del patch_pad_mask, patch, coord, token_pad_mask, batch\n    \n    return collate","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.419923Z","iopub.execute_input":"2021-05-20T07:22:52.420394Z","iopub.status.idle":"2021-05-20T07:22:52.443386Z","shell.execute_reply.started":"2021-05-20T07:22:52.420358Z","shell.execute_reply":"2021-05-20T07:22:52.442563Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_check_dataset():\n    tokenizer = load_tokenizer()\n    df_train, df_valid = make_fold('train-1')\n\n    # df_train = make_fold('test') #1616107\n    # dataset = BmsDataset(df_train, tokenizer, remote_augment)\n\n    #dataset = BmsDataset(df_valid, tokenizer)\n    #print(dataset)\n\n    train_dataset = BmsDataset(df_train,tokenizer)\n        \n        \n    # for i in range(len(dataset)):\n    \"\"\"\n    for i in range(5):\n        #i = np.random.choice(len(dataset))\n        r = dataset[i]\n\n        print(r['index'])\n        print(r['image_id'])\n        #print(r['formula'])\n        print(r['InChI'])\n        print(r['token'])\n\n        print('image : ')\n        #print('\\t', r['image'].shape)\n        print('')\n\n        #---\n        image = patch_to_image(r['patch'], r['coord'], width=1024, height=1024)\n        #image_show('image', image, resize=.1)\n        del image\n        gc.collect()\n        #cv2.waitKey(0)\n\n    \"\"\"\n  \n\n\n    #exit(0)\n    #\"\"\"\n    loader = DataLoader(\n        train_dataset,\n        #sampler = RandomSampler(train_dataset),\n        batch_size  = 8,\n        drop_last   = True,\n        num_workers = 8,\n        pin_memory  = True,\n        collate_fn  = null_collate,\n    )\n    for t,batch in enumerate(loader):\n        #if t>30: break\n        print(t, '-----------')\n        #print('index : ', batch['index'])\n        \"\"\"\n        print('image : ')\n        #print('\\t', batch['image'].shape, batch['image'].is_contiguous())\n        print('\\t', batch['patch'].shape, batch['patch'].is_contiguous())\n        print('\\t', batch['coord'].shape, batch['coord'].is_contiguous())\n        #print('\\t', batch['mask'].shape, batch['mask'].is_contiguous())\n        print('length  : ')\n        print('\\t',len( batch['length']))\n        print('\\t', batch['length'])\n        print('token  : ')\n        print('\\t', batch['token'].shape, batch['token'].is_contiguous())\n        \"\"\"\n        #print('\\t', batch['token'])\n\n        #print('')\n        del t, batch\n    del loader, train_dataset, df_train, df_valid, tokenizer\n    #\"\"\"    ","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.445761Z","iopub.execute_input":"2021-05-20T07:22:52.446265Z","iopub.status.idle":"2021-05-20T07:22:52.456267Z","shell.execute_reply.started":"2021-05-20T07:22:52.446227Z","shell.execute_reply":"2021-05-20T07:22:52.455468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 数据load到4000 条 左右就会挂 ， 讲道理 load数据应该不会不断加内存吧\n#run_check_dataset()","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.457742Z","iopub.execute_input":"2021-05-20T07:22:52.458165Z","iopub.status.idle":"2021-05-20T07:22:52.468442Z","shell.execute_reply.started":"2021-05-20T07:22:52.458127Z","shell.execute_reply":"2021-05-20T07:22:52.467609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Net Modeules","metadata":{}},{"cell_type":"markdown","source":"## 1D Encoder and Decoder","metadata":{}},{"cell_type":"code","source":"!pip install fairseq\nfrom typing import Tuple, Dict\nimport torch.nn as nn\nfrom fairseq import utils\nfrom fairseq.models import *\nfrom fairseq.modules import *","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:52.469786Z","iopub.execute_input":"2021-05-20T07:22:52.470392Z","iopub.status.idle":"2021-05-20T07:22:58.25379Z","shell.execute_reply.started":"2021-05-20T07:22:52.470352Z","shell.execute_reply":"2021-05-20T07:22:58.252711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ------------------------------------------------------\n# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need\n\nclass PositionEncode1D(nn.Module):\n    def __init__(self, dim, max_length):\n        super().__init__()\n        assert (dim % 2 == 0)\n        self.max_length = max_length\n\n        d = torch.exp(torch.arange(0., dim, 2)* (-math.log(10000.0) / dim))\n        position = torch.arange(0., max_length).unsqueeze(1)\n        pos = torch.zeros(1, max_length, dim)\n        pos[0, :, 0::2] = torch.sin(position * d)\n        pos[0, :, 1::2] = torch.cos(position * d)\n        self.register_buffer('pos', pos)\n\n    def forward(self, x):\n        batch_size, T, dim = x.shape\n        x = x + self.pos[:,:T]\n        return x\n\n    \n\"\"\"\n# https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py\n#https://github.com/pytorch/fairseq/issues/568\n# fairseq/fairseq/models/fairseq_encoder.py\n\n# https://github.com/pytorch/fairseq/blob/master/fairseq/modules/transformer_layer.py\nclass TransformerEncode(FairseqEncoder):\n\n    def __init__(self, dim, ff_dim, num_head, num_layer):\n        super().__init__({})\n        #print('my TransformerEncode()')\n\n        self.layer = nn.ModuleList([\n            TransformerEncoderLayer(Namespace({\n                'encoder_embed_dim': dim,\n                'encoder_attention_heads': num_head,\n                'attention_dropout': 0.1,\n                'dropout': 0.1,\n                'encoder_normalize_before': True,\n                'encoder_ffn_embed_dim': ff_dim,\n            })) for i in range(num_layer)\n        ])\n        self.layer_norm = nn.LayerNorm(dim)\n\n    def forward(self, x):# T x B x C\n        #print('my TransformerEncode forward()')\n        for layer in self.layer:\n            x = layer(x)\n        x = self.layer_norm(x)\n        return x\n\"\"\"    \n\n\n# https://mt.cs.upc.edu/2020/12/21/the-transformer-fairseq-edition/\n# for debug\n# class TransformerDecode(FairseqDecoder):\n#     def __init__(self, dim, ff_dim, num_head, num_layer):\n#         super().__init__({})\n#         print('my TransformerDecode()')\n#\n#         self.layer = nn.ModuleList([\n#             TransformerDecoderLayer(Namespace({\n#                 'decoder_embed_dim': dim,\n#                 'decoder_attention_heads': num_head,\n#                 'attention_dropout': 0.1,\n#                 'dropout': 0.1,\n#                 'decoder_normalize_before': True,\n#                 'decoder_ffn_embed_dim': ff_dim,\n#             })) for i in range(num_layer)\n#         ])\n#         self.layer_norm = nn.LayerNorm(dim)\n#\n#\n#     def forward(self, x, mem, x_mask):# T x B x C\n#         print('my TransformerDecode forward()')\n#         for layer in self.layer:\n#             x = layer(x, mem, self_attn_mask=x_mask)[0]\n#         x = self.layer_norm(x)\n#         return x  # T x B x C\n#https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute\nclass Namespace(object):\n    def __init__(self, adict):\n        self.__dict__.update(adict)\n        \n# https://fairseq.readthedocs.io/en/latest/tutorial_simple_lstm.html\n# see https://gitlab.maastrichtuniversity.nl/dsri-examples/dsri-pytorch-workspace/-/blob/c8a88cdeb8e1a0f3a2ccd3c6119f43743cbb01e9/examples/transformer/fairseq/models/transformer.py\nclass TransformerDecode(FairseqIncrementalDecoder):\n    def __init__(self, dim, ff_dim, num_head, num_layer):\n        super().__init__({})\n        #print('my TransformerDecode()')\n\n        self.layer = nn.ModuleList([\n            TransformerDecoderLayer(Namespace({\n                'decoder_embed_dim': dim,\n                'decoder_attention_heads': num_head,\n                'attention_dropout': 0.1,\n                'dropout': 0.1,\n                'decoder_normalize_before': True,\n                'decoder_ffn_embed_dim': ff_dim,\n            })) for i in range(num_layer)\n        ])\n        self.layer_norm = nn.LayerNorm(dim)\n\n\n    def forward(self, x, mem, x_mask, x_pad_mask, mem_pad_mask):\n            #print('my TransformerDecode forward()')\n            for layer in self.layer:\n                x = layer(\n                    x,\n                    mem,\n                    self_attn_mask=x_mask,\n                    self_attn_padding_mask=x_pad_mask,\n                    encoder_padding_mask=mem_pad_mask,\n                )[0]\n            x = self.layer_norm(x)\n            return x  # T x B x C\n\n    #def forward_one(self, x, mem, incremental_state):\n    def forward_one(self,\n            x   : Tensor,\n            mem : Tensor,\n            incremental_state : Optional[Dict[str, Dict[str, Optional[Tensor]]]]\n    )-> Tensor:\n        x = x[-1:]\n        for layer in self.layer:\n            x = layer(x, mem, incremental_state=incremental_state)[0]\n        x = self.layer_norm(x)\n        return x\n\n'''\nhttps://fairseq.readthedocs.io/en/latest/_modules/fairseq/modules/transformer_layer.html\nclass TransformerDecoderLayer(nn.Module):\ndef forward(\n        self,\n        x,\n        encoder_out: Optional[torch.Tensor] = None,\n        encoder_padding_mask: Optional[torch.Tensor] = None,\n        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,\n        prev_self_attn_state: Optional[List[torch.Tensor]] = None,\n        prev_attn_state: Optional[List[torch.Tensor]] = None,\n        self_attn_mask: Optional[torch.Tensor] = None,\n        self_attn_padding_mask: Optional[torch.Tensor] = None,\n        need_attn: bool = False,\n        need_head_weights: bool = False,\n    ):\n        \n        \nhttps://github.com/pytorch/fairseq/blob/05b86005bcca0155319fa9b81abfd69f63c06906/examples/simultaneous_translation/utils/data_utils.py#L31      \n'''\n","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:58.259099Z","iopub.execute_input":"2021-05-20T07:22:58.261277Z","iopub.status.idle":"2021-05-20T07:22:58.292493Z","shell.execute_reply.started":"2021-05-20T07:22:58.261233Z","shell.execute_reply":"2021-05-20T07:22:58.291635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TNT","metadata":{}},{"cell_type":"code","source":"!pip install timm\n\nfrom torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence\n#https://github.com/pytorch/pytorch/issues/1788\n# https://stackoverflow.com/questions/51030782/why-do-we-pack-the-sequences-in-pytorch\n\nfrom timm.models.vision_transformer import Mlp\nfrom timm.models.layers import DropPath, trunc_normal_\n#from timm.models.tnt import *","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:22:58.297148Z","iopub.execute_input":"2021-05-20T07:22:58.299196Z","iopub.status.idle":"2021-05-20T07:23:04.014749Z","shell.execute_reply.started":"2021-05-20T07:22:58.299157Z","shell.execute_reply":"2021-05-20T07:23:04.013729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 这个多头注意力会造成内存不断增加吗？\nclass Attention(nn.Module):\n    \"\"\" Multi-Head Attention\n    \"\"\"\n\n    def __init__(self, dim, hidden_dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):\n        super().__init__()\n        self.hidden_dim = hidden_dim\n        self.num_heads = num_heads\n        head_dim = hidden_dim // num_heads\n        self.head_dim = head_dim\n        self.scale = head_dim ** -0.5\n\n        self.qk = nn.Linear(dim, hidden_dim * 2, bias=qkv_bias)\n        self.v = nn.Linear(dim, dim, bias=qkv_bias)\n        self.attn_drop = nn.Dropout(attn_drop, inplace=True)\n        self.proj = nn.Linear(dim, dim)\n        self.proj_drop = nn.Dropout(proj_drop, inplace=True)\n\n    def forward(self,\n        x: Tensor,\n        mask: Optional[Tensor] = None\n    )-> Tensor:\n\n        B, N, C = x.shape\n        qk = self.qk(x).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n        q, k = qk[0], qk[1]  # make torchscript happy (cannot use tensor as tuple)\n        v = self.v(x).reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)\n\n        #---\n        attn = (q @ k.transpose(-2, -1)) * self.scale # B x self.num_heads x NxN\n        if mask is not None:\n            #mask = mask.unsqueeze(1).repeat(1,self.num_heads,1,1)\n            mask = mask.unsqueeze(1).expand(-1,self.num_heads,-1,-1)\n            attn = attn.masked_fill(mask == 0, -6e4)\n            # attn = attn.masked_fill(mask == 0, -half('inf'))\n            # https://github.com/NVIDIA/apex/issues/93\n            # How to use fp16 training with masked operations\n\n        attn = attn.softmax(dim=-1)\n        attn = self.attn_drop(attn)\n\n        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)\n        x = self.proj(x)\n        x = self.proj_drop(x)\n        del mask, attn\n        return x\n\n\n\n\nclass Block(nn.Module):\n    def __init__(self, dim, in_dim, num_pixel, num_heads=12, in_num_head=4, mlp_ratio=4.,\n                 qkv_bias=False, drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):\n        super().__init__()\n        # Inner transformer\n        self.norm_in = norm_layer(in_dim)\n        self.attn_in = Attention(\n            in_dim, in_dim, num_heads=in_num_head, qkv_bias=qkv_bias,\n            attn_drop=attn_drop, proj_drop=drop)\n\n        self.norm_mlp_in = norm_layer(in_dim)\n        self.mlp_in = Mlp(in_features=in_dim, hidden_features=int(in_dim * 4),\n                          out_features=in_dim, act_layer=act_layer, drop=drop)\n\n        self.norm1_proj = norm_layer(in_dim)\n        self.proj = nn.Linear(in_dim * num_pixel, dim, bias=True)\n        # Outer transformer\n        self.norm_out = norm_layer(dim)\n        self.attn_out = Attention(\n            dim, dim, num_heads=num_heads, qkv_bias=qkv_bias,\n            attn_drop=attn_drop, proj_drop=drop)\n        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n\n        self.norm_mlp = norm_layer(dim)\n        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio),\n                       out_features=dim, act_layer=act_layer, drop=drop)\n\n    def forward(self, pixel_embed, patch_embed, mask):\n        # inner\n        pixel_embed = pixel_embed + self.drop_path(self.attn_in(self.norm_in(pixel_embed)))\n        pixel_embed = pixel_embed + self.drop_path(self.mlp_in(self.norm_mlp_in(pixel_embed)))\n        # outer\n        B, N, C = patch_embed.size()\n        patch_embed[:, 1:] = patch_embed[:, 1:] + self.proj(self.norm1_proj(pixel_embed).reshape(B, N, -1))[:, 1:]\n        patch_embed = patch_embed + self.drop_path(self.attn_out(self.norm_out(patch_embed), mask))\n        patch_embed = patch_embed + self.drop_path(self.mlp(self.norm_mlp(patch_embed)))\n        return pixel_embed, patch_embed\n\n#---------------------------------\n\nclass PixelEmbed(nn.Module):\n\n    def __init__(self,  patch_size=16, in_dim=48, stride=4):\n        super().__init__()\n        self.in_dim = in_dim\n        self.proj = nn.Conv2d(3, self.in_dim, kernel_size=7, padding=0, stride=stride)\n\n    def forward(self, patch, pixel_pos):\n        BN = len(patch)\n        x = patch\n        x = self.proj(x)\n        #x = x.transpose(1, 2).reshape(B * self.num_patches, self.in_dim, self.new_patch_size, self.new_patch_size)\n        x = x + pixel_pos\n        x = x.reshape(BN, self.in_dim, -1).transpose(1, 2)\n        return x\n\n\n#---------------------------------\n\n\n\nclass TNT(nn.Module):\n    \"\"\" Transformer in Transformer - https://arxiv.org/abs/2103.00112\n    \"\"\"\n\n    def __init__(self,\n            patch_size=patch_size,\n            embed_dim =patch_dim,\n            in_dim=pixel_dim,\n            depth=12,\n            num_heads=6,\n            in_num_head=4,\n            mlp_ratio=4.,\n            qkv_bias=False,\n            drop_rate=0.,\n            attn_drop_rate=0.,\n            drop_path_rate=0.,\n            norm_layer=nn.LayerNorm,\n            first_stride=pixel_stride):\n        super().__init__()\n\n        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models\n\n        self.pixel_embed = PixelEmbed( patch_size=patch_size, in_dim=in_dim, stride=first_stride)\n        #num_patches = self.pixel_embed.num_patches\n        #self.num_patches = num_patches\n        new_patch_size = 4 #self.pixel_embed.new_patch_size\n        num_pixel = new_patch_size ** 2\n\n        self.norm1_proj = norm_layer(num_pixel * in_dim)\n        self.proj = nn.Linear(num_pixel * in_dim, embed_dim)\n        self.norm2_proj = norm_layer(embed_dim)\n\n        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n        self.patch_pos = nn.Embedding(100*100,embed_dim) #nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n        self.pixel_pos = nn.Parameter(torch.zeros(1, in_dim, new_patch_size, new_patch_size))\n        self.pos_drop = nn.Dropout(p=drop_rate)\n\n        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule\n        blocks = []\n        for i in range(depth):\n            blocks.append(Block(\n                dim=embed_dim, in_dim=in_dim, num_pixel=num_pixel, num_heads=num_heads, in_num_head=in_num_head,\n                mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate,\n                drop_path=dpr[i], norm_layer=norm_layer))\n        self.blocks = nn.ModuleList(blocks)\n        self.norm = norm_layer(embed_dim)\n\n        trunc_normal_(self.cls_token, std=.02)\n        #trunc_normal_(self.patch_pos, std=.02)\n        trunc_normal_(self.pixel_pos, std=.02)\n        self.apply(self._init_weights)\n\n    def _init_weights(self, m):\n        if isinstance(m, nn.Linear):\n            trunc_normal_(m.weight, std=.02)\n            if isinstance(m, nn.Linear) and m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n        elif isinstance(m, nn.LayerNorm):\n            nn.init.constant_(m.bias, 0)\n            nn.init.constant_(m.weight, 1.0)\n\n    @torch.jit.ignore\n    def no_weight_decay(self):\n        return {'patch_pos', 'pixel_pos', 'cls_token'}\n\n\n    def forward(self,  patch, coord, mask):\n        B = len(patch)\n        batch_size, max_of_num_patch, s, s = patch.shape\n\n        patch = patch.reshape(batch_size*max_of_num_patch, 1, s, s).repeat(1,3,1,1)\n        pixel_embed = self.pixel_embed(patch, self.pixel_pos)\n\n        patch_embed = self.norm2_proj(self.proj(self.norm1_proj(pixel_embed.reshape(B, max_of_num_patch, -1))))\n\n        #patch_embed = torch.cat((self.cls_token.expand(B, -1, -1), patch_embed), dim=1)\n        #patch_embed = patch_embed + self.patch_pos\n        #patch_embed[:, 1:] = patch_embed[:, 1:] + self.patch_pos(coord[:, :, 0] * 100 + coord[:, :, 1])\n\n        patch_embed[:,:1]= self.cls_token.expand(B, -1, -1)\n        patch_embed= patch_embed + self.patch_pos(coord[:, :, 0] * 100 + coord[:, :, 1])\n        patch_embed = self.pos_drop(patch_embed)\n\n        for blk in self.blocks:\n            pixel_embed, patch_embed = blk(pixel_embed, patch_embed, mask)\n\n        patch_embed = self.norm(patch_embed)\n        \n        del patch, mask\n        \n        return patch_embed","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.016529Z","iopub.execute_input":"2021-05-20T07:23:04.016885Z","iopub.status.idle":"2021-05-20T07:23:04.056664Z","shell.execute_reply.started":"2021-05-20T07:23:04.016827Z","shell.execute_reply":"2021-05-20T07:23:04.055896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_dummy_data():\n    # make dummy data\n    # image_id,width,height,scale,orientation\n    meta = [\n        ['000011a64c74', 325, 229, 2, 0, ],\n        ['000019cc0cd2', 288, 148, 1, 0, ],\n        ['0000252b6d2b', 509, 335, 2, 0, ],\n        ['000026b49b7e', 243, 177, 1, 0, ],\n        ['000026fc6c36', 294, 112, 1, 0, ],\n        ['000028818203', 402, 328, 2, 0, ],\n        ['000029a61c01', 395, 294, 2, 0, ],\n        ['000035624718', 309, 145, 1, 0, ],\n    ]\n    batch_size = 8\n\n    # <todo> check border for padding\n    # <todo> pepper noise\n\n    batch = {\n        'num_patch': [],\n        'patch': [],\n        'coord': [],\n    }\n    for b in range(batch_size):\n        image_id = meta[b][0]\n        scale = meta[b][3]\n\n        image_file = origin_data_dir + '/%s/%s/%s/%s/%s.png' % ('train', image_id[0], image_id[1], image_id[2], image_id)\n        #print(image_file)\n        image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)\n\n        image = resize_image(image, scale)\n        image = repad_image(image, patch_size)  # remove border and repad\n        # print(image.shape)\n\n        k, yx = image_to_patch(image, patch_size, pixel_pad, threshold=0)\n\n        for y, x in yx:\n            # cv2.circle(image,(x,y),8,128,1)\n            x = x * patch_size\n            y = y * patch_size\n            cv2.rectangle(image, (x, y), (x + patch_size, y + patch_size), 128, 1)\n\n        #image_show('image-%d' % b, image, resize=1)\n        del image\n        #cv2.waitKey(1)\n\n        batch['patch'].append(k)\n        batch['coord'].append(yx)\n        batch['num_patch'].append(len(k))\n\n    # ----\n    max_of_num_patch = max(batch['num_patch'])\n    mask = np.zeros((batch_size, max_of_num_patch, max_of_num_patch))\n    patch = np.zeros((batch_size, max_of_num_patch, patch_size + 2 * pixel_pad, patch_size + 2 * pixel_pad))\n    coord = np.zeros((batch_size, max_of_num_patch, 2))\n    for b in range(batch_size):\n        N = batch['num_patch'][b]\n        patch[b, :N] = batch['patch'][b]\n        coord[b, :N] = batch['coord'][b]\n        mask[b, :N, :N] = 1\n\n    num_patch = batch['num_patch']\n    patch = torch.from_numpy(patch).float()\n    coord = torch.from_numpy(coord).long()\n    mask = torch.from_numpy(mask).byte()\n\n    return patch,coord,num_patch,mask\n\n\ndef run_check_tnt_patch():\n    patch,coord,num_patch,mask = make_dummy_data()\n\n    tnt = TNT()\n    patch_embed = tnt(patch, coord, mask)\n    print(patch_embed.shape)\n","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.058074Z","iopub.execute_input":"2021-05-20T07:23:04.058453Z","iopub.status.idle":"2021-05-20T07:23:04.07598Z","shell.execute_reply.started":"2021-05-20T07:23:04.058412Z","shell.execute_reply":"2021-05-20T07:23:04.075181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#run_check_tnt_patch()","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.078534Z","iopub.execute_input":"2021-05-20T07:23:04.078913Z","iopub.status.idle":"2021-05-20T07:23:04.086743Z","shell.execute_reply.started":"2021-05-20T07:23:04.078876Z","shell.execute_reply":"2021-05-20T07:23:04.085925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TNT + PositionEncode1D + Transformer Decoder","metadata":{}},{"cell_type":"code","source":"from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence, PackedSequence\nvocab_size = 193\ntext_dim    = 384\ndecoder_dim = 384\nnum_layer = 3\nnum_head  = 8\nff_dim = 1024","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.088572Z","iopub.execute_input":"2021-05-20T07:23:04.089057Z","iopub.status.idle":"2021-05-20T07:23:04.096022Z","shell.execute_reply.started":"2021-05-20T07:23:04.08902Z","shell.execute_reply":"2021-05-20T07:23:04.095193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Net(nn.Module):\n\n    def __init__(self,):\n        super(Net, self).__init__()\n        self.cnn = TNT()\n        self.image_encode = nn.Identity()\n\n        #---\n        self.text_pos    = PositionEncode1D(text_dim,max_length)\n        self.token_embed = nn.Embedding(vocab_size, text_dim)\n        self.text_decode = TransformerDecode(decoder_dim, ff_dim, num_head, num_layer)\n\n        #---\n        self.logit  = nn.Linear(decoder_dim, vocab_size)\n        self.dropout = nn.Dropout(p=0.5)\n\n        #----\n        # initialization\n        self.token_embed.weight.data.uniform_(-0.1, 0.1)\n        self.logit.bias.data.fill_(0)\n        self.logit.weight.data.uniform_(-0.1, 0.1)\n\n\n    @torch.jit.unused\n    def forward(self, patch, coord, token, patch_pad_mask, token_pad_mask):\n        device = patch.device\n        batch_size = len(patch)\n        #---\n        patch = patch*2-1\n        image_embed = self.cnn(patch, coord, patch_pad_mask)\n        image_embed = self.image_encode(image_embed).permute(1,0,2).contiguous()\n\n        text_embed = self.token_embed(token)\n        text_embed = self.text_pos(text_embed).permute(1,0,2).contiguous()\n\n        max_of_length = token_pad_mask.shape[-1]\n        text_mask = np.triu(np.ones((max_of_length, max_of_length)), k=1).astype(np.uint8)\n        text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)==1).to(device)\n\n        #----\n        # <todo> perturb mask as aug\n        text_pad_mask = token_pad_mask[:,:,0]==0\n        image_pad_mask = patch_pad_mask[:,:,0]==0\n        x = self.text_decode(text_embed[:max_of_length], image_embed, text_mask, text_pad_mask, image_pad_mask)\n        x = x.permute(1,0,2).contiguous()\n        l = self.logit(x)\n\n        logit = torch.zeros((batch_size, max_length, vocab_size),device=device)\n        logit[:,:max_of_length]=l\n        \n        del image_embed, text_mask, text_pad_mask, image_pad_mask\n        \n        return logit\n\n    #submit function has not been coded. i will leave it as an exercise for the kaggler\n    #@torch.jit.export\n    # def forward_argmax_decode(self, patch, coord, mask):\n    #\n    #     image_dim   = 384\n    #     text_dim    = 384\n    #     decoder_dim = 384\n    #     num_layer = 3\n    #     num_head  = 8\n    #     ff_dim    = 1024\n    #\n    #     STOI = {\n    #         '<sos>': 190,\n    #         '<eos>': 191,\n    #         '<pad>': 192,\n    #     }\n    #     max_length = 278 # 275\n    #\n    #\n    #     #---------------------------------\n    #     device = patch.device\n    #     batch_size = len(patch)\n    #\n    #     patch = patch*2-1\n    #     image_embed = self.cnn(patch, coord, mask)\n    #     image_embed = self.image_encode(image_embed).permute(1,0,2).contiguous()\n    #\n    #     token = torch.full((batch_size, max_length), STOI['<pad>'],dtype=torch.long, device=device)\n    #     text_pos = self.text_pos.pos\n    #     token[:,0] = STOI['<sos>']\n    #\n    #\n    #     #-------------------------------------\n    #     eos = STOI['<eos>']\n    #     pad = STOI['<pad>']\n    #     # https://github.com/alexmt-scale/causal-transformer-decoder/blob/master/tests/test_consistency.py\n    #     # slow version\n    #     # if 0:\n    #     #     for t in range(max_length-1):\n    #     #         last_token = token [:,:(t+1)]\n    #     #         text_embed = self.token_embed(last_token)\n    #     #         text_embed = self.text_pos(text_embed).permute(1,0,2).contiguous() #text_embed + text_pos[:,:(t+1)] #\n    #     #\n    #     #         text_mask = np.triu(np.ones((t+1, t+1)), k=1).astype(np.uint8)\n    #     #         text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)==1).to(device)\n    #     #\n    #     #         x = self.text_decode(text_embed, image_embed, text_mask)\n    #     #         x = x.permute(1,0,2).contiguous()\n    #     #\n    #     #         l = self.logit(x[:,-1])\n    #     #         k = torch.argmax(l, -1)  # predict max\n    #     #         token[:, t+1] = k\n    #     #         if ((k == eos) | (k == pad)).all():  break\n    #\n    #     # fast version\n    #     if 1:\n    #         #incremental_state = {}\n    #         incremental_state = torch.jit.annotate(\n    #             Dict[str, Dict[str, Optional[Tensor]]],\n    #             torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {}),\n    #         )\n    #         for t in range(max_length-1):\n    #             #last_token = token [:,:(t+1)]\n    #             #text_embed = self.token_embed(last_token)\n    #             #text_embed = self.text_pos(text_embed) #text_embed + text_pos[:,:(t+1)] #\n    #\n    #             last_token = token[:, t]\n    #             text_embed = self.token_embed(last_token)\n    #             text_embed = text_embed + text_pos[:,t] #\n    #             text_embed = text_embed.reshape(1,batch_size,text_dim)\n    #\n    #             x = self.text_decode.forward_one(text_embed, image_embed, incremental_state)\n    #             x = x.reshape(batch_size,decoder_dim)\n    #             #print(incremental_state.keys())\n    #\n    #             l = self.logit(x)\n    #             k = torch.argmax(l, -1)  # predict max\n    #             token[:, t+1] = k\n    #             if ((k == eos) | (k == pad)).all():  break\n    #\n    #     predict = token[:, 1:]\n    #     return predict","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.097161Z","iopub.execute_input":"2021-05-20T07:23:04.099305Z","iopub.status.idle":"2021-05-20T07:23:04.11595Z","shell.execute_reply.started":"2021-05-20T07:23:04.099266Z","shell.execute_reply":"2021-05-20T07:23:04.115081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loss #################################################################\ndef seq_cross_entropy_loss(logit, token, length):\n    truth = token[:, 1:]\n    L = [l - 1 for l in length]\n    logit = pack_padded_sequence(logit, L, batch_first=True).data\n    truth = pack_padded_sequence(truth, L, batch_first=True).data\n    loss = F.cross_entropy(logit, truth, ignore_index=STOI['<pad>'])\n    return loss\n\n# https://www.aclweb.org/anthology/2020.findings-emnlp.276.pdf\ndef seq_focal_cross_entropy_loss(logit, token, length):\n    gamma = 0.5 # {0.5,1.0}\n    #label_smooth = 0.90\n\n    #---\n    truth = token[:, 1:]\n    L = [l - 1 for l in length]\n    logit = pack_padded_sequence(logit, L, batch_first=True).data\n    truth = pack_padded_sequence(truth, L, batch_first=True).data\n    #loss = F.cross_entropy(logit, truth, ignore_index=STOI['<pad>'])\n    #non_pad = torch.where(truth != STOI['<pad>'])[0]  # & (t!=STOI['<sos>'])\n\n\n    # ---\n    #p = F.softmax(logit,-1)\n    #logp = - torch.log(torch.clamp(p, 1e-4, 1 - 1e-4))\n\n    logp = F.log_softmax(logit, -1)\n    logp = logp.gather(1, truth.reshape(-1,1)).reshape(-1)\n    p = logp.exp()\n\n    loss = - ((1 - p) ** gamma)*logp  #focal\n    #loss = - ((1 + p) ** gamma)*logp  #anti-focal\n    loss = loss.mean()\n    return loss\n\n    \ndef np_loss_cross_entropy(probability, truth):\n    batch_size = len(probability)\n    truth = truth.reshape(-1)\n    p = probability[np.arange(batch_size),truth]\n    loss = -np.log(np.clip(p,1e-6,1))\n    loss = loss.mean()\n    return loss","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.117275Z","iopub.execute_input":"2021-05-20T07:23:04.117616Z","iopub.status.idle":"2021-05-20T07:23:04.129133Z","shell.execute_reply.started":"2021-05-20T07:23:04.117581Z","shell.execute_reply":"2021-05-20T07:23:04.127939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check #################################################################\n\n\ndef run_check_net():\n    patch, coord, num_patch, patch_pad_mask = make_dummy_data()\n    batch_size = len(patch)\n\n    token  = np.full((batch_size, max_length), STOI['<pad>'], np.int64) #token\n    length = np.random.randint(5, max_length-2, batch_size)\n    length = np.sort(length)[::-1].copy()\n\n\n    max_of_length = max(length)\n    token_pad_mask  = np.zeros((batch_size, max_of_length, max_of_length))\n    for b in range(batch_size):\n        l = length[b]\n        t = np.random.choice(vocab_size,l)\n        t = np.insert(t,0,     STOI['<sos>'])\n        t = np.insert(t,len(t),STOI['<eos>'])\n        L = len(t)\n        token[b,:L]=t\n        token_pad_mask [b, :L, :L] = 1\n\n    token = torch.from_numpy(token).long()\n    token_pad_mask = torch.from_numpy(token_pad_mask).byte()\n\n\n\n    #---\n    net = Net()\n    net.train()\n\n    logit = net(patch, coord, token, patch_pad_mask, token_pad_mask)\n    loss = seq_cross_entropy_loss(logit, token, length)\n\n\n    print('vocab_size',vocab_size)\n    print('max_length',max_length)\n    print('')\n    print(length)\n    print(length.shape)\n    print(token.shape)\n    print('---')\n\n    print(logit.shape)\n    print(loss)\n    print('---')\n    del net\n\n    #---\n    # print('torch.jit.script(net)')\n    # net.eval()\n    # net = torch.jit.script(net)\n    #\n    # predict = net.forward_argmax_decode(patch, coord, mask)\n    # print(predict.shape)","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.131701Z","iopub.execute_input":"2021-05-20T07:23:04.132041Z","iopub.status.idle":"2021-05-20T07:23:04.143699Z","shell.execute_reply.started":"2021-05-20T07:23:04.132004Z","shell.execute_reply":"2021-05-20T07:23:04.142951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#run_check_net()","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.145032Z","iopub.execute_input":"2021-05-20T07:23:04.14537Z","iopub.status.idle":"2021-05-20T07:23:04.155884Z","shell.execute_reply.started":"2021-05-20T07:23:04.145333Z","shell.execute_reply":"2021-05-20T07:23:04.154978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\nimport sys\n\nlocal_vars = list(locals().items())\nfor var, obj in local_vars:\n    print(var, sys.getsizeof(obj))\n\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.157158Z","iopub.execute_input":"2021-05-20T07:23:04.157538Z","iopub.status.idle":"2021-05-20T07:23:04.16652Z","shell.execute_reply.started":"2021-05-20T07:23:04.157463Z","shell.execute_reply":"2021-05-20T07:23:04.165551Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ----------------\nis_mixed_precision = True #False  #\n\n\n###################################################################################################\nimport torch.cuda.amp as amp\nif is_mixed_precision:\n    class AmpNet(Net):\n        @torch.cuda.amp.autocast()\n        def forward(self, *args):\n            return super(AmpNet, self).forward(*args)\nelse:\n    AmpNet = Net","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.167921Z","iopub.execute_input":"2021-05-20T07:23:04.168335Z","iopub.status.idle":"2021-05-20T07:23:04.175357Z","shell.execute_reply.started":"2021-05-20T07:23:04.168298Z","shell.execute_reply":"2021-05-20T07:23:04.174557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nfrom datetime import datetime\nimport Levenshtein\n\nIDENTIFIER   = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')\n\n# http://stackoverflow.com/questions/34950201/pycharm-print-end-r-statement-not-working\nclass Logger(object):\n    def __init__(self):\n        self.terminal = sys.stdout  #stdout\n        self.file = None\n\n    def open(self, file, mode=None):\n        if mode is None: mode ='w'\n        self.file = open(file, mode)\n\n    def write(self, message, is_terminal=1, is_file=1 ):\n        if '\\r' in message: is_file=0\n\n        if is_terminal == 1:\n            self.terminal.write(message)\n            self.terminal.flush()\n            #time.sleep(1)\n\n        if is_file == 1:\n            self.file.write(message)\n            self.file.flush()\n\n    def flush(self):\n        # this flush method is needed for python 3 compatibility.\n        # this handles the flush command by doing nothing.\n        # you might want to specify some extra behavior here.\n        pass\n\n# etc ------------------------------------\ndef time_to_str(t, mode='min'):\n    if mode=='min':\n        t  = int(t)/60\n        hr = t//60\n        min = t%60\n        return '%2d hr %02d min'%(hr,min)\n\n    elif mode=='sec':\n        t   = int(t)\n        min = t//60\n        sec = t%60\n        return '%2d min %02d sec'%(min,sec)\n\n    else:\n        raise NotImplementedError","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.176785Z","iopub.execute_input":"2021-05-20T07:23:04.177204Z","iopub.status.idle":"2021-05-20T07:23:04.187405Z","shell.execute_reply.started":"2021-05-20T07:23:04.177168Z","shell.execute_reply":"2021-05-20T07:23:04.186491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Optimizer","metadata":{}},{"cell_type":"code","source":"from torch.optim.optimizer import Optimizer\nimport itertools as it\n\ndef get_learning_rate(optimizer):\n    lr=[]\n    for param_group in optimizer.param_groups:\n       lr +=[ param_group['lr'] ]\n\n    assert(len(lr)==1) #we support only one param_group\n    lr = lr[0]\n\n    return lr\n\nclass RAdam(Optimizer):\n\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)\n        self.buffer = [[None, None, None] for ind in range(10)]\n        super(RAdam, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(RAdam, self).__setstate__(state)\n\n    def step(self, closure=None):\n\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data.float()\n                if grad.is_sparse:\n                    raise RuntimeError('RAdam does not support sparse gradients')\n\n                p_data_fp32 = p.data.float()\n\n                state = self.state[p]\n\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n                else:\n                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n                beta1, beta2 = group['betas']\n\n                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value = 1 - beta2)\n                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n\n                state['step'] += 1\n                buffered = self.buffer[int(state['step'] % 10)]\n                if state['step'] == buffered[0]:\n                    N_sma, step_size = buffered[1], buffered[2]\n                else:\n                    buffered[0] = state['step']\n                    beta2_t = beta2 ** state['step']\n                    N_sma_max = 2 / (1 - beta2) - 1\n                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)\n                    buffered[1] = N_sma\n\n                    # more conservative since it's an approximated value\n                    if N_sma >= 5:\n                        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'])\n                    else:\n                        step_size = 1.0 / (1 - beta1 ** state['step'])\n                    buffered[2] = step_size\n\n                if group['weight_decay'] != 0:\n                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n\n                # more conservative since it's an approximated value\n                if N_sma >= 5:\n                    denom = exp_avg_sq.sqrt().add_(group['eps'])\n                    p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size * group['lr'])\n                else:\n                    p_data_fp32.add_(exp_avg, alpha=-step_size * group['lr'])\n\n                p.data.copy_(p_data_fp32)\n\n        return loss\n\nclass Lookahead(Optimizer):\n    def __init__(self, optimizer, alpha=0.5, k=6):\n\n        if not 0.0 <= alpha <= 1.0:\n            raise ValueError(f'Invalid slow update rate: {alpha}')\n        if not 1 <= k:\n            raise ValueError(f'Invalid lookahead steps: {k}')\n\n        self.optimizer = optimizer\n        self.param_groups = self.optimizer.param_groups\n        self.alpha = alpha\n        self.k = k\n        for group in self.param_groups:\n            group[\"step_counter\"] = 0\n\n        self.slow_weights = [\n                [p.clone().detach() for p in group['params']]\n            for group in self.param_groups]\n\n        for w in it.chain(*self.slow_weights):\n            w.requires_grad = False\n        self.state = optimizer.state\n\n    def step(self, closure=None):\n        loss = None\n        if closure is not None:\n            loss = closure()\n        loss = self.optimizer.step()\n\n        for group,slow_weights in zip(self.param_groups,self.slow_weights):\n            group['step_counter'] += 1\n            if group['step_counter'] % self.k != 0:\n                continue\n            for p,q in zip(group['params'],slow_weights):\n                if p.grad is None:\n                    continue\n                q.data.add_(p.data - q.data, alpha=self.alpha )\n                p.data.copy_(q.data)\n        return loss","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.188953Z","iopub.execute_input":"2021-05-20T07:23:04.189298Z","iopub.status.idle":"2021-05-20T07:23:04.214729Z","shell.execute_reply.started":"2021-05-20T07:23:04.189263Z","shell.execute_reply":"2021-05-20T07:23:04.213909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"def do_valid(net, tokenizer, valid_loader):\n\n    valid_probability = []\n    valid_truth = []\n    valid_length = []\n    valid_num = 0\n\n    net.eval()\n    start_timer = timer()\n    for t, batch in enumerate(valid_loader):\n        batch_size = len(batch['index'])\n        length = batch['length']\n        token  = batch['token' ].cuda()\n        token_pad_mask = batch['token_pad_mask' ].cuda()\n        #image  = batch['image' ].cuda()\n        num_patch = batch['num_patch']\n        patch  = batch['patch' ].cuda()\n        coord  = batch['coord' ].cuda()\n        patch_pad_mask  = batch['patch_pad_mask' ].cuda()\n\n        with torch.no_grad():\n            logit = data_parallel(net, (patch, coord, token, patch_pad_mask, token_pad_mask)) #net(image, token, length)\n            probability = F.softmax(logit,-1)\n\n        valid_num += batch_size\n        valid_probability.append(probability.data.cpu().numpy())\n        valid_truth.append(token.data.cpu().numpy())\n        valid_length.extend(length)\n        print('\\r %8d / %d  %s'%(valid_num, len(valid_loader.sampler),time_to_str(timer() - start_timer,'sec')),end='',flush=True)\n        del t, batch\n        \n    assert(valid_num == len(valid_loader.sampler)) #len(valid_loader.dataset))\n    #print('')\n    #----------------------\n    probability = np.concatenate(valid_probability)\n    predict = probability.argmax(-1)\n    truth   = np.concatenate(valid_truth)\n    length  = valid_length\n\n\n    #----\n    p = probability[:,:-1].reshape(-1,vocab_size)\n    t = truth[:,1:].reshape(-1)\n\n    non_pad = np.where(t!=STOI['<pad>'])[0] #& (t!=STOI['<sos>'])\n    p = p[non_pad]\n    t = t[non_pad]\n    loss = np_loss_cross_entropy(p, t)\n\n    #----\n    lb_score = 0\n    if 1:\n        score = []\n        for i,(p, t) in enumerate(zip(predict, truth)):\n            t = truth[i][1:length[i]-1]\n            p = predict[i][1:length[i]-1]\n            t = tokenizer.one_predict_to_inchi(t)\n            p = tokenizer.one_predict_to_inchi(p)\n            s = Levenshtein.distance(p, t)\n            score.append(s)\n        lb_score = np.mean(score)\n        \n    \"\"\"\n    if 1:\n        score = []\n        for i, (p, t) in enumerate(zip(predict, truth)):\n            t = truth[i][1:length[i]-1]     # in the buggy version, i have used 1 instead of i\n            p = predict[i][1:length[i]-1]\n            t = tokenizer.one_predict_to_inchi(t)\n            p = tokenizer.one_predict_to_inchi(p)\n            s = Levenshtein.distance(p, t)\n            score.append(s)\n        lb_score = np.mean(score)\n    \"\"\"\n\n    #lb_score = compute_lb_score(k, t)\n    del valid_loader, net, predict, truth,valid_probability, valid_truth, valid_length, valid_num\n    return [loss, lb_score]\n","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.216137Z","iopub.execute_input":"2021-05-20T07:23:04.216524Z","iopub.status.idle":"2021-05-20T07:23:04.232049Z","shell.execute_reply.started":"2021-05-20T07:23:04.216488Z","shell.execute_reply":"2021-05-20T07:23:04.231293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_train():\n\n    fold = 3\n    out_dir = \\\n        './tnt-patch1-s0.8/fold%d' % fold\n    initial_checkpoint = None\n    #initial_checkpoint = \\\n      #out_dir + '/checkpoint/00755000_model.pth'#None #\n       #'/root/share1/kaggle/2021/bms-moleular-translation/result/try22/tnt-patch1/fold3/checkpoint/00697000_model.pth'\n\n    debug = 0\n    start_lr = 0.00001# 1\n    batch_size = 32   # 24\n\n\n    ## setup  ----------------------------------------\n    for f in ['checkpoint', 'train', 'valid', 'backup']: os.makedirs(out_dir + '/' + f, exist_ok=True)\n    # backup_project_as_zip(PROJECT_PATH, out_dir +'/backup/code.train.%s.zip'%IDENTIFIER)\n\n    log = Logger()\n    log.open(out_dir + '/log.train.txt', mode='a')\n    log.write('\\n--- [START %s] %s\\n\\n' % (IDENTIFIER, '-' * 64))\n    log.write('\\t%s\\n' % COMMON_STRING)\n    #log.write('\\t__file__ = %s\\n' % __file__)\n    log.write('\\tout_dir  = %s\\n' % out_dir)\n    log.write('\\n')\n\n    ## dataset ------------------------------------\n\n    df_train, df_valid = make_fold('train-%d' % fold)\n    df_valid = df_valid.iloc[:5_000]\n\n\n    tokenizer = load_tokenizer()\n    train_dataset = BmsDataset(df_train,tokenizer)\n    valid_dataset = BmsDataset(df_valid,tokenizer)\n\n    train_loader = DataLoader(\n        train_dataset,\n        sampler = RandomSampler(train_dataset),\n        #sampler=UniformLengthSampler(train_dataset, is_shuffle=True), #200_000\n        batch_size=batch_size,\n        drop_last=True,\n        num_workers=8,\n        pin_memory=True,\n        worker_init_fn=lambda id: np.random.seed(torch.initial_seed() // 2 ** 32 + id),\n        collate_fn=null_collate,\n    )\n    valid_loader = DataLoader(\n        valid_dataset,\n        #sampler=UniformLengthSampler(valid_dataset, 5_000),\n        sampler=SequentialSampler(valid_dataset),\n        batch_size=32,\n        drop_last=False,\n        num_workers=8,\n        pin_memory=True,\n        collate_fn=null_collate,\n    )\n\n    log.write('train_dataset : \\n%s\\n' % (train_dataset))\n    log.write('valid_dataset : \\n%s\\n' % (valid_dataset))\n    log.write('\\n')\n\n    ## net ----------------------------------------\n    log.write('** net setting **\\n')\n    scaler = amp.GradScaler()\n    net = AmpNet().cuda()\n\n\n    if initial_checkpoint is not None:\n        f = torch.load(initial_checkpoint, map_location=lambda storage, loc: storage)\n        start_iteration = f['iteration']\n        start_epoch     = f['epoch']\n        state_dict = f['state_dict']\n\n        #---\n        # state_dict = {k.replace('cnn.e.','cnn.'):v for k,v in state_dict.items()}\n        # del state_dict['text_pos.pos']\n        # del state_dict['cnn.head.weight']\n        # del state_dict['cnn.head.bias']\n        # net.load_state_dict(state_dict, strict=False)\n\n        #---\n        net.load_state_dict(state_dict, strict=True)  # True\n    else:\n        start_iteration = 0\n        start_epoch = 0\n\n    log.write('\\tinitial_checkpoint = %s\\n' % initial_checkpoint)\n    log.write('\\n')\n\n    # -----------------------------------------------\n    if 0:  ##freeze\n        for p in net.encoder.parameters(): p.requires_grad = False\n\n    optimizer = Lookahead(RAdam(filter(lambda p: p.requires_grad, net.parameters()), lr=start_lr), alpha=0.5, k=5)\n    # optimizer = RAdam(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr)\n\n    num_iteration = 10 * 1000\n    #num_iteration = 10\n    \n    iter_log = 1000\n    iter_valid = 1000\n    iter_save = list(range(0, num_iteration, 1000))  # 1*1000\n    #iter_save = [0]\n    \n    log.write('optimizer\\n  %s\\n' % (optimizer))\n    log.write('\\n')\n\n    ## start training here! ##############################################\n    log.write('** start training here! **\\n')\n    log.write('   is_mixed_precision = %s \\n' % str(is_mixed_precision))\n    log.write('   batch_size = %d\\n' % (batch_size))\n    #log.write('   experiment = %s\\n' % str(__file__.split('/')[-2:]))\n    log.write('                      |----- VALID ---|---- TRAIN/BATCH --------------\\n')\n    log.write('rate     iter   epoch | loss  lb(lev) | loss0  loss1  | time          \\n')\n    log.write('----------------------------------------------------------------------\\n')\n             # 0.00000   0.00* 0.00  | 0.000  0.000  | 0.000  0.000  |  0 hr 00 min\n\n    def message(mode='print'):\n        if mode == ('print'):\n            asterisk = ' '\n            loss = batch_loss\n        if mode == ('log'):\n            asterisk = '*' if iteration in iter_save else ' '\n            loss = train_loss\n        \n        #'%4.3f  %5.2f  | ' % (*valid_loss,) + \\\n        text = \\\n            '%0.5f  %5.4f%s %4.2f  | ' % (rate, iteration / 10000 , asterisk, epoch,) + \\\n            '%4.3f  %4.3f  %4.3f  | ' % (*loss,) + \\\n            '%s' % (time_to_str(timer() - start_timer, 'min'))\n\n        return text\n\n    # ----\n    valid_loss = np.zeros(2, np.float32)\n    train_loss = np.zeros(3, np.float32)\n    batch_loss = np.zeros_like(train_loss)\n    sum_train_loss = np.zeros_like(train_loss)\n    sum_train = 0\n    loss0 = torch.FloatTensor([0]).cuda().sum()\n    loss1 = torch.FloatTensor([0]).cuda().sum()\n    loss2 = torch.FloatTensor([0]).cuda().sum()\n\n    start_timer = timer()\n    iteration = start_iteration\n    epoch = start_epoch\n    rate = 0\n    #while iteration < num_iteration:\n\n    for t, batch in enumerate(train_loader):\n        if iteration in iter_save:\n            if iteration != start_iteration:\n                torch.save({\n                    'state_dict': net.state_dict(),\n                    'iteration': iteration,\n                    'epoch': epoch,\n                }, out_dir + '/checkpoint/%08d_model.pth' % (iteration))\n                pass\n\n        \"\"\"\n        if (iteration % iter_valid == 0):\n            if iteration != start_iteration:\n                valid_loss = do_valid(net, tokenizer, valid_loader)  #\n                pass\n        \"\"\"            \n\n        if (iteration % iter_log == 0):\n            print('\\r', end='', flush=True)\n            log.write(message(mode='log') + '\\n')\n\n        # learning rate schduler ------------\n        rate = get_learning_rate(optimizer)\n\n        # one iteration update  -------------\n        batch_size = len(batch['index'])\n        length = batch['length']\n        token  = batch['token' ].cuda()\n        token_pad_mask = batch['token_pad_mask' ].cuda()\n        #image  = batch['image' ].cuda()\n        num_patch = batch['num_patch']\n        patch  = batch['patch' ].cuda()\n        coord  = batch['coord' ].cuda()\n        patch_pad_mask = batch['patch_pad_mask' ].cuda()\n\n\n        # ----\n        net.train()\n        optimizer.zero_grad()\n\n        if is_mixed_precision:\n            with amp.autocast():\n                #assert(False)\n                logit = data_parallel(net, (patch, coord, token, patch_pad_mask, token_pad_mask)) #net(image, token, length)\n                loss0 = seq_cross_entropy_loss(logit, token, length)\n                #loss0 = seq_anti_focal_cross_entropy_loss(logit, token, length)\n\n            scaler.scale(loss0).backward()\n            #scaler.unscale_(optimizer)\n            #torch.nn.utils.clip_grad_norm_(net.parameters(), 2)\n            scaler.step(optimizer)\n            scaler.update()\n\n\n        else:\n            assert False\n            # print('fp32')\n            # image_embed = encoder(image)\n            logit, weight = decoder(image_embed, token, length)\n\n            (loss0).backward()\n            optimizer.step()\n\n        # print statistics  --------\n        epoch += 1 / len(train_loader)\n        iteration += 1\n        batch_loss = np.array([loss0.item(), loss1.item(), loss2.item()])\n        sum_train_loss += batch_loss\n        sum_train += 1\n        if iteration % 100 == 0:\n            train_loss = sum_train_loss / (sum_train + 1e-12)\n            sum_train_loss[...] = 0\n            sum_train = 0\n\n        print('\\r', end='', flush=True)\n        print(message(mode='print'), end='', flush=True)\n\n        # debug--------------------------\n        if debug:\n            pass\n\n        # delete per batch\n        del batch_size, length, token ,token_pad_mask, num_patch, patch, coord, patch_pad_mask, t, batch\n        \n        if iteration > num_iteration:\n            break\n            \n    log.write('\\n')","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:29:16.994585Z","iopub.execute_input":"2021-05-20T07:29:16.994927Z","iopub.status.idle":"2021-05-20T07:29:17.024835Z","shell.execute_reply.started":"2021-05-20T07:29:16.994893Z","shell.execute_reply":"2021-05-20T07:29:17.023757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run_train()","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:29:20.012069Z","iopub.execute_input":"2021-05-20T07:29:20.012381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#!ls ./tnt-patch1-s0.8/fold3/checkpoint/\n\n#list(range(0, 3000, 1000))","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.267501Z","iopub.status.idle":"2021-05-20T07:23:04.268336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\nimport sys\n\n# These are the usual ipython objects, including this one you are creating\nipython_vars = ['In', 'Out', 'exit', 'quit', 'get_ipython', 'ipython_vars']\n\n# Get a sorted list of the objects and their sizes\nsorted([(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)\n\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-05-20T07:23:04.270064Z","iopub.status.idle":"2021-05-20T07:23:04.270748Z"},"trusted":true},"execution_count":null,"outputs":[]}]}