{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd\nimport cv2\nimport os\nimport re\n\nimport torch\nimport torch.nn as nn\nimport torchvision\nimport torchvision.transforms as transforms\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom tqdm.auto import tqdm\ntqdm.pandas()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH = 16\nEPOCHS = 2\n\nLR = 0.0001\nIM_SIZE = 128\n\nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nTRAIN_LABELS_PATH = '../input/bms-molecular-translation/train_labels.csv'\nTRAIN_DIR = '../input/bms-molecular-translation/train/'\nTEST_DIR = '../input/bms-molecular-translation/test/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_labels = pd.read_csv('../input/bms-molecular-translation/train_labels.csv')\ntrain_labels.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_file_path(image_id):\n    return \"../input/bms-molecular-translation/train/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )\n\ndef get_test_file_path(image_id):\n    return \"../input/bms-molecular-translation/test/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )\n\ntrain_labels['file_path'] = train_labels['image_id'].progress_apply(get_train_file_path)\n\nprint(f'train.shape: {train_labels.shape}')\ntrain_labels.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get chemical formula\ntrain_labels['formula'] = train_labels['InChI'].progress_apply(lambda x: x.split('/')[1])\ntrain_labels.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get labels\nCL = [\"N\", \"Br\", \"I\", \"S\", \"Cl\", \"H\", \"C\", \"P\", \"O\", \"Si\", \"F\", \"B\"]\nNUM_CL = len(CL)\nNUM_CL","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Remove digits from chemical formula\ndef str_to_list(st):    \n    r2=[]\n    for e in ['Br', 'Cl', 'Si']:\n        if e in st:\n            r2.append(e)\n            st = st.replace(e,'')\n    return [chr for chr in st] + r2\n\n\ntrain_labels['f2'] = train_labels['formula'].progress_apply(lambda x: str_to_list(\"\".join(re.findall(\"[a-zA-Z]+\", x))))\ntrain_labels.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"md = dict(zip(CL, range(NUM_CL+1)))\nmd","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Replace 'lab' to digits\ntrain_labels['lab'] = train_labels['f2'].progress_apply(lambda x: \"|\".join([str(md[v]) for v in x]))\ntrain_labels.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Just for speed up\ntrain_df = train_labels[:1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"X_Train, Y_Train = train_df['file_path'].values, train_df['lab'].values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GetData(Dataset):\n    def __init__(self, Dir, FNames, Labels, Transform):\n        self.dir = Dir\n        self.fnames = FNames\n        self.transform = Transform\n        self.labels = Labels         \n        \n    def __len__(self):\n        return len(self.fnames)\n\n    def __getitem__(self, index):\n        img = cv2.imread(self.fnames[index])\n        img = cv2.resize(img, (IM_SIZE, IM_SIZE))\n        X = img/255.        \n    \n        if \"train\" in self.dir:              \n            y = self.labels[index]\n            y = y.split('|')\n            y = list(map(int, y))            \n            y = np.eye(NUM_CL, dtype='float')[y]                                    \n            y = y.sum(axis=0)\n\n            return self.transform(X), y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Transform = transforms.Compose(\n    [transforms.ToTensor()])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainset = GetData(TRAIN_DIR, X_Train, Y_Train, Transform)\ntrainloader = DataLoader(trainset, batch_size=BATCH, shuffle=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"next(iter(trainloader))[0].shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.resnet34()\nmodel.fc = nn.Linear(512, NUM_CL, bias=True)\nmodel = model.to(DEVICE)\n\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=LR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nfor epoch in range(EPOCHS):\n    tr_loss = 0.0\n\n    model = model.train()\n\n    for i, (images, labels) in enumerate(trainloader):        \n        images = images.to(DEVICE)\n        labels = labels.to(DEVICE)       \n        logits = model(images.float())       \n        loss = criterion(logits, labels)\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        tr_loss += loss.detach().item()\n    \n    model.eval()\n    print('Epoch: %d | Loss: %.4f'%(epoch, tr_loss / i))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}