{"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":"# !pip install timm\n# !pip install opencv-python\n# !pip install albumentations","metadata":{"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !pip install -U efficientnet_pytorch --no-cache","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport glob\nimport gc\ngc.enable()\nimport multiprocessing\nimport cv2\nimport copy\nimport time\nimport random\nfrom PIL import Image\nfrom PIL import ImageFile\nImageFile.LOAD_TRUNCATED_IMAGES = True\n\nimport base64\nimport pickle\nimport urllib\nfrom urllib import request\nfrom urllib.request import urlopen\nimport uuid\n\n# fold\nfrom sklearn.model_selection import StratifiedKFold\n\n# For downloading images\nfrom io import BytesIO\n\n# For data manipulation\nimport numpy as np\nimport pandas as pd\n\n# Pytorch Imports\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda import amp\nimport torchvision\n\n\n# Utils\nimport joblib\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\n# Sklearn Imports\nfrom sklearn.metrics import mean_squared_error\nfrom sklearn.model_selection import StratifiedKFold, KFold\n\n# For Image Models\nimport timm\n\n# For Transformer Models\nfrom transformers import AutoTokenizer, AutoModel\nfrom transformers import BertTokenizer, BertModel, BertConfig\n\n# Albumentations for augmentations\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# For colored terminal text\nfrom colorama import Fore, Back, Style\nr_ = Fore.RED\nb_ = Fore.BLUE\nc_ = Fore.CYAN\ng_ = Fore.GREEN\ny_ = Fore.YELLOW\nm_ = Fore.MAGENTA\nsr_ = Style.RESET_ALL\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n# For descriptive error messages\n# os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\"","metadata":{"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from efficientnet_pytorch import EfficientNet","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"saved_model_path = './saved_models/version 1.0/'\ndataset_path = '../input/wikipedia-image-caption/'","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def optimal_num_of_loader_workers():\n    num_cpus = multiprocessing.cpu_count()\n    num_gpus = torch.cuda.device_count()\n    optimal_value = min(num_cpus, num_gpus*4) if num_gpus else num_cpus - 1\n    return optimal_value","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CONFIG = {\n    \"seed\": 2021,\n    \"epochs\": 5,#20,\n    'nfolds':5,\n    \n    \"img_size\": 600,\n    \"text_model_name\": \"xlm-roberta-base\",\n    \n    \"embedding_size\": 256,\n    \"train_batch_size\": 4,\n    \"valid_batch_size\": 4,\n    \"learning_rate\": 1e-4,\n    \"scheduler\": 'CosineAnnealingLR',\n    \"min_lr\": 1e-6,\n    'num_workers':optimal_num_of_loader_workers(),\n    \n    \"T_max\": 500,\n    \"weight_decay\": 1e-6,\n    \"max_length\": 32,\n    \n    \"n_accumulate\": 1,\n}\n\nCONFIG[\"tokenizer\"] = AutoTokenizer.from_pretrained(CONFIG['text_model_name'])","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:52:56.772193Z","iopub.status.busy":"2021-11-03T15:52:56.771839Z","iopub.status.idle":"2021-11-03T15:53:03.529559Z","shell.execute_reply":"2021-11-03T15:53:03.528472Z","shell.execute_reply.started":"2021-11-03T15:52:56.772146Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CONFIG","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:53:03.533986Z","iopub.status.busy":"2021-11-03T15:53:03.533504Z","iopub.status.idle":"2021-11-03T15:53:03.548425Z","shell.execute_reply":"2021-11-03T15:53:03.545692Z","shell.execute_reply.started":"2021-11-03T15:53:03.533931Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(seed=42):\n    '''Sets the seed of the entire notebook so results are the same every time we run.\n    This is for REPRODUCIBILITY.'''\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    # When running on the CuDNN backend, two further options must be set\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # Set a fixed value for the hash seed\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    \nset_seed(CONFIG['seed'])","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:53:03.551700Z","iopub.status.busy":"2021-11-03T15:53:03.550899Z","iopub.status.idle":"2021-11-03T15:53:03.565160Z","shell.execute_reply":"2021-11-03T15:53:03.563922Z","shell.execute_reply.started":"2021-11-03T15:53:03.551623Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_feathers = glob.glob('./data/train_feather_files/' + 'train*')\nprint(train_feathers)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.DataFrame()\nfor file in train_feathers:\n    df = pd.read_feather(file)\n    train_df = pd.concat([train_df, df])\nprint(\"Before removing duplicate rows:\",train_df.shape)\ntrain_df = train_df.drop_duplicates() #Drop duplicate rows if any\nprint(\"After removing duplicate rows:\",train_df.shape)\ntrain_df = train_df.sample(frac = 1).reset_index(drop = True)\nprint(\"Null:\", train_df.isnull().any().any())\n\ntrain_df.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(train_df['language'].unique()), train_df['language'].unique())","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(train_df['language'].value_counts())","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_df)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.at[0, 'image_url']","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def url_to_image(img_url, file_name):\n    try:\n        file_name = str(uuid.uuid4())\n        file_name = f'./data/train_images/{file_name}.jpg'\n        req = request.Request(img_url)\n        req.add_header('User-Agent', 'abc-bot')\n        response = request.urlopen(req)\n        f= open(file_name, 'wb')\n        f.write(response.read())\n        f.close()\n        img = Image.open(file_name).convert(\"RGB\")\n        os.remove(file_name)\n        return img\n    except:\n        return None","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_transforms = {\n    \"train\": A.Compose([\n            A.Resize(CONFIG['img_size'], CONFIG['img_size']),\n            A.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n                max_pixel_value=255.0, \n                p=1.0\n            ),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.Rotate(limit=180, p=0.7),\n            A.ShiftScaleRotate(\n                shift_limit = 0.1, scale_limit=0.1, rotate_limit=45, p=0.5\n            ),\n            A.HueSaturationValue(\n                hue_shift_limit=0.2, sat_shift_limit=0.2,\n                val_shift_limit=0.2, p=0.5\n            ),\n            A.RandomBrightnessContrast(\n                brightness_limit=(-0.1, 0.1),\n                contrast_limit=(-0.1, 0.1), p=0.5\n            ),\n            ToTensorV2(p=1.0),\n        ]),\n    \n    \"valid\": A.Compose([\n            A.Resize(CONFIG['img_size'], CONFIG['img_size']),\n            A.Normalize(\n                mean=[0.485, 0.456, 0.406], \n                std=[0.229, 0.224, 0.225], \n                max_pixel_value=255.0, \n                p=1.0\n            ),\n            ToTensorV2(p=1.0),\n        ])\n}","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:53:37.340533Z","iopub.status.busy":"2021-11-03T15:53:37.339803Z","iopub.status.idle":"2021-11-03T15:53:37.351876Z","shell.execute_reply":"2021-11-03T15:53:37.350739Z","shell.execute_reply.started":"2021-11-03T15:53:37.340483Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class WikipediaDataset(Dataset):\n    def __init__(self, data, tokenizer, max_length, transforms=None):\n        self.data = data.reset_index(drop=True)\n        self.max_len = max_length\n        self.tokenizer = tokenizer\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, index):\n#         image_bytes = base64.b64decode(self.data[index][\"b64_bytes\"])\n#         img = np.asarray(Image.open(BytesIO(image_bytes)).convert(\"RGB\"))\n        if torch.is_tensor(index):\n            index = index.tolist()\n        img = url_to_image(self.data.at[index, \"image_url\"], self.data.at[index, \"page_title\"])\n        while img == None:\n            index = random.randint(0, len(self.data)-1)\n            img = url_to_image(self.data.at[index, \"image_url\"], self.data.at[index, \"page_title\"])\n        img = np.array(img)\n        caption = random.choice(self.data.at[index, \"caption_title_and_reference_description\"])\n        caption = caption.replace(\"[SEP]\", \"</s>\") # sep token for xlm-roberta\n        inputs = self.tokenizer.encode_plus(\n                caption,\n                truncation=True,\n                add_special_tokens=True,\n                max_length=self.max_len,\n                padding='max_length'\n            )\n        \n        ids = inputs['input_ids']\n        mask = inputs['attention_mask']\n        \n        if self.transforms:\n            img = self.transforms(image=img)[\"image\"]\n        \n        return {\n            'ids': torch.tensor(ids, dtype=torch.long),\n            'mask': torch.tensor(mask, dtype=torch.long),\n            'image': img\n        }","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:53:37.315063Z","iopub.status.busy":"2021-11-03T15:53:37.314644Z","iopub.status.idle":"2021-11-03T15:53:37.337699Z","shell.execute_reply":"2021-11-03T15:53:37.336664Z","shell.execute_reply.started":"2021-11-03T15:53:37.315012Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ImageFeatureExtractor(nn.Module):\n    def __init__(self):\n        super(ImageFeatureExtractor, self).__init__()\n        self.model = EfficientNet.from_pretrained('efficientnet-b7')\n        del self.model._fc, self.model._dropout\n        self.n_feature_vector_size = 2560\n#         self.bn1 = nn.BatchNorm1d(2560)\n\n    def forward(self, inputs):\n        output = self.model.extract_features(inputs)\n        output = self.model._avg_pooling(output)\n        output = torch.flatten(output, start_dim=1)\n#         output = self.bn1(output)\n        return output\n    \n# h = ImageFeatureExtractor().to(device)\n# print(h(torch.randn(1,3,600,600).to(device)).shape)\n# del h\n# gc.collect()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TextExtractorModel(nn.Module):\n    def __init__(self, text_model):\n        super(TextExtractorModel, self).__init__()\n        self.text_model = AutoModel.from_pretrained(text_model)\n        self.text_fc = nn.Sequential(\n            nn.Linear(768, 1024),\n            nn.Dropout(p=0.2),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Linear(1024, 2560),\n            nn.BatchNorm1d(2560)\n        )\n        self.init_weights(self.text_fc)\n        \n    def init_weights(self, m):\n        if type(m) == torch.nn.Linear:\n            torch.nn.init.xavier_uniform_(m.weight)\n            m.bias.data.fill_(0)\n            \n    def forward(self, ids, mask):\n        out = self.text_model(input_ids=ids, attention_mask=mask, output_hidden_states=False)[1]\n        text_embeddings = self.text_fc(out)\n        return text_embeddings\n    ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Loss(nn.Module):\n    def __init__(self, eps=1e-6):\n        super().__init__()\n        self.mse = nn.MSELoss()\n        self.cosine = nn.CosineEmbeddingLoss()\n        self.eps = eps\n        \n    def forward(self, inputs, targets):\n        c_loss = self.cosine(inputs, targets, torch.Tensor(inputs.size(0)).to(device).fill_(1.0))\n        m_loss = torch.sqrt(self.mse(inputs, targets) + self.eps)\n        return 0.75*m_loss + 0.25*c_loss","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:54:31.363152Z","iopub.status.busy":"2021-11-03T15:54:31.362722Z","iopub.status.idle":"2021-11-03T15:54:31.368876Z","shell.execute_reply":"2021-11-03T15:54:31.367608Z","shell.execute_reply.started":"2021-11-03T15:54:31.363101Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def fetch_scheduler(optimizer):\n    if CONFIG['scheduler'] == 'CosineAnnealingLR':\n        scheduler = lr_scheduler.CosineAnnealingLR(optimizer,T_max=CONFIG['T_max'], eta_min=CONFIG['min_lr'])\n    elif CONFIG['scheduler'] == 'CosineAnnealingWarmRestarts':\n        scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer,T_0=CONFIG['T_0'], eta_min=CONFIG['min_lr'])\n    elif CONFIG['scheduler'] == None:\n        return None\n        \n    return scheduler","metadata":{"execution":{"iopub.execute_input":"2021-11-03T15:54:31.434457Z","iopub.status.busy":"2021-11-03T15:54:31.434022Z","iopub.status.idle":"2021-11-03T15:54:31.450688Z","shell.execute_reply":"2021-11-03T15:54:31.449460Z","shell.execute_reply.started":"2021-11-03T15:54:31.434393Z"}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(CONFIG['nfolds']):\n    os.makedirs(os.path.join(saved_model_path, f'model{i}'),exist_ok=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"results = {}\n\n# K-fold Cross Validation model evaluation\nprint('Total number of folds:', CONFIG['nfolds'])\nprint('-'*50)\n\n# kfold = KFold(n_splits=CONFIG['nfolds'], shuffle=True, random_state=CONFIG['seed'])\n\nkfold = StratifiedKFold(n_splits=CONFIG['nfolds'], shuffle=True, random_state=CONFIG['seed'])\nfor k , (train_idx,valid_idx) in enumerate(kfold.split(X=train_df, y=train_df['language'])):\n    train_df.loc[valid_idx,'Fold'] = k\n\n\n# for fold, (train_ids, test_ids) in enumerate(kfold.split(train_dataset)):\nfor fold in range(CONFIG['nfolds']):\n    if fold==0:\n        continue\n    print('FOLD', fold)\n    print('-'*50)\n#     train_subsampler = SubsetRandomSampler(train_ids)\n#     test_subsampler = SubsetRandomSampler(test_ids)\n#     trainloader = DataLoader(train_dataset, batch_size=CONFIG['train_batch_size'], num_workers=CONFIG['num_workers'], \n#                              shuffle=True, pin_memory=True, drop_last=True, sampler=train_subsampler)\n#     testloader = DataLoader(valid_dataset, batch_size=CONFIG['valid_batch_size'], num_workers=CONFIG['num_workers'], \n#                             shuffle=False, pin_memory=True, sampler=test_subsampler)\n\n    x_train, x_valid = train_df.query(f\"Fold != {fold}\"), train_df.query(f\"Fold == {fold}\")\n    train_dataset = WikipediaDataset(x_train, CONFIG[\"tokenizer\"], CONFIG[\"max_length\"], transforms=data_transforms[\"train\"])\n    valid_dataset = WikipediaDataset(x_valid, CONFIG[\"tokenizer\"], CONFIG[\"max_length\"], transforms=data_transforms[\"valid\"])\n    trainloader = DataLoader(train_dataset, batch_size=CONFIG['train_batch_size'], num_workers=CONFIG['num_workers'], \n                             shuffle=True, pin_memory=True, drop_last=True)\n    testloader = DataLoader(valid_dataset, batch_size=CONFIG['valid_batch_size'], num_workers=CONFIG['num_workers'], \n                            shuffle=False, pin_memory=True)\n    \n#     model = Model(FeatureExtractor, CONFIG['text_model_name']).to(device)\n    image_model = ImageFeatureExtractor().to(device)\n    text_model = TextExtractorModel(CONFIG['text_model_name']).to(device)\n    \n    criterion = Loss().to(device)\n#     optimizer = optim.AdamW(model.parameters(), lr=CONFIG['learning_rate'], weight_decay=CONFIG['weight_decay'],\n#                            betas=(0.9, 0.999), amsgrad=True)\n    optimizer = optim.AdamW([\n                {'params': image_model.parameters()},\n                {'params': text_model.parameters(), 'lr':1e-5}\n            ], lr=CONFIG['learning_rate'], weight_decay=CONFIG['weight_decay'], betas=(0.9, 0.999), amsgrad=True)\n    scheduler = fetch_scheduler(optimizer)\n    \n    best_val_loss = 1000\n    \n    #### TRAINING & EVALUATION  ####\n    for epoch in range(CONFIG['epochs']):\n        #### TRAINING ####\n#         model.train()\n#         model.freeze_backbone()\n        image_model.train(); text_model.train();\n        current_loss = 0.0\n        t = tqdm(enumerate(trainloader), total=len(trainloader), desc=\"Train: \")\n        for batch_id, data in t:\n            optimizer.zero_grad()\n            ids = data['ids'].to(device, dtype = torch.long)\n            mask = data['mask'].to(device, dtype = torch.long)\n            images = data['image'].to(device, dtype=torch.float)\n            \n            image_outputs = image_model(images)\n            text_outputs = text_model(ids, mask)\n            loss = criterion(image_outputs, text_outputs)\n            loss.backward()\n            optimizer.step()\n            current_loss += float(loss.item())\n            t.set_postfix_str('Training Loss='+str(round(current_loss/(batch_id+1), 4)))\n            \n            if scheduler is not None:\n                scheduler.step()\n        \n        #### EVALUATION ####\n        files = glob.glob('./data/train_images/*')\n        for f in files:\n            os.remove(f)\n        with torch.no_grad():\n#             model.eval()\n            image_model.eval(); text_model.eval();\n            current_loss = 0.0\n            t = tqdm(enumerate(testloader), total=len(testloader), desc=\"Val: \")\n            for batch_id, data in t:\n                ids = data['ids'].to(device, dtype = torch.long)\n                mask = data['mask'].to(device, dtype = torch.long)\n                images = data['image'].to(device, dtype=torch.float)\n\n#                 image_outputs, text_outputs = model(images, ids, mask)\n                image_outputs = image_model(images)\n                text_outputs = text_model(ids, mask)\n                loss = criterion(image_outputs, text_outputs)\n                current_loss += float(loss.item())\n                t.set_postfix_str('Val Loss='+str(round(current_loss/(batch_id+1), 4)))\n                \n            if (current_loss/len(testloader))<=best_val_loss:\n                print(f'loss has been decreased from {best_val_loss} to {(current_loss/len(testloader))}')\n                best_val_loss = (current_loss/len(testloader))\n                results[fold] = best_val_loss\n                torch.save(image_model.state_dict(), os.path.join(saved_model_path, f'model{fold}/image_model{fold}.bin'))\n                torch.save(text_model.state_dict(), os.path.join(saved_model_path, f'model{fold}/text_model{fold}.bin'))\n                CONFIG[\"tokenizer\"].save_pretrained(os.path.join(saved_model_path, f'model{fold}'))\n                \n    torch.cuda.empty_cache()\n    del image_model, text_model, trainloader, testloader, optimizer, scheduler\n    gc.collect()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import scipy.io as sio\n\ncaptions_df = pd.read_csv(os.path.join(dataset_path, 'test_caption_list.csv'))\n# print(captions_df.shape)\ncaptions = captions_df['caption_title_and_reference_description'].tolist()  \n# print(captions)\ncaptions = [caption.replace(\"[SEP]\", \"</s>\") for caption in captions]\n# print(captions)\n\nfor fold in range(CONFIG['nfolds']):\n#     fold = 0 \n    caption_tokens = []\n    text_model = TextExtractorModel(CONFIG['text_model_name']).to(device)\n    text_model.load_state_dict(torch.load(os.path.join(saved_model_path, f'model{fold}/text_model{fold}.bin')))\n    with torch.no_grad():\n        text_model.eval()\n        for caption in captions:\n            inputs = CONFIG[\"tokenizer\"].encode_plus(\n                        caption,\n                        truncation=True,\n                        add_special_tokens=True,\n                        max_length=CONFIG[\"max_length\"],\n                        padding='max_length'\n                    )\n            ids = torch.tensor(inputs['input_ids'], dtype=torch.long).to(device).unsqueeze(0)\n            mask = torch.tensor(inputs['attention_mask'], dtype=torch.long).to(device).unsqueeze(0)\n            text_features = text_model(ids, mask)\n            text_features = text_features.cpu().detach().numpy()\n            caption_tokens.append(text_features.squeeze())\n    print('Fold:', fold, np.array(caption_tokens).shape)\n    sio.savemat(os.path.join(saved_model_path, f'model{fold}/text_embeddings{fold}.mat'), \n                {'text_embeddings': np.array(caption_tokens)})\n    del caption_tokens, text_model\n    gc.collect()\n    #         break","metadata":{},"execution_count":null,"outputs":[]}]}