├── .DS_Store ├── .gitignore ├── LICENSE ├── Notebooks ├── Create_mini-ImageNet-LT.ipynb ├── TailCalibX_Toy_Dataset.ipynb └── mini-imagenet-assets │ ├── all_classes.txt │ ├── equal_classes.txt │ ├── readme.md │ └── unequal_classes.txt ├── README.md ├── libs ├── .DS_Store ├── core │ ├── TailCalib.py │ ├── TailCalibX.py │ ├── ce.py │ ├── core_base.py │ ├── ecbd.py │ └── modals.py ├── data │ ├── .DS_Store │ ├── CIFAR100_LT │ │ ├── .gitattributes │ │ ├── Imb_0.01_phase_train_seed_1.data_target_tuple │ │ ├── Imb_0.01_phase_train_seed_2.data_target_tuple │ │ ├── Imb_0.01_phase_train_seed_3.data_target_tuple │ │ ├── Imb_0.02_phase_train_seed_1.data_target_tuple │ │ ├── Imb_0.02_phase_train_seed_2.data_target_tuple │ │ ├── Imb_0.02_phase_train_seed_3.data_target_tuple │ │ ├── Imb_0.1_phase_train_seed_1.data_target_tuple │ │ ├── Imb_0.1_phase_train_seed_2.data_target_tuple │ │ ├── Imb_0.1_phase_train_seed_3.data_target_tuple │ │ └── README.md │ ├── ImbalanceCIFAR.py │ └── dataloader.py ├── loss │ ├── CosineDistill.py │ └── SoftmaxLoss.py ├── models │ ├── CosineDotProductClassifier.py │ ├── DotProductClassifier.py │ ├── ResNet32Feature.py │ ├── ResNext50Feature.py │ ├── ResNextFeature.py │ └── ecbd_converter.py ├── samplers │ └── ClassAwareSampler.py └── utils │ ├── default_config.yaml │ ├── experiments_maker.py │ ├── globals.py │ ├── logger.py │ └── utils.py ├── main.py ├── readme_assets ├── long_tail-buzz_ss.png ├── method.svg └── toy_example_output.svg ├── run_TailCalibX_CIFAR100-LT.sh ├── run_TailCalibX_mini-imagenet-LT.sh ├── run_all_CIFAR100-LT.sh ├── run_all_mini-ImageNet-LT.sh └── tailcalib_pip ├── LICENSE ├── README.md ├── setup.py └── src └── tailcalib.py /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rahulvigneswaran/TailCalibX/b836fd93925efb48f610c08bd56bb0c5a3fca9b0/.DS_Store -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | logs 3 | run_logs 4 | logs_important 5 | logs_old1 6 | logs_all 7 | wandb 8 | datasets 9 | extra 10 | garbage 11 | datasets 12 | # Notebooks 13 | .vscode 14 | *.csv 15 | *.pt 16 | *.pth 17 | *.txt 18 | # *.png 19 | NotUsedAnymore 20 | .ipynb_checkpoints 21 | *.pdf 22 | # *.yaml 23 | .vscode -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Rahul Vigneswaran 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Notebooks/Create_mini-ImageNet-LT.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "# Imports\n", 10 | "import numpy as np\n", 11 | "from sklearn.model_selection import train_test_split\n", 12 | "from os import walk\n", 13 | "\n", 14 | "# inits\n", 15 | "output = f\"libs/data/mini-imagenet/\" # Should be the place you want to output the .txt file which the dataloader will read later.\n", 16 | "imb_ratio = 0.01 # Choose Imbalance ratio here\n", 17 | "root = \"/home/rahul_intern/Imagenet/mini_imagenet\" # Should point to the dir in which your ImagNet is.\n", 18 | "\n", 19 | "# Split ratios \n", 20 | "train_split_ratio = 0.8\n", 21 | "test_split_ratio = 1 - train_split_ratio\n", 22 | "val_split_ratio = 0.2\n", 23 | "\n", 24 | "final = []\n", 25 | "labels = []\n", 26 | "final_1 = []\n", 27 | "labels_1 = []\n", 28 | "final_2 = []\n", 29 | "labels_2 = []\n", 30 | "\n", 31 | "all_classes = []\n", 32 | "equal_classes = []\n", 33 | "unequal_classes = []\n", 34 | "\n", 35 | "mini_keys = ['n02110341', 'n01930112', 'n04509417', 'n04067472', 'n04515003', 'n02120079', 'n03924679', 'n02687172', 'n03075370', 'n07747607', 'n09246464', 'n02457408', 'n04418357', 'n03535780', 'n04435653', 'n03207743', 'n04251144', 'n03062245', 'n02174001', 'n07613480', 'n03998194', 'n02074367', 'n04146614', 'n04243546', 'n03854065', 'n03838899', 'n02871525', 'n03544143', 'n02108089', 'n13133613', 'n03676483', 'n03337140', 'n03272010', 'n01770081', 'n09256479', 'n02091244', 'n02116738', 'n04275548', 'n03773504', 'n02606052', 'n03146219', 'n04149813', 'n07697537', 'n02823428', 'n02089867', 'n03017168', 'n01704323', 'n01532829', 'n03047690', 'n03775546', 'n01843383', 'n02971356', 'n13054560', 'n02108551', 'n02101006', 'n03417042', 'n04612504', 'n01558993', 'n04522168', 'n02795169', 'n06794110', 'n01855672', 'n04258138', 'n02110063', 'n07584110', 'n02091831', 'n03584254', 'n03888605', 'n02113712', 'n03980874', 'n02219486', 'n02138441', 'n02165456', 'n02108915', 'n03770439', 'n01981276', 'n03220513', 'n02099601', 'n02747177', 'n01749939', 'n03476684', 'n02105505', 'n02950826', 'n04389033', 'n03347037', 'n02966193', 'n03127925', 'n03400231', 'n04296562', 'n03527444', 'n04443257', 'n02443484', 'n02114548', 'n04604644', 'n01910747', 'n04596742', 'n02111277', 'n03908618', 'n02129165', 'n02981792']\n", 36 | "\n", 37 | "with open(\"/home/rahul_intern/fb_dl_fresh/long_tail/libs/data/mini-imagenet/all_classes.txt\", \"r\") as f:\n", 38 | " for line in f:\n", 39 | " all_classes.append(str(line.strip()))\n", 40 | "\n", 41 | "with open(\"/home/rahul_intern/fb_dl_fresh/long_tail/libs/data/mini-imagenet/equal_classes.txt\", \"r\") as f:\n", 42 | " for line in f:\n", 43 | " equal_classes.append(str(line.strip()))\n", 44 | "\n", 45 | "with open(\"/home/rahul_intern/fb_dl_fresh/long_tail/libs/data/mini-imagenet/unequal_classes.txt\", \"r\") as f:\n", 46 | " for line in f:\n", 47 | " unequal_classes.append(str(line.strip()))\n", 48 | "\n", 49 | "filenames = next(walk(root), (None, None, []))[2] \n", 50 | "for name in filenames:\n", 51 | " label_temp = name.split(\"_\")[0]\n", 52 | " if label_temp in all_classes:\n", 53 | " final.append(name)\n", 54 | " labels.append(label_temp)\n", 55 | " if label_temp in equal_classes:\n", 56 | " final_1.append(name)\n", 57 | " labels_1.append(label_temp)\n", 58 | " else:\n", 59 | " final_2.append(name)\n", 60 | " labels_2.append(label_temp)\n", 61 | "\n", 62 | "actual_label = np.unique(labels)\n", 63 | "pseudo_label = np.arange(len(np.unique(labels)))\n", 64 | "\n", 65 | "# Converts the labels to range of 0 to max ints\n", 66 | "label_dict = {}\n", 67 | "inverse_label_dict = {}\n", 68 | "\n", 69 | "for i,j in zip(actual_label, pseudo_label):\n", 70 | " label_dict[i] = j\n", 71 | " inverse_label_dict[j] = i\n", 72 | " \n", 73 | "# Re-splitting the mini-imagenet which was made for few-shot into proper train, val, test sets.\n", 74 | "train_x, test_x, train_y, test_y = train_test_split(final_1, labels_1, train_size=train_split_ratio, test_size=test_split_ratio, stratify=labels_1)\n", 75 | "train_x, val_x, train_y, val_y = train_test_split(train_x, train_y, train_size=1-val_split_ratio, test_size=val_split_ratio, stratify=train_y)\n", 76 | "\n", 77 | "# Limiting train, val, test datapoints to 500, 100, 100 per class (Not a very clean code but gets the job done)\n", 78 | "def select_random_data(train_x, train_y, count=500):\n", 79 | " train_classwise_dict = {}\n", 80 | " for i, j in zip(train_y, train_x):\n", 81 | " if i in train_classwise_dict.keys():\n", 82 | " train_classwise_dict[i].extend([j])\n", 83 | " else:\n", 84 | " train_classwise_dict[i] = []\n", 85 | " train_classwise_dict[i].extend([j])\n", 86 | "\n", 87 | " new_train_x = []\n", 88 | " new_train_y = []\n", 89 | " for i in train_classwise_dict.keys():\n", 90 | " ind1 = np.random.permutation(len(train_classwise_dict[i]))[:count]\n", 91 | " new_train_x.append(list(np.array(train_classwise_dict[i])[ind1]))\n", 92 | " new_train_y.append([i]*count)\n", 93 | " return sum(new_train_x, []), sum(new_train_y, [])\n", 94 | "\n", 95 | "train_x, train_y = select_random_data(train_x, train_y, 500)\n", 96 | "val_x, val_y = select_random_data(val_x, val_y, 100)\n", 97 | "test_x, test_y = select_random_data(test_x, test_y, 100)\n", 98 | "\n", 99 | "# Randomly select and limit datapoints per class from unequal_classes, divide them into train, val, test and append it to the already limited and divided train, val, test of equal_classes\n", 100 | "train_classwise_dict = {}\n", 101 | "for i, j in zip(labels_2, final_2):\n", 102 | " if i in train_classwise_dict.keys():\n", 103 | " train_classwise_dict[i].extend([j])\n", 104 | " else:\n", 105 | " train_classwise_dict[i] = []\n", 106 | " train_classwise_dict[i].extend([j])\n", 107 | "\n", 108 | "new_train_x = []\n", 109 | "new_train_y = []\n", 110 | "new_val_x = []\n", 111 | "new_val_y = []\n", 112 | "new_test_x = []\n", 113 | "new_test_y = []\n", 114 | "for i in train_classwise_dict.keys():\n", 115 | " ind1 = np.random.permutation(len(train_classwise_dict[i]))[:500]\n", 116 | " ind2 = np.random.permutation(len(train_classwise_dict[i]))[500:600]\n", 117 | " ind3 = np.random.permutation(len(train_classwise_dict[i]))[600:700]\n", 118 | " new_train_x.append(list(np.array(train_classwise_dict[i])[ind1]))\n", 119 | " new_val_x.append(list(np.array(train_classwise_dict[i])[ind2]))\n", 120 | " new_test_x.append(list(np.array(train_classwise_dict[i])[ind3]))\n", 121 | " new_train_y.append([i]*500)\n", 122 | " new_val_y.append([i]*100)\n", 123 | " new_test_y.append([i]*100)\n", 124 | "\n", 125 | "train_x.extend(sum(new_train_x, []))\n", 126 | "train_y.extend(sum(new_train_y, []))\n", 127 | "val_x.extend(sum(new_val_x, []))\n", 128 | "val_y.extend(sum(new_val_y, []))\n", 129 | "test_x.extend(sum(new_test_x, []))\n", 130 | "test_y.extend(sum(new_test_y, []))\n", 131 | "\n", 132 | "\n", 133 | "# print(np.unique(train_y, return_counts=True)[1], len(np.unique(train_y, return_counts=True)[1]))\n", 134 | "# print(np.unique(val_y, return_counts=True)[1], len(np.unique(val_y, return_counts=True)[1]))\n", 135 | "# print(np.unique(test_y, return_counts=True)[1], len(np.unique(test_y, return_counts=True)[1]))\n", 136 | "\n", 137 | "# Making Imbalanced train data\n", 138 | "def get_img_num_per_cls(cls_num, imb_type, imb_factor, data_length):\n", 139 | " img_max = data_length / cls_num\n", 140 | " img_num_per_cls = []\n", 141 | " if imb_type == \"exp\":\n", 142 | " for cls_idx in range(cls_num):\n", 143 | " num = img_max * (imb_factor ** (cls_idx / (cls_num - 1.0)))\n", 144 | " img_num_per_cls.append(int(num))\n", 145 | " elif imb_type == \"step\":\n", 146 | " for cls_idx in range(cls_num // 2):\n", 147 | " img_num_per_cls.append(int(img_max))\n", 148 | " for cls_idx in range(cls_num // 2):\n", 149 | " img_num_per_cls.append(int(img_max * imb_factor))\n", 150 | " else:\n", 151 | " img_num_per_cls.extend([int(img_max)] * cls_num)\n", 152 | " return img_num_per_cls\n", 153 | "\n", 154 | "def gen_imbalanced_data(img_num_per_cls, data, targets):\n", 155 | " new_data = []\n", 156 | " new_targets = []\n", 157 | " targets_np = np.array(targets, dtype=np.int64)\n", 158 | " data = np.array(data)\n", 159 | " classes = np.unique(targets_np)\n", 160 | "\n", 161 | " num_per_cls_dict = dict()\n", 162 | " for the_class, the_img_num in zip(classes, img_num_per_cls):\n", 163 | " num_per_cls_dict[the_class] = the_img_num\n", 164 | " idx = np.where(targets_np == the_class)[0]\n", 165 | " np.random.shuffle(idx)\n", 166 | " selec_idx = idx[:the_img_num]\n", 167 | " new_data.extend(data[selec_idx, ...])\n", 168 | " new_targets.extend([the_class,]* the_img_num)\n", 169 | " \n", 170 | " # print(len(new_data[-1]))\n", 171 | " # new_data = np.stack(new_data)\n", 172 | " return new_data, new_targets\n", 173 | "\n", 174 | "# Convert WordNetID labels to a range from 0 to 100\n", 175 | "train_y = [label_dict[i] for i in train_y]\n", 176 | "val_y = [label_dict[i] for i in val_y]\n", 177 | "test_y = [label_dict[i] for i in test_y]\n", 178 | "\n", 179 | "img_num_per_cls = get_img_num_per_cls(100, \"exp\", imb_ratio, len(train_x))\n", 180 | "train_x, train_y = gen_imbalanced_data(img_num_per_cls, train_x, train_y)\n", 181 | "\n", 182 | "# Writing as txt into \"output\" dir in inits\n", 183 | "dataxy = [(train_x, train_y), (val_x, val_y), (test_x, test_y)]\n", 184 | "for i, j in enumerate([\"train\", \"val\", \"test\"]):\n", 185 | " with open(f'{output}{imb_ratio}_{j}.txt', 'w') as f:\n", 186 | " for line, lab in zip(dataxy[i][0], dataxy[i][1]):\n", 187 | " f.write(line + \" \" + str(lab))\n", 188 | " f.write('\\n')\n", 189 | "\n", 190 | "for phase in [\"train\",\"val\", \"test\"]:\n", 191 | " finals = []\n", 192 | " labels = []\n", 193 | " input = f\"{output}{imb_ratio}_{phase}.txt\"\n", 194 | " with open(input) as f:\n", 195 | " for line in f:\n", 196 | " finals.append(line.split()[0])\n", 197 | " labels.append(line.split()[-1])\n", 198 | "\n", 199 | "\n", 200 | " print(np.unique(labels, return_counts=True)[1])\n", 201 | " max_val = max(np.unique(labels, return_counts=True)[1])\n", 202 | " min_val = min(np.unique(labels, return_counts=True)[1])\n", 203 | " sum_val = sum(np.unique(labels, return_counts=True)[1])\n", 204 | " cls_count = len(np.unique(labels, return_counts=True)[1])\n", 205 | " print(f\"{phase} -> Max: {max_val} | Min: {min_val} | Sum: {sum_val} | Imb: {max_val/min_val} | Class count: {cls_count}\")" 206 | ] 207 | } 208 | ], 209 | "metadata": { 210 | "language_info": { 211 | "name": "python" 212 | }, 213 | "orig_nbformat": 4 214 | }, 215 | "nbformat": 4, 216 | "nbformat_minor": 2 217 | } 218 | -------------------------------------------------------------------------------- /Notebooks/mini-imagenet-assets/all_classes.txt: -------------------------------------------------------------------------------- 1 | n02110341 2 | n01930112 3 | n04509417 4 | n04067472 5 | n04515003 6 | n02120079 7 | n03924679 8 | n02687172 9 | n03075370 10 | n07747607 11 | n09246464 12 | n02457408 13 | n04418357 14 | n03535780 15 | n04435653 16 | n03207743 17 | n04251144 18 | n03062245 19 | n02174001 20 | n07613480 21 | n03998194 22 | n02074367 23 | n04146614 24 | n04243546 25 | n03854065 26 | n03838899 27 | n02871525 28 | n03544143 29 | n02108089 30 | n13133613 31 | n03676483 32 | n03337140 33 | n03272010 34 | n01770081 35 | n09256479 36 | n02091244 37 | n02116738 38 | n04275548 39 | n03773504 40 | n02606052 41 | n03146219 42 | n04149813 43 | n07697537 44 | n02823428 45 | n02089867 46 | n03017168 47 | n01704323 48 | n01532829 49 | n03047690 50 | n03775546 51 | n01843383 52 | n02971356 53 | n13054560 54 | n02108551 55 | n02101006 56 | n03417042 57 | n04612504 58 | n01558993 59 | n04522168 60 | n02795169 61 | n06794110 62 | n01855672 63 | n04258138 64 | n02110063 65 | n07584110 66 | n02091831 67 | n03584254 68 | n03888605 69 | n02113712 70 | n03980874 71 | n02219486 72 | n02138441 73 | n02165456 74 | n02108915 75 | n03770439 76 | n01981276 77 | n03220513 78 | n02099601 79 | n02747177 80 | n01749939 81 | n03476684 82 | n02105505 83 | n02950826 84 | n04389033 85 | n03347037 86 | n02966193 87 | n03127925 88 | n03400231 89 | n04296562 90 | n03527444 91 | n04443257 92 | n02443484 93 | n02114548 94 | n04604644 95 | n01910747 96 | n04596742 97 | n02111277 98 | n03908618 99 | n02129165 100 | n02981792 101 | -------------------------------------------------------------------------------- /Notebooks/mini-imagenet-assets/equal_classes.txt: -------------------------------------------------------------------------------- 1 | n01532829 2 | n01558993 3 | n01749939 4 | n01770081 5 | n01843383 6 | n01855672 7 | n01910747 8 | n01930112 9 | n01981276 10 | n02074367 11 | n02091244 12 | n02091831 13 | n02099601 14 | n02101006 15 | n02105505 16 | n02108089 17 | n02108551 18 | n02108915 19 | n02110063 20 | n02110341 21 | n02111277 22 | n02113712 23 | n02114548 24 | n02116738 25 | n02120079 26 | n02129165 27 | n02138441 28 | n02165456 29 | n02174001 30 | n02219486 31 | n02443484 32 | n02457408 33 | n02687172 34 | n02747177 35 | n02795169 36 | n02823428 37 | n02871525 38 | n02950826 39 | n02966193 40 | n02971356 41 | n02981792 42 | n03017168 43 | n03047690 44 | n03127925 45 | n03146219 46 | n03207743 47 | n03220513 48 | n03272010 49 | n03337140 50 | n03347037 51 | n03417042 52 | n03476684 53 | n03527444 54 | n03535780 55 | n03544143 56 | n03584254 57 | n03676483 58 | n03770439 59 | n03773504 60 | n03775546 61 | n03838899 62 | n03854065 63 | n03888605 64 | n03908618 65 | n03924679 66 | n03980874 67 | n03998194 68 | n04067472 69 | n04146614 70 | n04149813 71 | n04243546 72 | n04251144 73 | n04258138 74 | n04275548 75 | n04296562 76 | n04389033 77 | n04435653 78 | n04509417 79 | n04515003 80 | n04522168 81 | n04596742 82 | n04604644 83 | n06794110 84 | n07613480 85 | n07697537 86 | n07747607 87 | n09246464 88 | n09256479 89 | n13054560 90 | n13133613 91 | -------------------------------------------------------------------------------- /Notebooks/mini-imagenet-assets/readme.md: -------------------------------------------------------------------------------- 1 | This folder hass all the required assets for generating mini-imagene-lt on your own! 2 | If you face any issue, raise an issue. 3 | -------------------------------------------------------------------------------- /Notebooks/mini-imagenet-assets/unequal_classes.txt: -------------------------------------------------------------------------------- 1 | n01704323 2 | n02089867 3 | n02606052 4 | n03062245 5 | n03075370 6 | n03400231 7 | n04418357 8 | n04443257 9 | n04612504 10 | n07584110 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TailCalibX : Feature Generation for Long-tail Classification 2 | by [Rahul Vigneswaran](https://rahulvigneswaran.github.io/), [Marc T. Law](http://www.cs.toronto.edu/~law/), [Vineeth N. Balasubramanian](https://lab1055.github.io/), [Makarand Tapaswi](https://makarandtapaswi.github.io/) 3 | 4 | [[arXiv](https://arxiv.org/abs/2111.05956#:~:text=The%20visual%20world%20naturally%20exhibits,models%20based%20on%20deep%20learning.)] [[Code](https://github.com/rahulvigneswaran/TailCalibX)] [[pip Package](https://pypi.org/project/tailcalib/0.0.1/)] [[Video](https://www.youtube.com/watch?v=BhOtYW2a_pU)] 5 | ![TailCalibX methodology](readme_assets/method.svg "TailCalibX methodology") 6 | 7 | 8 | # Table of contents 9 | - [🐣 Easy Usage (Recommended way to use our method)](#-easy-usage-recommended-way-to-use-our-method) 10 | - [💻 Installation](#-installation) 11 | - [👨‍💻 Example Code](#-example-code) 12 | - [🧪 Advanced Usage](#-advanced-usage) 13 | - [✔ Things to do before you run the code from this repo](#-things-to-do-before-you-run-the-code-from-this-repo) 14 | - [📀 How to use?](#-how-to-use) 15 | - [📚 How to create the mini-ImageNet-LT dataset?](#-how-to-create-the-mini-imagenet-lt-dataset) 16 | - [⚙ Arguments](#-arguments) 17 | - [🏋️‍♂️ Trained weights](#%EF%B8%8F%EF%B8%8F-trained-weights) 18 | - [🪀 Results on a Toy Dataset](#-results-on-a-toy-dataset) 19 | - [🌴 Directory Tree](#-directory-tree) 20 | - [📃 Citation](#-citation) 21 | - [👁 Contributing](#-contributing) 22 | - [❤ About me](#-about-me) 23 | - [✨ Extras](#-extras) 24 | - [📝 License](#-license) 25 | 26 | ## 🐣 Easy Usage (Recommended way to use our method) 27 | ⚠ **Caution:** TailCalibX is just TailCalib employed multiple times. Specifically, we generate a set of features once every epoch and use them to train the classifier. In order to mimic that, three things must be done at __every epoch__ in the following order: 28 | 1. Collect all the features from your dataloader. 29 | 2. Use the `tailcalib` package to make the features balanced by generating samples. 30 | 3. Train the classifier. 31 | 4. Repeat. 32 | 33 | ### 💻 Installation 34 | Use the package manager [pip](https://pip.pypa.io/en/stable/) to install __tailcalib__. 35 | 36 | ```bash 37 | pip install tailcalib 38 | ``` 39 | 40 | ### 👨‍💻 Example Code 41 | Check the instruction [here](tailcalib_pip/README.md) for a much more detailed python package information. 42 | 43 | ```python 44 | # Import 45 | from tailcalib import tailcalib 46 | 47 | # Initialize 48 | a = tailcalib(base_engine="numpy") # Options: "numpy", "pytorch" 49 | 50 | # Imbalanced random fake data 51 | import numpy as np 52 | X = np.random.rand(200,100) 53 | y = np.random.randint(0,10, (200,)) 54 | 55 | # Balancing the data using "tailcalib" 56 | feat, lab, gen = a.generate(X=X, y=y) 57 | 58 | # Output comparison 59 | print(f"Before: {np.unique(y, return_counts=True)}") 60 | print(f"After: {np.unique(lab, return_counts=True)}") 61 | ``` 62 | 63 | ## 🧪 Advanced Usage 64 | 65 | ### ✔ Things to do before you run the code from this repo 66 | - Change the `data_root` for your dataset in `main.py`. 67 | - If you are using wandb logging ([Weights & Biases](https://docs.wandb.ai/quickstart)), make sure to change the `wandb.init` in `main.py` accordingly. 68 | 69 | ### 📀 How to use? 70 | - For just the methods proposed in this paper : 71 | - For CIFAR100-LT: `run_TailCalibX_CIFAR100-LT.sh` 72 | - For mini-ImageNet-LT : `run_TailCalibX_mini-ImageNet-LT.sh` 73 | - For all the results show in the paper : 74 | - For CIFAR100-LT: `run_all_CIFAR100-LT.sh` 75 | - For mini-ImageNet-LT : `run_all_mini-ImageNet-LT.sh` 76 | 77 | ### 📚 How to create the mini-ImageNet-LT dataset? 78 | Check `Notebooks/Create_mini-ImageNet-LT.ipynb` for the script that generates the mini-ImageNet-LT dataset with varying imbalance ratios and train-test-val splits. 79 | ### ⚙ Arguments 80 | - `--seed` : Select seed for fixing it. 81 | - Default : `1` 82 | - `--gpu` : Select the GPUs to be used. 83 | - Default : `"0,1,2,3"` 84 | 85 | - `--experiment`: Experiment number (Check 'libs/utils/experiment_maker.py'). 86 | - Default : `0.1` 87 | - `--dataset` : Dataset number. 88 | - Choices : `0 - CIFAR100, 1 - mini-imagenet` 89 | - Default : `0` 90 | - `--imbalance` : Select Imbalance factor. 91 | - Choices : `0: 1, 1: 100, 2: 50, 3: 10` 92 | - Default : `1` 93 | - `--type_of_val` : Choose which dataset split to use. 94 | - Choices: `"vt": val_from_test, "vtr": val_from_train, "vit": val_is_test` 95 | - Default : `"vit"` 96 | 97 | - `--cv1` to `--cv9` : Custom variable to use in experiments - purpose changes according to the experiment. 98 | - Default : `"1"` 99 | 100 | - `--train` : Run training sequence 101 | - Default : `False` 102 | - `--generate` : Run generation sequence 103 | - Default : `False` 104 | - `--retraining` : Run retraining sequence 105 | - Default : `False` 106 | - `--resume` : Will resume from the 'latest_model_checkpoint.pth' and wandb if applicable. 107 | - Default : `False` 108 | 109 | - `--save_features` : Collect feature representations. 110 | - Default : `False` 111 | - `--save_features_phase` : Dataset split of representations to collect. 112 | - Choices : `"train", "val", "test"` 113 | - Default : `"train"` 114 | 115 | - `--config` : If you have a yaml file with appropriate config, provide the path here. Will override the 'experiment_maker'. 116 | - Default : `None` 117 | 118 | ## 🏋️‍♂️ Trained weights 119 | | Experiment | CIFAR100-LT (ResNet32, seed 1, Imb 100) | mini-ImageNet-LT (ResNeXt50)| 120 | | ----------- | ----------- | ----------- | 121 | | TailCalib | [Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/CIFAR100-LT/TailCalib) | [Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/mini-ImageNet-LT/TailCalib) | 122 | | TailCalibX | [Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/CIFAR100-LT/TailCalibX) |[Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/mini-ImageNet-LT/TailCalibX) | 123 | | CBD + TailCalibX | [Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/CIFAR100-LT/CBD%2BTailCalibX) |[Git-LFS](https://downgit.github.io/#/home?url=https://github.com/rahulvigneswaran/trained_models/tree/main/TailCalibX/mini-ImageNet-LT/CBD%2BTailCalibX) | 124 | 125 | ## 🪀 Results on a Toy Dataset 126 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1Yj2qymSm3NgCBqvKn5r_cOiEFl9wGp3J?usp=sharing) 127 | 128 | The higher the `Imb ratio`, the more imbalanced the dataset is. 129 | `Imb ratio = maximum_sample_count / minimum_sample_count`. 130 | 131 | Check [this notebook](https://colab.research.google.com/drive/1Yj2qymSm3NgCBqvKn5r_cOiEFl9wGp3J?usp=sharing) to play with the toy example from which the plot below was generated. 132 | ![](readme_assets/toy_example_output.svg) 133 | 134 | ## 🌴 Directory Tree 135 | ```bash 136 | TailCalibX 137 | ├── libs 138 | │ ├── core 139 | │ │ ├── ce.py 140 | │ │ ├── core_base.py 141 | │ │ ├── ecbd.py 142 | │ │ ├── modals.py 143 | │ │ ├── TailCalib.py 144 | │ │ └── TailCalibX.py 145 | │ ├── data 146 | │ │ ├── dataloader.py 147 | │ │ ├── ImbalanceCIFAR.py 148 | │ │ └── mini-imagenet 149 | │ │ ├── 0.01_test.txt 150 | │ │ ├── 0.01_train.txt 151 | │ │ └── 0.01_val.txt 152 | │ ├── loss 153 | │ │ ├── CosineDistill.py 154 | │ │ └── SoftmaxLoss.py 155 | │ ├── models 156 | │ │ ├── CosineDotProductClassifier.py 157 | │ │ ├── DotProductClassifier.py 158 | │ │ ├── ecbd_converter.py 159 | │ │ ├── ResNet32Feature.py 160 | │ │ ├── ResNext50Feature.py 161 | │ │ └── ResNextFeature.py 162 | │ ├── samplers 163 | │ │ └── ClassAwareSampler.py 164 | │ └── utils 165 | │ ├── Default_config.yaml 166 | │ ├── experiments_maker.py 167 | │ ├── globals.py 168 | │ ├── logger.py 169 | │ └── utils.py 170 | ├── LICENSE 171 | ├── main.py 172 | ├── Notebooks 173 | │ ├── Create_mini-ImageNet-LT.ipynb 174 | │ └── toy_example.ipynb 175 | ├── readme_assets 176 | │ ├── method.svg 177 | │ └── toy_example_output.svg 178 | ├── README.md 179 | ├── run_all_CIFAR100-LT.sh 180 | ├── run_all_mini-ImageNet-LT.sh 181 | ├── run_TailCalibX_CIFAR100-LT.sh 182 | └── run_TailCalibX_mini-imagenet-LT.sh 183 | ``` 184 | Ignored `tailcalib_pip` as it is for the `tailcalib` pip package. 185 | 186 | 187 | ## 📃 Citation 188 | ``` 189 | @inproceedings{rahul2021tailcalibX, 190 | title = {{Feature Generation for Long-tail Classification}}, 191 | author = {Rahul Vigneswaran and Marc T. Law and Vineeth N. Balasubramanian and Makarand Tapaswi}, 192 | booktitle = {ICVGIP}, 193 | year = {2021} 194 | } 195 | ``` 196 | 197 | ## 👁 Contributing 198 | Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change. 199 | 200 | ## ❤ About me 201 | [Rahul Vigneswaran](https://rahulvigneswaran.github.io/) 202 | 203 | ## ✨ Extras 204 | [🐝 Long-tail buzz](https://rahulvigneswaran.github.io/longtail-buzz/) : If you are interested in deep learning research which involves __long-tailed / imbalanced__ dataset, take a look at [Long-tail buzz](https://rahulvigneswaran.github.io/longtail-buzz/) to learn about the recent trending papers in this field. 205 | 206 | ![](/readme_assets/long_tail-buzz_ss.png) 207 | 208 | ## 📝 License 209 | [MIT](LICENSE) 210 | -------------------------------------------------------------------------------- /libs/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rahulvigneswaran/TailCalibX/b836fd93925efb48f610c08bd56bb0c5a3fca9b0/libs/.DS_Store -------------------------------------------------------------------------------- /libs/core/TailCalib.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from torch.optim import optimizer 3 | import os 4 | import copy 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from torch.utils.data import TensorDataset, DataLoader 10 | from tqdm import tqdm 11 | import numpy as np 12 | 13 | # Custom imports 14 | from libs.core.core_base import model as base_model 15 | from libs.utils.utils import * 16 | from libs.utils.logger import Logger 17 | import libs.utils.globals as g 18 | if g.wandb_log: 19 | import wandb 20 | 21 | class model(base_model): 22 | def batch_forward(self, inputs, retrain=False): 23 | """Batch Forward 24 | 25 | Args: 26 | inputs (float Tensor): batch_size x image_size 27 | retrain (bool, optional): Incase of retraining, different dataloaders are used. Defaults to False. 28 | """ 29 | 30 | # Calculate Features and outputs 31 | if not(retrain): 32 | self.features = self.networks["feat_model"](inputs) 33 | self.features = F.normalize(self.features, dim=1) 34 | else: 35 | self.features = inputs 36 | 37 | self.logits = self.networks["classifier"](self.features) 38 | 39 | def accumulate(self, phase): 40 | """Accumulates features of all the datapoints in a particular split 41 | 42 | Args: 43 | phase ([type]): Which split of dataset should be accumulated? 44 | """ 45 | print_str = ['Accumulating features: %s' % (phase)] 46 | print_write(print_str, self.log_file) 47 | time.sleep(0.25) 48 | 49 | torch.cuda.empty_cache() 50 | 51 | # In validation or testing mode, set model to eval() and initialize running loss/correct 52 | for model in self.networks.values(): 53 | model.eval() 54 | 55 | # Iterate over dataset 56 | self.feat = {} 57 | self.labs = {} 58 | 59 | accum_features = [] 60 | accum_labels = [] 61 | 62 | for inputs, labels, paths in tqdm(self.data[phase]): 63 | inputs, labels = inputs.cuda(), labels.cuda() 64 | # If on training phase, enable gradients 65 | with torch.set_grad_enabled(False): 66 | 67 | # In validation or testing 68 | self.batch_forward(inputs, labels, phase=phase) 69 | 70 | accum_features.append(self.features) 71 | accum_labels.append(labels) 72 | 73 | # Appending and stacking is same as just concatenating. In this you just dont have to declare a torch.empty of particular size. 74 | accum_features = torch.vstack(accum_features) 75 | accum_labels = torch.hstack(accum_labels) 76 | 77 | for i in accum_labels.unique().cpu().numpy(): 78 | self.feat[i] = accum_features[accum_labels == i] 79 | self.labs[i] = torch.full((self.feat[i].size()[0],), i).cuda() 80 | 81 | def generate_points(self, tailcalibX=False): 82 | """Generate new datapoints 83 | """ 84 | 85 | if not os.path.isdir(self.config["training_opt"]["log_generate"]): 86 | os.makedirs(self.config["training_opt"]["log_generate"]) 87 | else: 88 | if not(tailcalibX): 89 | raise Exception("Generation Directory already exists!!") 90 | 91 | # Class statistics 92 | self.base_means = [] 93 | self.base_covs = [] 94 | for i in self.feat.keys(): 95 | self.base_means.append(self.feat[i].mean(dim=0)) 96 | self.base_covs.append(self.get_cov(self.feat[i]).unsqueeze(dim=0)) 97 | 98 | self.base_means = torch.vstack(self.base_means) 99 | self.base_covs = torch.vstack(self.base_covs) 100 | 101 | # Tukey's transform 102 | if self.config["pg"]["tukey"]: 103 | for i in self.feat.keys(): 104 | self.feat[i] = self.tukey_transform(self.feat[i], self.config["pg"]["tukey_value"]) 105 | 106 | # Distribution calibration and feature sampling 107 | if self.config["pg"]["generate"]: 108 | sample_from_each = self.get_sample_count(self.config["training_opt"]["data_count"],self.feat.keys()) 109 | 110 | K = self.config["pg"]["topk"] 111 | self.generated_points = {} 112 | 113 | for i in tqdm(self.feat.keys()): 114 | if np.sum(sample_from_each[i]) == 0 and self.config["pg"]["extra_points"] == 0 : 115 | continue 116 | 117 | self.generated_points[i] = [] 118 | for k, x_ij in zip(sample_from_each[i], self.feat[i]): 119 | if k == 0: 120 | continue 121 | # Getting the top k nearest classes based on l2 distance 122 | distances = torch.cdist(self.base_means, x_ij.unsqueeze(0)).squeeze() 123 | topk_idx_nn_analysis = torch.topk(-distances, k=int(self.config["pg"]["nn_analysis_k"]))[1] 124 | topk_idx = topk_idx_nn_analysis[:K] 125 | 126 | # Calibrating mean and covariance 127 | calibrated_mean, calibrated_cov = self.calibrate_distribution(self.base_means[topk_idx], self.base_covs[topk_idx], K, x_ij, self.config["pg"]["alpha"]) 128 | 129 | # Trick to avoid cholesky decomposition from failing. Look at https://juanitorduz.github.io/multivariate_normal/ 130 | EPS = 1e-4 131 | calibrated_cov += (torch.eye(calibrated_cov.shape[0])*EPS).cuda() 132 | 133 | # Note that providng the scal_tril is faster than providing the covariance matrix directly. 134 | new_dist = torch.distributions.multivariate_normal.MultivariateNormal(loc=calibrated_mean, scale_tril=torch.linalg.cholesky(calibrated_cov)) 135 | 136 | gen = new_dist.sample((int(k),)) 137 | 138 | self.generated_points[i].append(gen) 139 | 140 | self.generated_points[i] = torch.vstack(self.generated_points[i]) 141 | torch.cuda.empty_cache() 142 | 143 | torch.save(self.generated_points, self.config["training_opt"]["log_generate"] + "/generated_points.pt") 144 | print("\nPoint Generation Completed!\n") 145 | 146 | else: 147 | self.generated_points ={} 148 | print("\nPoint Generation is False!\n") 149 | 150 | def get_cov(self, X): 151 | """Calculate the covariance matrix for X 152 | 153 | Args: 154 | X (torch.tensor): Features 155 | 156 | Returns: 157 | [torch.tensor]: Covariance matrix of X 158 | """ 159 | n = X.shape[0] 160 | mu = X.mean(dim=0) 161 | X = (X - mu) 162 | return 1/(n-1) * (X.transpose(0, 1) @ X) # X^TX -> feat_size x num_of_samples @ num_of_samples x feat_size -> feat_size x feat_size 163 | 164 | def get_sample_count(self, count, keys): 165 | """Decides how many samples must be generated based on each existing train datapoints. 166 | 167 | Args: 168 | count (list): Number of samples in each class 169 | keys (dict.keys): Class keys 170 | 171 | Returns: 172 | dict: dict consists that has the info as to how many samples must be generated based on each existing train datapoints. 173 | """ 174 | sample_count_dict = {} 175 | for i in keys: 176 | current = count[i] 177 | head = max(count) 178 | # head is the sample count that we must match after the generation. This can be offset by "self.config["pg"]["extra_points"]". In our experiments this is set to 0 as it worked better. 179 | num_sample = head - current + self.config["pg"]["extra_points"] 180 | ratio = num_sample / current 181 | # Makes sure each datapoint is being used atleast once 182 | new_sample_from_each = [np.floor(ratio)] * current 183 | 184 | # Rest of the datapoints used for generation are decided randomly 185 | while True: 186 | if sum(new_sample_from_each) == num_sample: 187 | break 188 | idx = np.random.randint(0, current) 189 | new_sample_from_each[idx] += 1 190 | 191 | # Sanity checks 192 | assert sum(new_sample_from_each) == num_sample 193 | assert len(new_sample_from_each) == current 194 | 195 | sample_count_dict[i] = new_sample_from_each 196 | 197 | return sample_count_dict 198 | 199 | 200 | def calibrate_distribution(self, base_means, base_cov, k, x_ij, alpha=0.0): 201 | """Calibration of the distribution for generation. Check equation 7 and 8 from our paper - Feature Generation for Long-tail Classification. 202 | 203 | Args: 204 | base_means (torch.tensor): List of all the means that are used for calibration. 205 | base_cov (torch.tensor): List of all the covariance matrices used for calibraton. 206 | k (int): Number of classes used for calibration. 207 | x_ij (torch.tensor): Datapoint chosen to be used for generation. 208 | alpha (float, optional): Decides the spread of the generated samples. Defaults to 0.0. 209 | 210 | Returns: 211 | torch.tensor : Calibrated mean and covariance matrix 212 | """ 213 | 214 | calibrated_mean = (base_means.sum(dim=0) + x_ij)/(k+1) 215 | calibrated_cov = base_cov.sum(dim=0)/k + alpha 216 | 217 | return calibrated_mean, calibrated_cov 218 | 219 | def tukey_transform(self, x, lam=0.2): 220 | """Transforms any distribution into normal-distribution like. 221 | 222 | Args: 223 | x (torch.tensor): Features 224 | lam (float, optional): Adjusts how close the transformed features will be to the origin. Defaults to 0.2. 225 | 226 | Returns: 227 | torch.tensor: Normal distribution like features. 228 | """ 229 | if lam == 0: 230 | EPS = 1e-6 231 | x = x + EPS 232 | return x.log() 233 | else : 234 | return x**lam 235 | 236 | def retrain(self,): 237 | """Creates a new dataloader, reinits everything and trains just the classifier part. 238 | """ 239 | # Prepare a dataloader for all splits which includes the generated points and is also tukey transformed. 240 | self.prepare_updated_dataset(include_generated_points = self.config["pg"]["generate"]) 241 | 242 | # Change log_dir so it wont change the weights of the parent model 243 | self.config["training_opt"]["log_dir"] = self.config["training_opt"]["log_retrain"] 244 | 245 | # Create retrain directory 246 | if not os.path.isdir(self.config["training_opt"]["log_dir"]): 247 | os.makedirs(self.config["training_opt"]["log_dir"]) 248 | else: 249 | raise Exception("Retrained Directory already exists!!") 250 | 251 | g.log_dir = self.config["training_opt"]["log_dir"] 252 | if g.log_offline: 253 | if not os.path.isdir(f"{g.log_dir}/metrics"): 254 | os.makedirs(f"{g.log_dir}/metrics") 255 | 256 | # Reinitialize everything 257 | print("Using steps for training.") 258 | self.training_data_num = len(self.my_dataloader["train"].dataset) 259 | self.epoch_steps = int( 260 | self.training_data_num / self.training_opt["batch_size"] 261 | ) 262 | 263 | # Init logger 264 | self.logger = Logger(self.training_opt["log_dir"]) 265 | self.log_file = os.path.join(self.training_opt["log_dir"], "log.txt") 266 | self.logger.log_cfg(self.config) 267 | 268 | # Initialize loss 269 | self.init_criterions() 270 | 271 | # Initialize model 272 | self.init_models() 273 | 274 | # Initialize model optimizer and scheduler 275 | print("Initializing model optimizer.") 276 | self.init_optimizers(self.model_optim_params_dict) 277 | 278 | self.train(retrain=True) 279 | 280 | def prepare_updated_dataset(self, include_generated_points = True): 281 | """Prepares a dataloader for all splits which includes the generated points and is also tukey transformed. 282 | 283 | 284 | Args: 285 | include_generated_points (bool, optional): Do you wanna include the newly generated points. Defaults to True. 286 | """ 287 | self.my_dataloader = {} 288 | self.for_distance_analysis = {} 289 | for phase in ["train", "val", "test"]: 290 | 291 | self.accumulate(phase=phase) 292 | 293 | feat_all = [] 294 | labs_all = [] 295 | 296 | if self.config["pg"]["tukey"]: 297 | for i in self.labs.keys(): 298 | self.feat[i] = self.tukey_transform(self.feat[i], lam=self.config["pg"]["tukey_value"]) 299 | 300 | for i in self.labs.keys(): 301 | feat_all.append(self.feat[i]) 302 | labs_all.append(self.labs[i]) 303 | 304 | feat_all = torch.vstack(feat_all) 305 | labs_all = torch.hstack(labs_all).cuda() 306 | 307 | if include_generated_points and phase == "train": 308 | generated_points = torch.load(self.config["training_opt"]["log_generate"] + "/generated_points.pt") 309 | for i in generated_points.keys(): 310 | feat_all = torch.cat((feat_all, generated_points[i].cuda())) 311 | labs_all = torch.cat((labs_all, torch.full((generated_points[i].size()[0],), int(i)).cuda())) 312 | 313 | # Create dataloader 314 | # Note that here we l2 normalize the features before creating the dataloader 315 | my_dataset = TensorDataset(F.normalize(feat_all, dim=1), labs_all, labs_all) 316 | self.my_dataloader[phase] = DataLoader(my_dataset, batch_size=self.config["training_opt"]["batch_size"], shuffle=True) 317 | 318 | # This is there so that we can use source_import from the utils to import model 319 | def get_core(*args): 320 | return model(*args) 321 | 322 | -------------------------------------------------------------------------------- /libs/core/TailCalibX.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from torch.optim import optimizer 3 | import os 4 | import copy 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from torch.utils.data import TensorDataset, DataLoader 10 | from tqdm import tqdm 11 | import numpy as np 12 | 13 | # Custom imports 14 | from libs.core.TailCalib import model as base_model # Note that we are trying to inherit the TailCalib (which in turn inherits core_base) instead of just core_base 15 | from libs.utils.utils import * 16 | from libs.utils.logger import Logger 17 | import libs.utils.globals as g 18 | if g.wandb_log: 19 | import wandb 20 | 21 | class model(base_model): 22 | def batch_forward(self, inputs): 23 | """Batch Forward 24 | 25 | Args: 26 | inputs (float Tensor): batch_size x image_size 27 | """ 28 | # Calculate Features and outputs 29 | if self.accumulation: 30 | self.features = self.networks["feat_model"](inputs) 31 | self.features = F.normalize(self.features, dim=1) 32 | else: 33 | self.features = inputs 34 | 35 | self.logits = self.networks["classifier"](self.features) 36 | 37 | def train(self, retrain=False): 38 | """Main training 39 | 40 | Args: 41 | retrain (bool, optional): Incase of retraining different dataloaders are used. Defaults to False. 42 | """ 43 | phase = "train" 44 | print_str = ["Phase: train"] 45 | print_write(print_str, self.log_file) 46 | 47 | # Inits 48 | best_acc = 0.0 49 | best_epoch = 0 50 | self.retrain = retrain 51 | self.end_epoch = self.training_opt["num_epochs"] 52 | self.accumulation = False 53 | self.accumulation_step = 1 54 | 55 | # Initialize best model and other variables 56 | self.best_model_weights = {} 57 | for key, _ in self.config["networks"].items(): 58 | if self.config["networks"][key]["trainable"]: 59 | self.best_model_weights[key] = copy.deepcopy(self.networks[key].state_dict()) 60 | 61 | # Loop over epochs 62 | for epoch in range(self.start_epoch, self.end_epoch + 1): 63 | # global config 64 | g.epoch_global = epoch 65 | 66 | # "Accumulate features -> Generate points -> Prepare a new dataloader" cycle. 67 | self.accumulate(phase="train") 68 | self.generate_points(tailcalibX=True) 69 | self.prepare_updated_dataset(include_generated_points = self.config["pg"]["generate"]) 70 | data_load = self.my_dataloader["train"] 71 | 72 | # Switch to train mode 73 | for key, model in self.networks.items(): 74 | if self.config["networks"][key]["trainable"]: 75 | # only train the module with lr > 0 76 | if self.config["networks"][key]["optim_params"]["lr"] == 0.0: 77 | model.eval() 78 | else: 79 | model.train() 80 | 81 | # Empty cuda cache 82 | torch.cuda.empty_cache() 83 | 84 | # Step the schedulers 85 | if self.model_scheduler_dict: 86 | for key, scheduler in self.model_scheduler_dict.items(): 87 | scheduler.step() 88 | if self.criterion_optimizer_scheduler: 89 | self.criterion_optimizer_scheduler.step() 90 | 91 | print_write([self.training_opt["log_dir"]], self.log_file) 92 | 93 | # print learning rate 94 | current_lr = self.show_current_lr() 95 | current_lr = min(current_lr * 50, 1.0) 96 | 97 | self.step = 0 98 | total_preds = [] 99 | total_labels = [] 100 | for inputs, labels, _ in data_load: 101 | # Break when step equal to epoch step 102 | if self.step == self.epoch_steps: 103 | break 104 | 105 | # Force shuffle option 106 | if self.do_shuffle: 107 | inputs, labels = self.shuffle_batch(inputs, labels) 108 | 109 | # Pushing to GPU 110 | inputs, labels = inputs, labels.cuda() 111 | 112 | with torch.set_grad_enabled(True): 113 | # If training, forward with loss, and no top 5 accuracy calculation 114 | self.batch_forward(inputs, labels, phase="train") 115 | self.batch_loss(labels) 116 | self.batch_backward() 117 | 118 | # Tracking and printing predictions 119 | _, preds = torch.max(self.logits, 1) 120 | total_preds.append(torch2numpy(preds)) 121 | total_labels.append(torch2numpy(labels)) 122 | 123 | # Output minibatch training results 124 | if self.step % self.training_opt['display_step'] == 0: 125 | 126 | minibatch_loss_classifier = self.loss_classifier.item() if 'ClassifierLoss' in self.criterions else None 127 | minibatch_loss_embed = self.loss_embed.item() if 'EmbeddingLoss' in self.criterions else None 128 | minibatch_loss_embed_proto = self.loss_embed_proto.item() if 'EmbeddingLoss' in self.criterions else None 129 | minibatch_loss_embed_biasreduc = self.loss_embed_biasreduc.item() if 'EmbeddingLoss' in self.criterions else None 130 | minibatch_loss_total = self.loss.item() 131 | minibatch_acc = mic_acc_cal(preds, labels) 132 | 133 | 134 | print_str = ['Epoch: [%d/%d]' 135 | % (epoch, self.training_opt['num_epochs']), 136 | 'Step: [%d/%d]' 137 | % (self.step, self.epoch_steps), 138 | 'Minibatch_loss_embedding: %.3f' 139 | % (minibatch_loss_embed) if minibatch_loss_embed else '', 140 | 'Minibatch_loss_classifier: %.3f' 141 | % (minibatch_loss_classifier) if minibatch_loss_classifier else '', 142 | 'Minibatch_accuracy_micro: %.3f' 143 | % (minibatch_acc)] 144 | print_write(print_str, self.log_file) 145 | 146 | loss_info = { 147 | 'epoch': epoch, 148 | 'Step': self.step, 149 | 'Total': minibatch_loss_total, 150 | 'Embedding (Total)': minibatch_loss_embed, 151 | 'Proto': minibatch_loss_embed_proto, 152 | 'BiasReduc': minibatch_loss_embed_biasreduc, 153 | 'Classifier': minibatch_loss_classifier, 154 | } 155 | 156 | self.logger.log_loss(loss_info) 157 | 158 | # wandb logging 159 | wandb_log({"Training Loss": minibatch_loss_total}) 160 | 161 | # batch-level: sampler update 162 | if hasattr(self.data["train"].sampler, "update_weights"): 163 | if hasattr(self.data["train"].sampler, "ptype"): 164 | ptype = self.data["train"].sampler.ptype 165 | else: 166 | ptype = "score" 167 | ws = get_priority(ptype, self.logits.detach(), labels) 168 | 169 | inlist = [indexes.cpu().numpy(), ws] 170 | if self.training_opt["sampler"]["type"] == "ClassPrioritySampler": 171 | inlist.append(labels.cpu().numpy()) 172 | self.data["train"].sampler.update_weights(*inlist) 173 | 174 | # Clear things out (optional) 175 | del inputs, labels, self.logits, self.features, preds 176 | 177 | # Update steps 178 | self.step+=1 179 | g.step_global += 1 180 | 181 | # epoch-level: reset sampler weight 182 | if hasattr(self.data["train"].sampler, "get_weights"): 183 | self.logger.log_ws(epoch, self.data["train"].sampler.get_weights()) 184 | if hasattr(self.data["train"].sampler, "reset_weights"): 185 | self.data["train"].sampler.reset_weights(epoch) 186 | 187 | # After every epoch, validation 188 | rsls = {'epoch': epoch} 189 | rsls_train = self.eval_with_preds(total_preds, total_labels) 190 | rsls_eval, _ , _ , _ = self.eval(phase='val') 191 | rsls.update(rsls_train) 192 | rsls.update(rsls_eval) 193 | 194 | # Reset class weights for sampling if pri_mode is valid 195 | if hasattr(self.data["train"].sampler, "reset_priority"): 196 | ws = get_priority( 197 | self.data["train"].sampler.ptype, 198 | self.total_logits.detach(), 199 | self.total_labels, 200 | ) 201 | self.data["train"].sampler.reset_priority( 202 | ws, self.total_labels.cpu().numpy() 203 | ) 204 | 205 | self.logger.log_acc(rsls) 206 | 207 | # # Under validation, the best model need to be updated 208 | if rsls_eval["val_all"] > best_acc: 209 | best_epoch = epoch 210 | best_acc = rsls_eval["val_all"] 211 | for key, _ in self.config["networks"].items(): 212 | if self.config["networks"][key]["trainable"]: 213 | self.best_model_weights[key] = copy.deepcopy(self.networks[key].state_dict()) 214 | 215 | # wandb log best epoch, train accuracy, based on best validation accuracy 216 | wandb_log({"Best Val": 100*best_acc, "Best Epoch": best_epoch}) 217 | wandb_log({"Best train": 100*rsls_train["train_all"], "Best Epoch": best_epoch}) 218 | 219 | wandb_log({'B_val_all': self.eval_acc_mic_top1, 220 | 'B_val_many': self.many_acc_top1, 221 | 'B_val_median': self.median_acc_top1, 222 | 'B_val_low': self.low_acc_top1}) 223 | 224 | wandb_log({'B_train_all': rsls_train["train_all"], 225 | 'B_train_many': rsls_train["train_many"], 226 | 'B_train_median': rsls_train["train_median"], 227 | 'B_train_low': rsls_train["train_low"]}) 228 | 229 | print("===> Saving checkpoint") 230 | self.save_latest(epoch) 231 | 232 | # Clear things out (optional) 233 | del rsls_eval 234 | del rsls_train 235 | del rsls 236 | 237 | # Resetting the model with the best weights 238 | self.reset_model(self.best_model_weights) 239 | 240 | # Save the best model 241 | self.save_model(epoch, best_epoch, self.best_model_weights, best_acc) 242 | 243 | # After training is complete, gets the classwise accuracies of all the splits and saves it based on the based model 244 | for i in list(self.data.keys()): 245 | # wandb is switched off temprorily so that the this validation loop is not logged 246 | g.wandb_log = False 247 | accs_dict , _ , _ , cls_acc = self.eval(phase=i) 248 | if g.log_offline: 249 | torch.save((accs_dict,cls_acc),g.log_dir+f"/metrics/{i}_cls_acc.pt") 250 | print(accs_dict) 251 | g.wandb_log = True 252 | 253 | print("Training Complete.") 254 | print_str = [f"Best validation accuracy is {best_acc} at epoch {best_epoch}"] 255 | print_write(print_str, self.log_file) 256 | 257 | # Empty cuda cache 258 | torch.cuda.empty_cache() 259 | 260 | def accumulate(self, phase): 261 | """Accumulates features of all the datapoints in a particular split 262 | 263 | Args: 264 | phase ([type]): Which split of dataset should be accumulated? 265 | """ 266 | print_str = ['Accumulating features: %s' % (phase)] 267 | print_write(print_str, self.log_file) 268 | time.sleep(0.25) 269 | self.accumulation = True 270 | 271 | torch.cuda.empty_cache() 272 | 273 | # In validation or testing mode, set model to eval() and initialize running loss/correct 274 | for model in self.networks.values(): 275 | model.eval() 276 | 277 | # Iterate over dataset 278 | self.feat = {} 279 | self.labs = {} 280 | 281 | accum_features = [] 282 | accum_labels = [] 283 | 284 | for inputs, labels, _ in tqdm(self.data[phase]): 285 | inputs, labels = inputs.cuda(), labels.cuda() 286 | # If on training phase, enable gradients 287 | with torch.set_grad_enabled(False): 288 | 289 | # In validation or testing 290 | self.batch_forward(inputs, labels, phase=phase) 291 | accum_features.append(self.features) 292 | accum_labels.append(labels) 293 | 294 | accum_features = torch.vstack(accum_features) 295 | accum_labels = torch.hstack(accum_labels) 296 | 297 | for i in accum_labels.unique().cpu().numpy(): 298 | self.feat[i] = accum_features[accum_labels == i] 299 | self.labs[i] = torch.full((self.feat[i].size()[0],), i).cuda() 300 | 301 | self.accumulation = False 302 | 303 | # This is there so that we can use source_import from the utils to import model 304 | def get_core(*args): 305 | return model(*args) 306 | -------------------------------------------------------------------------------- /libs/core/ce.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | import os 3 | import copy 4 | import pickle 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from torch.utils.data import TensorDataset, DataLoader 10 | from tqdm import tqdm 11 | import time 12 | import numpy as np 13 | import warnings 14 | import pdb 15 | 16 | # Custom imports 17 | from libs.core.core_base import model as base_model 18 | from libs.utils.utils import * 19 | from libs.utils.logger import Logger 20 | import libs.utils.globals as g 21 | if g.wandb_log: 22 | import wandb 23 | 24 | class model(base_model): 25 | """Basic CrossEntropy los training 26 | 27 | Args: 28 | base_model (): [description] 29 | """ 30 | def batch_forward(self, inputs): 31 | """Batch-wise forward prop 32 | 33 | Args: 34 | inputs (float Tensor): batch_size x image_size 35 | labels ([type], optional): [description]. Defaults to None. 36 | """ 37 | # Calculate Features and outputs 38 | self.features = self.networks["feat_model"](inputs) 39 | self.logits = self.networks["classifier"](self.features) 40 | 41 | # This is there so that we can use source_import from the utils to import model 42 | def get_core(*args): 43 | return model(*args) -------------------------------------------------------------------------------- /libs/core/core_base.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from torch.optim import optimizer 3 | import os 4 | import copy 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from tqdm import tqdm 10 | import numpy as np 11 | 12 | # Custom imports 13 | from libs.utils.utils import * 14 | from libs.utils.logger import Logger 15 | import libs.utils.globals as g 16 | if g.wandb_log: 17 | import wandb 18 | 19 | class model: 20 | def __init__(self, config, data): 21 | """Initialize 22 | 23 | Args: 24 | config (Dict): Dictionary of all the configurations 25 | data (list): Train, val, test splits 26 | """ 27 | 28 | self.config = config 29 | self.training_opt = self.config["training_opt"] 30 | self.data = data 31 | self.num_gpus = torch.cuda.device_count() 32 | self.do_shuffle = config["shuffle"] if "shuffle" in config else False 33 | self.start_epoch = 1 34 | 35 | # For gradient accumulation 36 | self.accumulation_step = self.training_opt["accumulation_step"] 37 | 38 | #-----Offline Logger 39 | self.logger = Logger(self.training_opt["log_dir"]) 40 | self.log_file = os.path.join(self.training_opt["log_dir"], "log.txt") 41 | self.logger.log_cfg(self.config) 42 | 43 | # If using steps for training, we need to calculate training steps 44 | # for each epoch based on actual number of training data instead of 45 | # oversampled data number 46 | print("Using steps for training.") 47 | self.training_data_num = len(self.data["train"].dataset) 48 | self.epoch_steps = int(self.training_data_num / self.training_opt["batch_size"]) 49 | 50 | # Initialize loss 51 | self.init_criterions() 52 | 53 | # Initialize model 54 | self.init_models() 55 | 56 | # Initialize model optimizer and scheduler 57 | print("Initializing model optimizer.") 58 | self.init_optimizers(self.model_optim_params_dict) 59 | 60 | 61 | def init_models(self): 62 | """Initialize models 63 | """ 64 | networks_defs = self.config["networks"] 65 | self.networks = {} 66 | self.model_optim_params_dict = {} 67 | 68 | print("Using", torch.cuda.device_count(), "GPUs.") 69 | 70 | # Create the models in loop 71 | for key, val in networks_defs.items(): 72 | def_file = val["def_file"] 73 | model_args = val["params"] 74 | 75 | # Create/load model 76 | self.networks[key] = source_import(def_file).create_model(**model_args) 77 | if networks_defs[key]["trainable"]: 78 | self.networks[key] = nn.DataParallel(self.networks[key]).cuda() 79 | 80 | # Freezing part or entire model 81 | if "fix" in val and val["fix"]: 82 | print(f"Freezing weights of module {key}") 83 | for param_name, param in self.networks[key].named_parameters(): 84 | param.requires_grad = False 85 | if "fix_set" in val: 86 | for fix_layer in val["fix_set"]: 87 | for param_name, param in self.networks[key].named_parameters(): 88 | if fix_layer == param_name: 89 | param.requires_grad = False 90 | print(f"=====> Freezing: {param_name} | {param.requires_grad}") 91 | 92 | # wandb logging 93 | if g.wandb_log: 94 | wandb.watch(self.networks[key], log="all") 95 | 96 | # Optimizer list to add to the optimizer in the "init_optimizer" step 97 | optim_params = val["optim_params"] 98 | self.model_optim_params_dict[key] = { 99 | "params": self.networks[key].parameters(), 100 | "lr": optim_params["lr"], 101 | "momentum": optim_params["momentum"], 102 | "weight_decay": optim_params["weight_decay"], 103 | } 104 | 105 | def init_optimizers(self, optim_params_dict): 106 | """Init optimizer with/without scheduler for it 107 | 108 | Args: 109 | optim_params_dict (Dict): A dictonary with all the params for the optimizer 110 | """ 111 | networks_defs = self.config["networks"] 112 | self.model_optimizer_dict = {} 113 | self.model_scheduler_dict = {} 114 | 115 | for key, val in networks_defs.items(): 116 | if networks_defs[key]["trainable"]: 117 | # optimizer 118 | if ("optimizer" in self.training_opt and self.training_opt["optimizer"] == "adam"): 119 | print("=====> Using Adam optimizer") 120 | optimizer = optim.Adam([optim_params_dict[key],]) 121 | else: 122 | print("=====> Using SGD optimizer") 123 | optimizer = optim.SGD([optim_params_dict[key],]) 124 | self.model_optimizer_dict[key] = optimizer 125 | 126 | # scheduler 127 | if val["scheduler_params"]: 128 | scheduler_params = val["scheduler_params"] 129 | 130 | if scheduler_params["coslr"]: 131 | self.model_scheduler_dict[key] = torch.optim.lr_scheduler.CosineAnnealingLR( 132 | optimizer, 133 | self.training_opt["num_epochs"], 134 | ) 135 | elif scheduler_params['warmup']: 136 | print("===> Module {} : Using warmup".format(key)) 137 | self.model_scheduler_dict[key] = WarmupMultiStepLR(optimizer, scheduler_params['lr_step'], 138 | gamma=scheduler_params['lr_factor'], warmup_epochs=scheduler_params['warm_epoch']) 139 | else: 140 | self.model_scheduler_dict[key] = optim.lr_scheduler.StepLR( 141 | optimizer, 142 | step_size=scheduler_params["step_size"], 143 | gamma=scheduler_params["gamma"], 144 | ) 145 | 146 | def init_criterions(self): 147 | """Initialize criterion (loss) and if required optimizer, scheduler for trainable params in it. 148 | """ 149 | criterion_defs = self.config["criterions"] 150 | self.criterions = {} 151 | self.criterion_weights = {} 152 | 153 | for key, val in criterion_defs.items(): 154 | def_file = val["def_file"] 155 | loss_args = list(val["loss_params"].values()) 156 | 157 | self.criterions[key] = (source_import(def_file).create_loss(*loss_args).cuda()) 158 | self.criterion_weights[key] = val["weight"] 159 | 160 | if val["optim_params"]: 161 | print("Initializing criterion optimizer.") 162 | optim_params = val["optim_params"] 163 | optim_params = [ 164 | { 165 | "params": self.criterions[key].parameters(), 166 | "lr": optim_params["lr"], 167 | "momentum": optim_params["momentum"], 168 | "weight_decay": optim_params["weight_decay"], 169 | } 170 | ] 171 | 172 | # Initialize criterion optimizer 173 | if ("optimizer" in self.training_opt and self.training_opt["optimizer"] == "adam"): 174 | print("=====> Using Adam optimizer") 175 | optimizer = optim.Adam(optim_params) 176 | else: 177 | print("=====> Using SGD optimizer") 178 | optimizer = optim.SGD(optim_params) 179 | self.criterion_optimizer = optimizer 180 | 181 | # Initialize criterion scheduler 182 | if "scheduler_params" in val and val["scheduler_params"]: 183 | scheduler_params = val["scheduler_params"] 184 | if scheduler_params["coslr"]: 185 | self.criterion_optimizer_scheduler = ( 186 | torch.optim.lr_scheduler.CosineAnnealingLR( 187 | optimizer, 188 | self.training_opt["num_epochs"], 189 | ) 190 | ) 191 | elif scheduler_params['warmup']: 192 | print("===> Module {} : Using warmup".format(key)) 193 | self.criterion_optimizer_scheduler = WarmupMultiStepLR(optimizer, scheduler_params['lr_step'], 194 | gamma=scheduler_params['lr_factor'], warmup_epochs=scheduler_params['warm_epoch']) 195 | else: 196 | self.criterion_optimizer_scheduler = optim.lr_scheduler.StepLR( 197 | optimizer, 198 | step_size=scheduler_params["step_size"], 199 | gamma=scheduler_params["gamma"], 200 | ) 201 | 202 | else: 203 | self.criterion_optimizer_scheduler = None 204 | else: 205 | self.criterion_optimizer = None 206 | self.criterion_optimizer_scheduler = None 207 | 208 | 209 | def show_current_lr(self): 210 | """Shows current learning rate 211 | 212 | Returns: 213 | float: Current learning rate 214 | """ 215 | max_lr = 0.0 216 | for key, val in self.model_optimizer_dict.items(): 217 | lr_set = list(set([para["lr"] for para in val.param_groups])) 218 | if max(lr_set) > max_lr: 219 | max_lr = max(lr_set) 220 | lr_set = ",".join([str(i) for i in lr_set]) 221 | print_str = [f"=====> Current Learning Rate of model {key} : {str(lr_set)}"] 222 | 223 | print_write(print_str, self.log_file) 224 | wandb_log({f"LR - {key}": float(lr_set)}) 225 | 226 | if self.criterion_optimizer: 227 | lr_set_rad = list(set([para["lr"] for para in self.criterion_optimizer.param_groups])) 228 | lr_set_rad = ",".join([str(i) for i in lr_set_rad]) 229 | 230 | wandb_log({f"LR - Radius": float(lr_set_rad)}) 231 | 232 | return max_lr 233 | 234 | def batch_forward(self, inputs): 235 | """Batch Forward 236 | 237 | Args: 238 | inputs (float Tensor): batch_size x image_size 239 | """ 240 | 241 | # Calculate Features and outputs 242 | self.features = self.networks["feat_model"](inputs) 243 | self.features = F.normalize(self.features, p=2, dim=1) 244 | 245 | self.logits = self.networks["classifier"](self.features) 246 | 247 | def batch_backward(self): 248 | """Backprop 249 | """ 250 | if self.accumulation_step == 1: 251 | # Zero out optimizer gradients 252 | for key, optimizer in self.model_optimizer_dict.items(): 253 | optimizer.zero_grad() 254 | if self.criterion_optimizer: 255 | self.criterion_optimizer.zero_grad() 256 | 257 | # Back-propagation from loss outputs 258 | self.loss.backward() 259 | 260 | # Gradient accumulation incase the batch doesnt fit inside a single GPU 261 | if (self.step+1) % self.accumulation_step == 0: 262 | # Step optimizers 263 | for key, optimizer in self.model_optimizer_dict.items(): 264 | optimizer.step() 265 | if self.criterion_optimizer: 266 | self.criterion_optimizer.step() 267 | 268 | if self.accumulation_step != 1: 269 | # Zero out optimizer gradients 270 | for key, optimizer in self.model_optimizer_dict.items(): 271 | optimizer.zero_grad() 272 | if self.criterion_optimizer: 273 | self.criterion_optimizer.zero_grad() 274 | 275 | def batch_loss(self, labels): 276 | """Calculate training loss 277 | 278 | Args: 279 | labels (int): Dim = Batch_size 280 | """ 281 | self.loss = 0 282 | 283 | # Calculating loss 284 | if "ClassifierLoss" in self.criterions.keys(): 285 | self.loss_classifier = self.criterions["ClassifierLoss"](self.logits, labels) 286 | self.loss_classifier *= self.criterion_weights["ClassifierLoss"] 287 | self.loss += self.loss_classifier 288 | 289 | self.loss = self.loss / self.accumulation_step 290 | 291 | def shuffle_batch(self, x, y): 292 | """Force shuffle data 293 | 294 | Args: 295 | x (float Tensor): Datapoints 296 | y (int): Labels 297 | 298 | Returns: 299 | floatTensor, int: Return shuffled datapoints and corresponding labels 300 | """ 301 | index = torch.randperm(x.size(0)) 302 | x = x[index] 303 | y = y[index] 304 | return x, y 305 | 306 | def print_test(self, path=None): 307 | """Loads best model and prints accuracies of all the splits 308 | 309 | Args: 310 | path ([type], optional): [description]. Defaults to None. 311 | """ 312 | if path != None: 313 | self.reset_model(torch.load(path)["state_dict_best"]) 314 | else: 315 | self.reset_model(torch.load(f"{self.training_opt['log_dir']}/final_model_checkpoint.pth")) 316 | 317 | for i in list(self.data.keys()): 318 | accs_dict , _ , _ , cls_acc = self.eval(phase=i) 319 | print(accs_dict) 320 | 321 | def train(self, retrain=False): 322 | """Main training 323 | 324 | Args: 325 | retrain (bool, optional): Incase of retraining different dataloaders are used. Defaults to False. 326 | """ 327 | phase = "train" 328 | print_str = ["Phase: train"] 329 | print_write(print_str, self.log_file) 330 | 331 | # Inits 332 | best_acc = 0.0 333 | best_epoch = 0 334 | self.retrain = retrain 335 | self.end_epoch = self.training_opt["num_epochs"] 336 | 337 | # Initialize best model and other variables 338 | self.best_model_weights = {} 339 | for key, _ in self.config["networks"].items(): 340 | if self.config["networks"][key]["trainable"]: 341 | self.best_model_weights[key] = copy.deepcopy(self.networks[key].state_dict()) 342 | 343 | # Loop over epochs 344 | for epoch in range(self.start_epoch, self.end_epoch + 1): 345 | # global config 346 | g.epoch_global = epoch 347 | 348 | # Switch to train mode 349 | for key, model in self.networks.items(): 350 | if self.config["networks"][key]["trainable"]: 351 | # only train the module with lr > 0 352 | if self.config["networks"][key]["optim_params"]["lr"] == 0.0: 353 | model.eval() 354 | else: 355 | model.train() 356 | 357 | # Empty cuda cache 358 | torch.cuda.empty_cache() 359 | 360 | # Step the schedulers 361 | if self.model_scheduler_dict: 362 | for key, scheduler in self.model_scheduler_dict.items(): 363 | scheduler.step() 364 | if self.criterion_optimizer_scheduler: 365 | self.criterion_optimizer_scheduler.step() 366 | 367 | print_write([self.training_opt["log_dir"]], self.log_file) 368 | 369 | # print learning rate 370 | current_lr = self.show_current_lr() 371 | current_lr = min(current_lr * 50, 1.0) 372 | 373 | # Choose a different dataloader based on whether this is training or retraining 374 | if self.retrain: 375 | data_enum = self.my_dataloader[phase] 376 | else: 377 | data_enum = self.data[phase] 378 | 379 | self.step = 0 380 | total_preds = [] 381 | total_labels = [] 382 | for inputs, labels, indexes in data_enum: 383 | # Break when step equal to epoch step 384 | if self.step == self.epoch_steps: 385 | break 386 | 387 | # Force shuffle option 388 | if self.do_shuffle: 389 | inputs, labels = self.shuffle_batch(inputs, labels) 390 | 391 | # Pushing to GPU 392 | inputs, labels = inputs.cuda(), labels.cuda() 393 | 394 | with torch.set_grad_enabled(True): 395 | # If training, forward with loss, and no top 5 accuracy calculation 396 | self.batch_forward(inputs) 397 | self.batch_loss(labels) 398 | self.batch_backward() 399 | 400 | # Tracking and printing predictions 401 | _, preds = torch.max(self.logits, 1) 402 | total_preds.append(torch2numpy(preds)) 403 | total_labels.append(torch2numpy(labels)) 404 | 405 | # Output minibatch training results 406 | if self.step % self.training_opt['display_step'] == 0: 407 | 408 | minibatch_loss_classifier = self.loss_classifier.item() if 'ClassifierLoss' in self.criterions else None 409 | minibatch_loss_embed = self.loss_embed.item() if 'EmbeddingLoss' in self.criterions else None 410 | minibatch_loss_embed_proto = self.loss_embed_proto.item() if 'EmbeddingLoss' in self.criterions else None 411 | minibatch_loss_embed_biasreduc = self.loss_embed_biasreduc.item() if 'EmbeddingLoss' in self.criterions else None 412 | minibatch_loss_total = self.loss.item() 413 | minibatch_acc = mic_acc_cal(preds, labels) 414 | 415 | 416 | print_str = ['Epoch: [%d/%d]' 417 | % (epoch, self.training_opt['num_epochs']), 418 | 'Step: [%d/%d]' 419 | % (self.step, self.epoch_steps), 420 | 'Minibatch_loss_embedding: %.3f' 421 | % (minibatch_loss_embed) if minibatch_loss_embed else '', 422 | 'Minibatch_loss_classifier: %.3f' 423 | % (minibatch_loss_classifier) if minibatch_loss_classifier else '', 424 | 'Minibatch_accuracy_micro: %.3f' 425 | % (minibatch_acc)] 426 | print_write(print_str, self.log_file) 427 | 428 | loss_info = { 429 | 'epoch': epoch, 430 | 'Step': self.step, 431 | 'Total': minibatch_loss_total, 432 | 'Embedding (Total)': minibatch_loss_embed, 433 | 'Proto': minibatch_loss_embed_proto, 434 | 'BiasReduc': minibatch_loss_embed_biasreduc, 435 | 'Classifier': minibatch_loss_classifier, 436 | } 437 | 438 | self.logger.log_loss(loss_info) 439 | 440 | # wandb logging 441 | wandb_log({"Training Loss": minibatch_loss_total}) 442 | 443 | # batch-level: sampler update 444 | if hasattr(self.data["train"].sampler, "update_weights"): 445 | if hasattr(self.data["train"].sampler, "ptype"): 446 | ptype = self.data["train"].sampler.ptype 447 | else: 448 | ptype = "score" 449 | ws = get_priority(ptype, self.logits.detach(), labels) 450 | 451 | inlist = [indexes.cpu().numpy(), ws] 452 | if self.training_opt["sampler"]["type"] == "ClassPrioritySampler": 453 | inlist.append(labels.cpu().numpy()) 454 | self.data["train"].sampler.update_weights(*inlist) 455 | 456 | # Clear things out (optional) 457 | del inputs, labels, self.logits, self.features, preds, indexes 458 | 459 | # Update steps 460 | self.step+=1 461 | g.step_global += 1 462 | 463 | # epoch-level: reset sampler weight 464 | if hasattr(self.data["train"].sampler, "get_weights"): 465 | self.logger.log_ws(epoch, self.data["train"].sampler.get_weights()) 466 | if hasattr(self.data["train"].sampler, "reset_weights"): 467 | self.data["train"].sampler.reset_weights(epoch) 468 | 469 | # After every epoch, validation 470 | rsls = {'epoch': epoch} 471 | rsls_train = self.eval_with_preds(total_preds, total_labels) 472 | rsls_eval, _ , _ , _ = self.eval(phase='val') 473 | rsls.update(rsls_train) 474 | rsls.update(rsls_eval) 475 | 476 | # Reset class weights for sampling if pri_mode is valid 477 | if hasattr(self.data["train"].sampler, "reset_priority"): 478 | ws = get_priority( 479 | self.data["train"].sampler.ptype, 480 | self.total_logits.detach(), 481 | self.total_labels, 482 | ) 483 | self.data["train"].sampler.reset_priority( 484 | ws, self.total_labels.cpu().numpy() 485 | ) 486 | 487 | self.logger.log_acc(rsls) 488 | 489 | # Under validation, the best model need to be updated 490 | if rsls_eval["val_all"] > best_acc: 491 | best_epoch = epoch 492 | best_acc = rsls_eval["val_all"] 493 | for key, _ in self.config["networks"].items(): 494 | if self.config["networks"][key]["trainable"]: 495 | self.best_model_weights[key] = copy.deepcopy(self.networks[key].state_dict()) 496 | 497 | # wandb log best epoch, train accuracy, based on best validation accuracy 498 | wandb_log({"Best Val": 100*best_acc, "Best Epoch": best_epoch}) 499 | wandb_log({"Best train": 100*rsls_train["train_all"], "Best Epoch": best_epoch}) 500 | 501 | wandb_log({'B_val_all': self.eval_acc_mic_top1, 502 | 'B_val_many': self.many_acc_top1, 503 | 'B_val_median': self.median_acc_top1, 504 | 'B_val_low': self.low_acc_top1}) 505 | 506 | wandb_log({'B_train_all': rsls_train["train_all"], 507 | 'B_train_many': rsls_train["train_many"], 508 | 'B_train_median': rsls_train["train_median"], 509 | 'B_train_low': rsls_train["train_low"]}) 510 | 511 | print("===> Saving checkpoint") 512 | self.save_latest(epoch) 513 | 514 | # Clear things out (optional) 515 | del rsls_eval 516 | del rsls_train 517 | del rsls 518 | 519 | # Resetting the model with the best weights 520 | self.reset_model(self.best_model_weights) 521 | 522 | # Save the best model 523 | self.save_model(epoch, best_epoch, self.best_model_weights, best_acc) 524 | 525 | # After training is complete, gets the classwise accuracies of all the splits and saves it based on the based model 526 | for i in list(self.data.keys()): 527 | # wandb is switched off temprorily so that the this validation loop is not logged 528 | g.wandb_log = False 529 | accs_dict , _ , _ , cls_acc = self.eval(phase=i) 530 | if g.log_offline: 531 | torch.save((accs_dict,cls_acc),g.log_dir+f"/metrics/{i}_cls_acc.pt") 532 | print(accs_dict) 533 | g.wandb_log = True 534 | 535 | print("Training Complete.") 536 | print_str = [f"Best validation accuracy is {best_acc} at epoch {best_epoch}"] 537 | print_write(print_str, self.log_file) 538 | 539 | # Empty cuda cache 540 | torch.cuda.empty_cache() 541 | 542 | def eval_with_preds(self, preds, labels): 543 | """Train accuracy 544 | 545 | Args: 546 | preds (int): Predictions 547 | labels (int): Ground Truth 548 | 549 | Returns: 550 | dict: dictionary of all training accuracies 551 | """ 552 | # Count the number of examples 553 | n_total = sum([len(p) for p in preds]) 554 | 555 | # Split the examples into normal and mixup 556 | normal_preds, normal_labels = [], [] 557 | mixup_preds, mixup_labels1, mixup_labels2, mixup_ws = [], [], [], [] 558 | for p, l in zip(preds, labels): 559 | if isinstance(l, tuple): 560 | mixup_preds.append(p) 561 | mixup_labels1.append(l[0]) 562 | mixup_labels2.append(l[1]) 563 | mixup_ws.append(l[2] * np.ones_like(l[0])) 564 | else: 565 | normal_preds.append(p) 566 | normal_labels.append(l) 567 | 568 | # Calculate normal prediction accuracy 569 | rsl = { 570 | "train_all": 0.0, 571 | "train_many": 0.0, 572 | "train_median": 0.0, 573 | "train_low": 0.0, 574 | } 575 | 576 | if len(normal_preds) > 0: 577 | normal_preds, normal_labels = list( 578 | map(np.concatenate, [normal_preds, normal_labels]) 579 | ) 580 | n_top1 = mic_acc_cal(normal_preds, normal_labels) 581 | ( 582 | n_top1_many, 583 | n_top1_median, 584 | n_top1_low, 585 | ) = shot_acc(normal_preds, normal_labels, self.data["train"]) 586 | rsl["train_all"] += len(normal_preds) / n_total * n_top1 587 | rsl["train_many"] += len(normal_preds) / n_total * n_top1_many 588 | rsl["train_median"] += len(normal_preds) / n_total * n_top1_median 589 | rsl["train_low"] += len(normal_preds) / n_total * n_top1_low 590 | 591 | # Calculate mixup prediction accuracy 592 | if len(mixup_preds) > 0: 593 | mixup_preds, mixup_labels, mixup_ws = list( 594 | map( 595 | np.concatenate, 596 | [mixup_preds * 2, mixup_labels1 + mixup_labels2, mixup_ws], 597 | ) 598 | ) 599 | mixup_ws = np.concatenate([mixup_ws, 1 - mixup_ws]) 600 | n_top1 = weighted_mic_acc_cal(mixup_preds, mixup_labels, mixup_ws) 601 | n_top1_many, n_top1_median, n_top1_low, = weighted_shot_acc( 602 | mixup_preds, mixup_labels, mixup_ws, self.data["train"] 603 | ) 604 | rsl["train_all"] += len(mixup_preds) / 2 / n_total * n_top1 605 | rsl["train_many"] += len(mixup_preds) / 2 / n_total * n_top1_many 606 | rsl["train_median"] += len(mixup_preds) / 2 / n_total * n_top1_median 607 | rsl["train_low"] += len(mixup_preds) / 2 / n_total * n_top1_low 608 | 609 | # Top-1 accuracy and additional string 610 | print_str = [ 611 | "\n Training acc Top1: %.3f \n" % (rsl["train_all"]), 612 | "Many_top1: %.3f" % (rsl["train_many"]), 613 | "Median_top1: %.3f" % (rsl["train_median"]), 614 | "Low_top1: %.3f" % (rsl["train_low"]), 615 | "\n", 616 | ] 617 | 618 | print_write(print_str, self.log_file) 619 | phase = "train" 620 | wandb_log({phase + '_all': rsl["train_all"]*100, 621 | phase + '_many': rsl["train_many"]*100, 622 | phase + '_median': rsl["train_median"]*100, 623 | phase + '_low': rsl["train_low"]*100, 624 | phase + ' Accuracy': rsl["train_all"]*100,}) 625 | 626 | return rsl 627 | 628 | def eval(self, phase='val'): 629 | print_str = ['Phase: %s' % (phase)] 630 | print_write(print_str, self.log_file) 631 | 632 | torch.cuda.empty_cache() 633 | 634 | # In validation or testing mode, set model to eval() and initialize running loss/correct 635 | for model in self.networks.values(): 636 | model.eval() 637 | 638 | self.total_labels = torch.empty(0, dtype=torch.long) 639 | self.total_preds = [] 640 | minibatch_loss_total = [] 641 | 642 | # Choose a different dataloader based on whether this is training or retraining 643 | if self.retrain: 644 | data_enum = self.my_dataloader[phase] 645 | else: 646 | data_enum = self.data[phase] 647 | 648 | # Iterate over dataset 649 | stepval = 0 650 | for inputs, labels, paths in tqdm(data_enum): 651 | inputs, labels = inputs.cuda(), labels.cuda() 652 | 653 | # If on training phase, enable gradients 654 | with torch.set_grad_enabled(False): 655 | 656 | # In validation or testing 657 | self.batch_forward(inputs) 658 | self.batch_loss(labels) 659 | minibatch_loss_total.append(self.loss.item()) 660 | 661 | _, preds = F.softmax(self.logits, dim=1).max(dim=1) 662 | self.total_preds.append(preds.cpu()) 663 | self.total_labels = torch.cat((self.total_labels, labels.cpu())) 664 | 665 | #----Clear things out 666 | del preds, inputs, labels 667 | torch.cuda.empty_cache() 668 | 669 | stepval+=1 670 | 671 | wandb_log({"Validation Loss": np.mean(minibatch_loss_total)}) 672 | 673 | preds = torch.hstack(self.total_preds) 674 | 675 | # Calculate the overall accuracy and F measurement 676 | self.eval_acc_mic_top1= mic_acc_cal(preds, 677 | self.total_labels) 678 | self.eval_f_measure = F_measure(preds, self.total_labels, theta=self.training_opt['open_threshold']) 679 | 680 | self.many_acc_top1, \ 681 | self.median_acc_top1, \ 682 | self.low_acc_top1, \ 683 | self.cls_accs = shot_acc(preds, 684 | self.total_labels, 685 | self.data['train'], 686 | acc_per_cls=True) 687 | 688 | # Top-1 accuracy and additional string 689 | print_str = ['\n\n', 690 | 'Phase: %s' 691 | % (phase), 692 | '\n\n', 693 | 'Evaluation_accuracy_micro_top1: %.3f' 694 | % (self.eval_acc_mic_top1), 695 | '\n', 696 | 'Averaged F-measure: %.3f' 697 | % (self.eval_f_measure), 698 | '\n', 699 | 'Many_shot_accuracy_top1: %.3f' 700 | % (self.many_acc_top1), 701 | 'Median_shot_accuracy_top1: %.3f' 702 | % (self.median_acc_top1), 703 | 'Low_shot_accuracy_top1: %.3f' 704 | % (self.low_acc_top1), 705 | '\n'] 706 | 707 | rsl = {phase + '_all': self.eval_acc_mic_top1, 708 | phase + '_many': self.many_acc_top1, 709 | phase + '_median': self.median_acc_top1, 710 | phase + '_low': self.low_acc_top1, 711 | phase + '_fscore': self.eval_f_measure, 712 | phase + '_loss': self.loss.item()} 713 | 714 | wandb_log({phase + '_all': self.eval_acc_mic_top1*100, 715 | phase + '_many': self.many_acc_top1*100, 716 | phase + '_median': self.median_acc_top1*100, 717 | phase + '_low': self.low_acc_top1*100, 718 | phase + ' Accuracy': self.eval_acc_mic_top1*100, 719 | phase + ' Loss': self.loss.item(),}) 720 | 721 | 722 | print_write(print_str, self.log_file) 723 | print(f"------------->{self.eval_acc_mic_top1 * 100}") 724 | 725 | return rsl, preds, self.total_labels, self.cls_accs 726 | 727 | def save_latest(self, epoch): 728 | """Saves necessary model states for resuming. 729 | 730 | Args: 731 | epoch (int): Epoch number 732 | """ 733 | # Model's state_dict 734 | model_weights = {} 735 | for key, _ in self.config["networks"].items(): 736 | if self.config["networks"][key]["trainable"]: 737 | model_weights[key] = copy.deepcopy( 738 | self.networks[key].state_dict() 739 | ) 740 | 741 | # Optimizer's state_dict 742 | optimizer_state_dict = {} 743 | for key, _ in self.model_optimizer_dict.items(): 744 | optimizer_state_dict[key] = copy.deepcopy( 745 | self.model_optimizer_dict[key].state_dict() 746 | ) 747 | 748 | # Criterion's Optimizer's state_dict 749 | criterion_optimizer_state_dict = self.criterion_optimizer.state_dict() if self.criterion_optimizer else None 750 | 751 | Scheduler's state dict 752 | scheduler_state_dict = {} 753 | if self.model_scheduler_dict: 754 | for key, _ in self.model_scheduler_dict.items(): 755 | scheduler_state_dict[key] = copy.deepcopy( 756 | self.model_scheduler_dict[key].state_dict() 757 | ) 758 | else: 759 | scheduler_state_dict = None 760 | 761 | # Criterion's Scheduler's state_dict 762 | criterion_scheduler_state_dict = self.criterion_optimizer_scheduler.state_dict() if self.criterion_optimizer_scheduler else None 763 | 764 | 765 | model_states = { 766 | "epoch": epoch, 767 | "state_dict": model_weights, 768 | "opt_state_dict": optimizer_state_dict, 769 | "opt_crit_state_dict": criterion_optimizer_state_dict, 770 | "sch_state_dict": scheduler_state_dict, 771 | "sch_crit_state_dict": criterion_scheduler_state_dict, 772 | "wandb_id": self.config["wandb_id"], 773 | } 774 | 775 | model_dir = os.path.join( 776 | self.training_opt["log_dir"], "latest_model_checkpoint.pth" 777 | ) 778 | torch.save(model_states, model_dir) 779 | 780 | def save_model(self, epoch, best_epoch, best_model_weights, best_acc): 781 | """Saves the best model's weights 782 | 783 | Args: 784 | epoch (int): Epoch number 785 | best_epoch (int): Epoch with the best accuracy or val loss 786 | best_model_weights (float Tensor): Best model's weights 787 | best_acc (float): Best accuracy 788 | """ 789 | 790 | model_states = { 791 | "epoch": epoch, 792 | "best_epoch": best_epoch, 793 | "state_dict_best": best_model_weights, 794 | "best_acc": best_acc, 795 | } 796 | 797 | model_dir = os.path.join( 798 | self.training_opt["log_dir"], "final_model_checkpoint.pth" 799 | ) 800 | 801 | torch.save(model_states, model_dir) 802 | 803 | def reset_model(self, model_state): 804 | """Resets the model with the best weight 805 | 806 | Args: 807 | model_state (dict): dict with best weight 808 | """ 809 | for key, model in self.networks.items(): 810 | if self.config["networks"][key]["trainable"]: 811 | weights = model_state[key] 812 | weights = {k: weights[k] for k in weights if k in model.state_dict()} 813 | model.load_state_dict(weights) 814 | 815 | def resume_run(self, saved_dict): 816 | """Resumes the run based on the states saved by "self.save_latest()" 817 | 818 | Args: 819 | model_state (dict): dict with best weight 820 | """ 821 | loaded_dict = torch.load(saved_dict) 822 | model_state = loaded_dict["state_dict"] 823 | optimizer_state_dict = loaded_dict["opt_state_dict"] 824 | criterion_optimizer_state_dict = loaded_dict["opt_crit_state_dict"] 825 | scheduler_state_dict = loaded_dict["sch_state_dict"] 826 | criterion_scheduler_state_dict = loaded_dict["sch_crit_state_dict"] 827 | 828 | for key, model in self.networks.items(): 829 | if self.config["networks"][key]["trainable"]: 830 | weights = model_state[key] 831 | weights = {k: weights[k] for k in weights if k in model.state_dict()} 832 | model.load_state_dict(weights) 833 | 834 | # Optimizer's state_dict 835 | for key, _ in self.model_optimizer_dict.items(): 836 | self.model_optimizer_dict[key].load_state_dict(optimizer_state_dict[key]) 837 | 838 | # Criterion's Optimizer's state_dict 839 | if self.criterion_optimizer : 840 | self.criterion_optimizer.load_state_dict(criterion_optimizer_state_dict) 841 | 842 | # Scheduler's state dict 843 | if self.model_scheduler_dict: 844 | for key, _ in self.model_scheduler_dict.items(): 845 | self.model_scheduler_dict[key].load_state_dict(scheduler_state_dict[key]) 846 | 847 | # Criterion's Scheduler's state_dict 848 | if self.criterion_optimizer_scheduler : 849 | self.criterion_optimizer_scheduler.load_state_dict(criterion_scheduler_state_dict) 850 | 851 | self.start_epoch = loaded_dict["epoch"] + 1 852 | 853 | print(f"\nResuming from Epoch: {self.start_epoch}!\n") 854 | 855 | #----------------------------------------------------- 856 | 857 | # This is there so that we can use source_import from the utils to import model 858 | def get_core(*args): 859 | return model(*args) 860 | -------------------------------------------------------------------------------- /libs/core/ecbd.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from torch.optim import optimizer 3 | import os 4 | import copy 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from tqdm import tqdm 10 | import numpy as np 11 | 12 | # Custom imports 13 | from libs.core.core_base import model as base_model 14 | from libs.utils.utils import * 15 | from libs.utils.logger import Logger 16 | import libs.utils.globals as g 17 | if g.wandb_log: 18 | import wandb 19 | 20 | 21 | # For repo specific to CBD paper with much more detailed instructions, check https://github.com/rahulvigneswaran/Class-Balanced-Distillation-for-Long-Tailed-Visual-Recognition.pytorch 22 | class model(base_model): 23 | def batch_forward(self, inputs): 24 | """Batch Forward 25 | 26 | Args: 27 | inputs (float Tensor): batch_size x image_size 28 | """ 29 | 30 | self.features_temp = self.networks["feat_model"](inputs) 31 | self.features_temp = F.normalize(self.features_temp, p=2, dim=1) 32 | 33 | if len(self.networks.keys()) > 3 : 34 | # Convert student feature to match with concatenated teachers' features 35 | self.features = self.networks["ecbd_converter"](self.features_temp) 36 | self.features = F.normalize(self.features, p=2, dim=1) 37 | else: 38 | self.features = self.features_temp 39 | 40 | # Calculate Features and outputs 41 | self.features_teacher = [] 42 | for i in self.networks.keys(): 43 | if not(("feat_model" in i) or ("classifier" in i) or ("ecbd_converter" in i)): 44 | self.temp = self.networks[i](inputs) 45 | self.features_teacher.append(F.normalize(self.temp, p=2, dim=1)) 46 | self.features_teacher = torch.hstack(self.features_teacher) 47 | self.features_teacher = F.normalize(self.features_teacher, p=2, dim=1) 48 | 49 | self.logits = self.networks["classifier"](self.features_temp) 50 | 51 | def batch_loss(self, labels): 52 | """Calculate training loss 53 | """ 54 | self.loss = 0 55 | 56 | # Calculating loss 57 | if "DistillLoss" in self.criterions.keys(): 58 | self.loss_distill = self.criterions["DistillLoss"](self.features, self.features_teacher) 59 | self.loss_distill *= self.criterion_weights["DistillLoss"] 60 | self.loss += self.loss_distill 61 | 62 | # Calculating loss 63 | if "ClassifierLoss" in self.criterions.keys(): 64 | self.loss_classifier = self.criterions["ClassifierLoss"](self.logits, labels) 65 | self.loss_classifier *= self.criterion_weights["ClassifierLoss"] 66 | self.loss += self.loss_classifier 67 | 68 | # This is there so that we can use source_import from the utils to import model 69 | def get_core(*args): 70 | return model(*args) -------------------------------------------------------------------------------- /libs/core/modals.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from torch.optim import optimizer 3 | import os 4 | import copy 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | import torch.nn.functional as F 9 | from torch.utils.data import TensorDataset, DataLoader 10 | from tqdm import tqdm 11 | import numpy as np 12 | 13 | # Custom imports 14 | from libs.core.core_base import model as base_model 15 | from libs.utils.utils import * 16 | from libs.utils.logger import Logger 17 | import libs.utils.globals as g 18 | if g.wandb_log: 19 | import wandb 20 | 21 | class model(base_model): 22 | def batch_forward(self, inputs, labels=None, phase="train", retrain= False): 23 | """Batch Forward 24 | 25 | Args: 26 | inputs (float Tensor): batch_size x image_size 27 | labels (int, optional): Labels. Defaults to None. 28 | phase (str, optional): Train or Test?. Defaults to "train". 29 | retrain (bool, optional): When retraining, only the classifier is forward propagated. Defaults to False. 30 | """ 31 | 32 | # Calculate Features and outputs 33 | if not(retrain): 34 | self.features = self.networks["feat_model"](inputs) 35 | else: 36 | self.features = inputs 37 | 38 | self.logits = self.networks["classifier"](self.features, labels) 39 | 40 | def accumulate(self, phase): 41 | """Accumulates features of all the datapoints in a particular split 42 | 43 | Args: 44 | phase ([type]): Which split of dataset should be accumulated? 45 | """ 46 | print_str = ['Accumulating features: %s' % (phase)] 47 | print_write(print_str, self.log_file) 48 | time.sleep(0.25) 49 | 50 | torch.cuda.empty_cache() 51 | 52 | # In validation or testing mode, set model to eval() and initialize running loss/correct 53 | for model in self.networks.values(): 54 | model.eval() 55 | 56 | # Iterate over dataset 57 | self.feat = {} 58 | self.labs = {} 59 | 60 | accum_features = [] 61 | accum_labels = [] 62 | 63 | for inputs, labels, paths in tqdm(self.data[phase]): 64 | inputs, labels = inputs.cuda(), labels.cuda() 65 | # If on training phase, enable gradients 66 | with torch.set_grad_enabled(False): 67 | 68 | # In validation or testing 69 | self.batch_forward(inputs, labels, phase=phase) 70 | 71 | accum_features.append(self.features) 72 | accum_labels.append(labels) 73 | torch.cuda.empty_cache() 74 | 75 | # Appending and stacking is same as just concatenating. In this you just dont have to declare a torch.empty of particular size. 76 | accum_features = torch.vstack(accum_features) 77 | accum_labels = torch.hstack(accum_labels) 78 | 79 | for i in accum_labels.unique().cpu().numpy(): 80 | self.feat[i] = accum_features[accum_labels == i] 81 | self.labs[i] = torch.full((self.feat[i].size()[0],), i).cuda() 82 | 83 | def generate_points(self): 84 | """Generate new datapoints 85 | """ 86 | 87 | if not os.path.isdir(self.config["training_opt"]["log_generate"]): 88 | os.makedirs(self.config["training_opt"]["log_generate"]) 89 | else: 90 | raise Exception("Generation Directory already exists!!") 91 | 92 | # Class statistics 93 | self.base_std = [] 94 | for i in self.feat.keys(): 95 | self.base_std.append(torch.std(self.feat[i], dim=0)) 96 | 97 | self.base_std = torch.vstack(self.base_std) 98 | 99 | 100 | if self.config["pg"]["generate"]: 101 | sample_from_each = self.get_sample_count(self.config["training_opt"]["data_count"],self.feat.keys()) 102 | 103 | self.generated_points = {} 104 | 105 | for i in tqdm(self.feat.keys()): 106 | if np.sum(sample_from_each[i]) == 0 and self.config["pg"]["extra_points"] == 0 : 107 | continue 108 | 109 | self.generated_points[i] = [] 110 | for k, x_ij in zip(sample_from_each[i], self.feat[i]): 111 | if k == 0: 112 | continue 113 | # Adding a small gaussian noise to the existing train datapoints 114 | gen = x_ij + self.config["pg"]["lambda"]*torch.randn((int(k), 128), device='cuda')*self.base_std[i] #FIXME should be standard deviation across samples or dims or dim_samples? 115 | self.generated_points[i].append(gen) 116 | 117 | self.generated_points[i] = torch.vstack(self.generated_points[i]) 118 | 119 | torch.save(self.generated_points, self.config["training_opt"]["log_generate"] + "/generated_points.pt") 120 | print("\nPoint Generation Completed!\n") 121 | 122 | else: 123 | self.generated_points ={} 124 | print("\nPoint Generation is False!\n") 125 | 126 | def retrain(self,): 127 | """Creates a new dataloader, reinits everything and trains just the classifier part. 128 | """ 129 | # Prepare a dataloader for all splits which includes the generated points and is also tukey transformed. 130 | self.prepare_updated_dataset(include_generated_points = self.config["pg"]["generate"]) 131 | 132 | # Change log_dir so it wont change the weights of the parent model 133 | self.config["training_opt"]["log_dir"] = self.config["training_opt"]["log_retrain"] 134 | 135 | # Create retrain directory 136 | if not os.path.isdir(self.config["training_opt"]["log_dir"]): 137 | os.makedirs(self.config["training_opt"]["log_dir"]) 138 | else: 139 | raise Exception("Retrained Directory already exists!!") 140 | 141 | g.log_dir = self.config["training_opt"]["log_dir"] 142 | if g.log_offline: 143 | if not os.path.isdir(f"{g.log_dir}/metrics"): 144 | os.makedirs(f"{g.log_dir}/metrics") 145 | 146 | # Reinitialize everything 147 | print("Using steps for training.") 148 | self.training_data_num = len(self.my_dataloader["train"].dataset) 149 | self.epoch_steps = int( 150 | self.training_data_num / self.training_opt["batch_size"] 151 | ) 152 | 153 | # Init logger 154 | self.logger = Logger(self.training_opt["log_dir"]) 155 | self.log_file = os.path.join(self.training_opt["log_dir"], "log.txt") 156 | self.logger.log_cfg(self.config) 157 | 158 | # Initialize loss 159 | self.init_criterions() 160 | 161 | # Initialize model 162 | self.init_models() 163 | 164 | # Initialize model optimizer and scheduler 165 | print("Initializing model optimizer.") 166 | self.init_optimizers(self.model_optim_params_dict) 167 | 168 | self.train(retrain=True) 169 | 170 | def prepare_updated_dataset(self, include_generated_points = True): 171 | """Prepares a dataloader for all splits which includes the generated points and is also tukey transformed. 172 | 173 | 174 | Args: 175 | include_generated_points (bool, optional): Do you wanna include the newly generated points. Defaults to True. 176 | """ 177 | self.my_dataloader = {} 178 | self.for_distance_analysis = {} 179 | for phase in ["train", "val", "test"]: 180 | 181 | self.accumulate(phase=phase) 182 | 183 | feat_all = [] 184 | labs_all = [] 185 | 186 | if self.config["pg"]["tukey"]: 187 | for i in self.labs.keys(): 188 | self.feat[i] = self.tukey_transform(self.feat[i], lam=self.config["pg"]["tukey_value"]) 189 | 190 | for i in self.labs.keys(): 191 | feat_all.append(self.feat[i]) 192 | labs_all.append(self.labs[i]) 193 | 194 | feat_all = torch.vstack(feat_all) 195 | labs_all = torch.hstack(labs_all).cuda() 196 | 197 | if include_generated_points and phase == "train": 198 | generated_points = torch.load(self.config["training_opt"]["log_generate"] + "/generated_points.pt") 199 | for i in generated_points.keys(): 200 | feat_all = torch.cat((feat_all, generated_points[i].cuda())) 201 | labs_all = torch.cat((labs_all, torch.full((generated_points[i].size()[0],), int(i)).cuda())) 202 | 203 | # Create dataloader 204 | my_dataset = TensorDataset(feat_all, labs_all, labs_all) 205 | self.my_dataloader[phase] = DataLoader(my_dataset, batch_size=self.config["training_opt"]["batch_size"], shuffle=True) 206 | 207 | # This is there so that we can use source_import from the utils to import model 208 | def get_core(*args): 209 | return model(*args) 210 | -------------------------------------------------------------------------------- /libs/data/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rahulvigneswaran/TailCalibX/b836fd93925efb48f610c08bd56bb0c5a3fca9b0/libs/data/.DS_Store -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/.gitattributes: -------------------------------------------------------------------------------- 1 | *.data_target_tuple filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.01_phase_train_seed_1.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:9586ff290b6f4be75e07db22c615bd20fe1a90ce7448b3a2ac52bc3cbb5bc21c 3 | size 48585007 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.01_phase_train_seed_2.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:694ce7e2d044431c8dd95ece33c12255c0f3e98a74f9480a96fd86d52eca556e 3 | size 48587759 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.01_phase_train_seed_3.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:01464f4279d2c4d197dbd30cac042fd8e525fef403e878092c15474285608c1e 3 | size 48579503 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.02_phase_train_seed_1.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e90c754566e10c7c462941b3e1b7ccf3f684a24b4f997f5c2ca2216a9cfb8fe5 3 | size 56476335 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.02_phase_train_seed_2.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:ec8400036a13542d3813cff7123b75362d096d96c549b48539f2ffd29af4a31c 3 | size 56468847 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.02_phase_train_seed_3.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:562adf56accb87e0b376dbd20c1f873e668315ce271e1ffa10c5db17877c2cf8 3 | size 56428527 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.1_phase_train_seed_1.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:25c8499467c9647c99c93245aa71634d4b550f7867bda920ff39e0ebdfd09c36 3 | size 87528751 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.1_phase_train_seed_2.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:b1f658f0791419a693ebe1f87a7fed9694961535ccae21fef1266cc6fb35b613 3 | size 87483887 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/Imb_0.1_phase_train_seed_3.data_target_tuple: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1b86d8a2665f5ed0a2bca453c23be2351d69d8aaa951c1d3aa2a6c5f6d1a33ed 3 | size 87488303 4 | -------------------------------------------------------------------------------- /libs/data/CIFAR100_LT/README.md: -------------------------------------------------------------------------------- 1 | 2 | Reproducibility is an issue when it comes to CIFAR100_LT. This is because different seeds select different data points when creating the imbalanced dataset. This will not be resolved even after choosing the same seed as ours `[1,2,3]`, because the randomness caused by the system (ours vs yours) still exists. 3 | 4 | To evade this issue, we fixed and saved the data points we used for each imbalance factor and seed. 5 | 6 | So, if you want to exactly reproduce our results, we suggest you to load the data, target tuple for respective phase, imb_factor, seed from here. 7 | -------------------------------------------------------------------------------- /libs/data/ImbalanceCIFAR.py: -------------------------------------------------------------------------------- 1 | """ 2 | Adopted from https://github.com/Megvii-Nanjing/BBN 3 | Customized by Kaihua Tang 4 | """ 5 | 6 | # Imports 7 | import torchvision 8 | import torchvision.transforms as transforms 9 | import numpy as np 10 | from PIL import Image 11 | import random 12 | import os 13 | import pickle 14 | import torch 15 | 16 | class IMBALANCECIFAR10(torchvision.datasets.CIFAR10): 17 | cls_num = 10 18 | 19 | def __init__(self, phase, imbalance_ratio, root = '/gruntdata5/kaihua/datasets', imb_type='exp', class_order=None, balanced=None, normal=False, special_aug=False, seed=1): 20 | train = True if phase == "train" else False 21 | super(IMBALANCECIFAR10, self).__init__(root, train, transform=None, target_transform=None, download=True) 22 | self.train = train 23 | self.phase = phase 24 | self.seed = seed 25 | self.imb_ratio = imbalance_ratio 26 | self.get_filenames() 27 | if self.train: 28 | if special_aug: 29 | aug = [ 30 | transforms.RandomApply([transforms.ColorJitter(brightness=(0.1, 0.3), contrast=(0.1, 0.3), saturation=(0.1, 0.3), hue=(0.1, 0.3))], p=0.8), 31 | transforms.RandomCrop(32, padding=4), 32 | transforms.RandomHorizontalFlip(), 33 | transforms.ToTensor(), 34 | transforms.Normalize( 35 | (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010) 36 | ), 37 | AddGaussianNoise(0., 0.01) 38 | ] 39 | else: 40 | aug = [ 41 | transforms.RandomCrop(32, padding=4), 42 | transforms.RandomHorizontalFlip(), 43 | # transforms.Resize(224), 44 | transforms.ToTensor(), 45 | transforms.Normalize( 46 | (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010) 47 | ), 48 | ] 49 | self.img_num_list = self.get_img_num_per_cls(self.cls_num, imb_type, imbalance_ratio) 50 | self.gen_imbalanced_data(self.img_num_list) 51 | self.transform = transforms.Compose(aug) 52 | else: 53 | self.transform = transforms.Compose([ 54 | #transforms.Resize(224), 55 | transforms.ToTensor(), 56 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 57 | ]) 58 | 59 | self.labels = self.targets 60 | 61 | print("{} Mode: Contain {} images".format(phase, len(self.data))) 62 | 63 | def _get_class_dict(self): 64 | class_dict = dict() 65 | for i, anno in enumerate(self.get_annotations()): 66 | cat_id = anno["category_id"] 67 | if not cat_id in class_dict: 68 | class_dict[cat_id] = [] 69 | class_dict[cat_id].append(i) 70 | return class_dict 71 | 72 | 73 | def get_img_num_per_cls(self, cls_num, imb_type, imb_factor): 74 | img_max = len(self.data) / cls_num 75 | img_num_per_cls = [] 76 | if imb_type == 'exp': 77 | for cls_idx in range(cls_num): 78 | num = img_max * (imb_factor**(cls_idx / (cls_num - 1.0))) 79 | img_num_per_cls.append(int(num)) 80 | elif imb_type == 'step': 81 | for cls_idx in range(cls_num // 2): 82 | img_num_per_cls.append(int(img_max)) 83 | for cls_idx in range(cls_num // 2): 84 | img_num_per_cls.append(int(img_max * imb_factor)) 85 | else: 86 | img_num_per_cls.extend([int(img_max)] * cls_num) 87 | return img_num_per_cls 88 | 89 | def gen_imbalanced_data(self, img_num_per_cls): 90 | """Synthetically imbalancing the dataset. This function would make the dataset long-tailed based on "img_num_per_cls" 91 | """ 92 | new_data = [] 93 | new_targets = [] 94 | new_filenames = [] 95 | targets_np = np.array(self.targets, dtype=np.int64) 96 | classes = np.unique(targets_np) 97 | self.num_per_cls_dict = dict() 98 | for the_class, the_img_num in zip(classes, img_num_per_cls): 99 | self.num_per_cls_dict[the_class] = the_img_num 100 | idx = np.where(targets_np == the_class)[0] 101 | np.random.shuffle(idx) 102 | selec_idx = idx[:the_img_num] 103 | new_data.append(self.data[selec_idx, ...]) 104 | new_targets.extend([the_class, ] * the_img_num) 105 | new_filenames.extend([self.filenames[i] for i in selec_idx]) 106 | new_data = np.vstack(new_data) 107 | self.data = new_data 108 | self.targets = new_targets 109 | self.filenames = new_filenames 110 | 111 | def __getitem__(self, index): 112 | img, label = self.data[index], self.labels[index] 113 | filename = self.filenames[index] 114 | # doing this so that it is consistent with all other datasets 115 | # to return a PIL Image 116 | img = Image.fromarray(img) 117 | 118 | if self.transform is not None: 119 | img = self.transform(img) 120 | 121 | if self.target_transform is not None: 122 | label = self.target_transform(label) 123 | 124 | return img, label, index 125 | 126 | def __len__(self): 127 | return len(self.labels) 128 | 129 | def get_num_classes(self): 130 | return self.cls_num 131 | 132 | def get_annotations(self): 133 | annos = [] 134 | for label in self.labels: 135 | annos.append({'category_id': int(label)}) 136 | return annos 137 | 138 | def get_cls_num_list(self): 139 | cls_num_list = [] 140 | for i in range(self.cls_num): 141 | cls_num_list.append(self.num_per_cls_dict[i]) 142 | return cls_num_list 143 | 144 | def get_filenames(self): 145 | self.filenames = [] 146 | if self.train: 147 | downloaded_list = self.train_list 148 | else: 149 | downloaded_list = self.test_list 150 | 151 | for file_name, checksum in downloaded_list: 152 | file_path = os.path.join(self.root, self.base_folder, file_name) 153 | with open(file_path, 'rb') as f: 154 | self.entry = pickle.load(f, encoding='latin1') 155 | self.filenames.extend(self.entry["filenames"]) 156 | 157 | class IMBALANCECIFAR100(IMBALANCECIFAR10): 158 | """`CIFAR100 `_ Dataset. 159 | This is a subclass of the `CIFAR10` Dataset. 160 | """ 161 | cls_num = 100 162 | base_folder = 'cifar-100-python' 163 | url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz" 164 | filename = "cifar-100-python.tar.gz" 165 | tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85' 166 | train_list = [ 167 | ['train', '16019d7e3df5f24257cddd939b257f8d'], 168 | ] 169 | 170 | test_list = [ 171 | ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'], 172 | ] 173 | meta = { 174 | 'filename': 'meta', 175 | 'key': 'fine_label_names', 176 | 'md5': '7973b15100ade9c7d40fb424638fde48', 177 | } 178 | def gen_imbalanced_data(self, img_num_per_cls): 179 | """Synthetically imbalancing the dataset. This function would make the dataset long-tailed based on "img_num_per_cls" 180 | """ 181 | 182 | new_data = [] 183 | new_targets = [] 184 | new_filenames = [] 185 | targets_np = np.array(self.targets, dtype=np.int64) 186 | classes = np.unique(targets_np) 187 | self.num_per_cls_dict = dict() 188 | for the_class, the_img_num in zip(classes, img_num_per_cls): 189 | self.num_per_cls_dict[the_class] = the_img_num 190 | idx = np.where(targets_np == the_class)[0] 191 | np.random.shuffle(idx) 192 | selec_idx = idx[:the_img_num] 193 | new_data.append(self.data[selec_idx, ...]) 194 | new_targets.extend([the_class, ] * the_img_num) 195 | new_filenames.extend([self.filenames[i] for i in selec_idx]) 196 | new_data = np.vstack(new_data) 197 | self.data = new_data 198 | self.targets = new_targets 199 | self.filenames = new_filenames 200 | 201 | class AddGaussianNoise(object): 202 | def __init__(self, mean=0., std=1.): 203 | self.std = std 204 | self.mean = mean 205 | 206 | def __call__(self, tensor): 207 | return tensor + torch.randn(tensor.size()) * self.std + self.mean 208 | 209 | def __repr__(self): 210 | return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std) -------------------------------------------------------------------------------- /libs/data/dataloader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torchvision 4 | from torch.utils.data import Dataset, DataLoader, ConcatDataset 5 | from torchvision import transforms 6 | import os 7 | from PIL import Image 8 | 9 | # Image statistics 10 | RGB_statistics = { 11 | "iNaturalist18": {"mean": [0.466, 0.471, 0.380], "std": [0.195, 0.194, 0.192]}, 12 | "default": {"mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]}, 13 | } 14 | 15 | # Data transformation with augmentation 16 | def get_data_transform(split, rgb_mean, rbg_std, key="default", jitter=True, special_aug=False): 17 | if split == "train" and key=="default" : 18 | if jitter and not(special_aug): 19 | return transforms.Compose( 20 | [ 21 | transforms.RandomResizedCrop(224), 22 | transforms.RandomHorizontalFlip(), 23 | transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0), 24 | transforms.ToTensor(), 25 | transforms.Normalize(rgb_mean, rbg_std), 26 | ] 27 | ) 28 | elif special_aug: 29 | return transforms.Compose( 30 | [ transforms.RandomApply([transforms.ColorJitter(brightness=(0.1, 0.3), contrast=(0.1, 0.3), saturation=(0.1, 0.3), hue=(0.1, 0.3))], p=0.8), 31 | transforms.RandomResizedCrop(224), 32 | transforms.RandomHorizontalFlip(), 33 | transforms.ToTensor(), 34 | transforms.Normalize(rgb_mean, rbg_std), 35 | AddGaussianNoise(0., 0.01), 36 | ] 37 | ) 38 | else : 39 | return transforms.Compose( 40 | [ 41 | transforms.RandomResizedCrop(224), 42 | transforms.RandomHorizontalFlip(), 43 | transforms.ToTensor(), 44 | transforms.Normalize(rgb_mean, rbg_std), 45 | ] 46 | ) 47 | elif split == "train" and key=="iNaturalist18" : 48 | return transforms.Compose( 49 | [ 50 | transforms.RandomResizedCrop(224), 51 | transforms.RandomHorizontalFlip(), 52 | transforms.ToTensor(), 53 | transforms.Normalize(rgb_mean, rbg_std), 54 | ] 55 | ) 56 | else : 57 | data_transforms = { 58 | "val": transforms.Compose( 59 | [ 60 | transforms.Resize(256), 61 | transforms.CenterCrop(224), 62 | transforms.ToTensor(), 63 | transforms.Normalize(rgb_mean, rbg_std), 64 | ] 65 | ), 66 | "test": transforms.Compose( 67 | [ 68 | transforms.Resize(256), 69 | transforms.CenterCrop(224), 70 | transforms.ToTensor(), 71 | transforms.Normalize(rgb_mean, rbg_std), 72 | ] 73 | ), 74 | } 75 | return data_transforms[split] 76 | 77 | 78 | # Dataset 79 | class LT_Dataset(Dataset): 80 | def __init__(self, root, txt, transform=None, template=None, top_k=None): 81 | self.img_path = [] 82 | self.labels = [] 83 | self.transform = transform 84 | with open(txt) as f: 85 | for line in f: 86 | rootalt = root 87 | self.img_path.append(os.path.join(rootalt, line.split()[0])) 88 | self.labels.append(int(line.split()[1])) 89 | # select top k class 90 | if top_k: 91 | # only select top k in training, in case train/val/test not matching. 92 | if "train" in txt: 93 | max_len = max(self.labels) + 1 94 | dist = [[i, 0] for i in range(max_len)] 95 | for i in self.labels: 96 | dist[i][-1] += 1 97 | dist.sort(key=lambda x: x[1], reverse=True) 98 | # saving 99 | torch.save(dist, template + "_top_{}_mapping".format(top_k)) 100 | else: 101 | # loading 102 | dist = torch.load(template + "_top_{}_mapping".format(top_k)) 103 | selected_labels = {item[0]: i for i, item in enumerate(dist[:top_k])} 104 | # replace original path and labels 105 | self.new_img_path = [] 106 | self.new_labels = [] 107 | for path, label in zip(self.img_path, self.labels): 108 | if label in selected_labels: 109 | self.new_img_path.append(path) 110 | self.new_labels.append(selected_labels[label]) 111 | self.img_path = self.new_img_path 112 | self.labels = self.new_labels 113 | self.img_num_list = list(np.unique(self.labels, return_counts=True)[1]) 114 | def __len__(self): 115 | return len(self.labels) 116 | 117 | def __getitem__(self, index): 118 | 119 | path = self.img_path[index] 120 | label = self.labels[index] 121 | 122 | with open(path, "rb") as f: 123 | sample = Image.open(f).convert("RGB") 124 | 125 | if self.transform is not None: 126 | sample = self.transform(sample) 127 | 128 | return sample, label, index 129 | 130 | 131 | # Load datasets 132 | def load_data( 133 | data_root, 134 | dataset, 135 | phase, 136 | batch_size, 137 | top_k_class=None, 138 | sampler_dic=None, 139 | num_workers=4, 140 | shuffle=True, 141 | cifar_imb_ratio=None, 142 | imb_type="exp", 143 | class_order=None, 144 | balanced=None, 145 | type_of_val="vft", #val_from_test or val_from_train or val_is_test 146 | special_aug=False, 147 | seed = 1, 148 | jitter = True, 149 | ): 150 | if imb_type == None: 151 | imb_type = "exp" 152 | txt_split = phase 153 | txt = "./libs/data/%s/%s_%s.txt" % (dataset, dataset, txt_split) 154 | template = "./libs/data/%s/%s" % (dataset, dataset) 155 | 156 | print("Loading data from %s" % (txt)) 157 | 158 | if "CIFAR" in dataset: 159 | from libs.data.ImbalanceCIFAR import IMBALANCECIFAR10, IMBALANCECIFAR100 160 | 161 | if dataset in ["iNaturalist18","iNaturalist18_insecta"]: 162 | print("===> Loading iNaturalist18 statistics") 163 | key = "iNaturalist18" 164 | else: 165 | key = "default" 166 | 167 | 168 | if dataset == "CIFAR100_LT": 169 | if cifar_imb_ratio == 1: 170 | print("====> CIFAR100 No Imbalance") 171 | else: 172 | print("====> CIFAR100 Imbalance Ratio: ", cifar_imb_ratio) 173 | set_ = IMBALANCECIFAR100( 174 | phase, imbalance_ratio=cifar_imb_ratio, root=data_root, imb_type=imb_type, class_order=class_order, balanced=balanced, special_aug=special_aug, seed=seed 175 | ) 176 | else: 177 | rgb_mean, rgb_std = RGB_statistics[key]["mean"], RGB_statistics[key]["std"] 178 | if phase not in ["train", "val"]: 179 | transform = get_data_transform("test", rgb_mean, rgb_std, key) 180 | else: 181 | transform = get_data_transform(phase, rgb_mean, rgb_std, key, jitter, special_aug) 182 | print("Use data transformation:", transform) 183 | 184 | set_ = LT_Dataset( 185 | data_root, txt, transform, template=template, top_k=top_k_class 186 | ) 187 | 188 | print(len(set_)) 189 | 190 | if sampler_dic and phase == "train": 191 | print("=====> Using sampler: ", sampler_dic["sampler"]) 192 | print("=====> Sampler parameters: ", sampler_dic["params"]) 193 | return DataLoader( 194 | dataset=set_, 195 | batch_size=batch_size, 196 | shuffle=False, 197 | sampler=sampler_dic["sampler"](set_, **sampler_dic["params"]), 198 | num_workers=num_workers, 199 | ) 200 | elif phase == "train": 201 | print("=====> No sampler.") 202 | print("=====> Shuffle is %s." % (shuffle)) 203 | return DataLoader( 204 | dataset=set_, 205 | batch_size=batch_size, 206 | shuffle=shuffle, 207 | num_workers=num_workers, 208 | ) 209 | else: 210 | print("=====> No sampler.") 211 | print("=====> Shuffle is %s." % (shuffle)) 212 | return DataLoader( 213 | dataset=set_, 214 | batch_size=batch_size, 215 | shuffle=True, 216 | num_workers=num_workers, 217 | ) 218 | 219 | # ECBD - # For repo specific to CBD paper with much more detailed instructions, check https://github.com/rahulvigneswaran/Class-Balanced-Distillation-for-Long-Tailed-Visual-Recognition.pytorch 220 | class AddGaussianNoise(object): 221 | def __init__(self, mean=0., std=1.): 222 | self.std = std 223 | self.mean = mean 224 | 225 | def __call__(self, tensor): 226 | return tensor + torch.randn(tensor.size()) * self.std + self.mean 227 | 228 | def __repr__(self): 229 | return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std) -------------------------------------------------------------------------------- /libs/loss/CosineDistill.py: -------------------------------------------------------------------------------- 1 | # Cosine Distillation loss 2 | 3 | # Imports 4 | import torch.nn as nn 5 | 6 | class CosineDistill(nn.Module): 7 | def __init__(self, beta=100): 8 | super(CosineDistill, self).__init__() 9 | self.beta = beta 10 | 11 | def forward(self, student, teacher): 12 | cos = nn.CosineSimilarity(dim=1) 13 | return self.beta*(1-cos(student, teacher)).mean(dim=0) 14 | 15 | def create_loss(*args): 16 | print("Loading Cosine Distance Loss.") 17 | return CosineDistill(*args) 18 | -------------------------------------------------------------------------------- /libs/loss/SoftmaxLoss.py: -------------------------------------------------------------------------------- 1 | # Softmax Loss 2 | 3 | # Imports 4 | import torch.nn as nn 5 | 6 | 7 | def create_loss(): 8 | """Generic Cross entropy loss 9 | """ 10 | print("Loading Softmax Loss.") 11 | return nn.CrossEntropyLoss() 12 | -------------------------------------------------------------------------------- /libs/models/CosineDotProductClassifier.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | import torch.nn as nn 3 | from os import path 4 | import torch 5 | import torch.nn.functional as F 6 | 7 | # Custom imports 8 | from libs.utils.utils import * 9 | 10 | class CosineDotProduct_Classifier(nn.Module): 11 | def __init__(self, num_classes=1000, feat_dim=2048, scale=10.0, *args): 12 | super(CosineDotProduct_Classifier, self).__init__() 13 | self.fc = nn.Linear(feat_dim, num_classes) 14 | self.scale = nn.Parameter(torch.FloatTensor(1).fill_(scale), requires_grad=True) 15 | 16 | def forward(self, x, *args): 17 | x = F.softplus(self.scale)*(torch.mm(x, F.normalize(self.fc.weight.T, dim=1))) + self.fc.bias # (Batch, feature) x (out, feature).T -> (Batch, out) 18 | wandb_log({"Classifier Scale": self.scale.item()}) 19 | return x 20 | 21 | 22 | def create_model(feat_dim, num_classes=1000, scale=10.0, pretrain=False, pretrain_dir=None, *args): 23 | """Initialize the model 24 | 25 | Args: 26 | feat_dim (int): output dimension of the previous feature extractor 27 | num_classes (int, optional): Number of classes. Defaults to 1000. 28 | 29 | Returns: 30 | Class: Model 31 | """ 32 | print("Loading Cosine Dot Product Classifier.") 33 | clf = CosineDotProduct_Classifier(num_classes, feat_dim, scale) 34 | 35 | if pretrain: 36 | if path.exists(pretrain_dir): 37 | print("===> Load Pretrain Initialization for CosineDotProductClassfier") 38 | weights = torch.load(pretrain_dir)["state_dict_best"]["classifier"] 39 | 40 | weights = { 41 | k: weights["module." + k] 42 | if "module." + k in weights 43 | else clf.state_dict()[k] 44 | for k in clf.state_dict() 45 | } 46 | clf.load_state_dict(weights) 47 | else: 48 | raise Exception(f"Pretrain path doesn't exist!!-{pretrain_dir}") 49 | else: 50 | print("===> Train classifier from the scratch") 51 | 52 | return clf 53 | -------------------------------------------------------------------------------- /libs/models/DotProductClassifier.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | import torch.nn as nn 3 | from os import path 4 | import torch 5 | import torch.nn.functional as F 6 | 7 | class DotProduct_Classifier(nn.Module): 8 | def __init__(self, num_classes=1000, feat_dim=2048, *args): 9 | super(DotProduct_Classifier, self).__init__() 10 | self.fc = nn.Linear(feat_dim, num_classes) 11 | 12 | def forward(self, x, *args): 13 | x = self.fc(x) 14 | return x 15 | 16 | 17 | def create_model(feat_dim, num_classes=1000, pretrain=False, pretrain_dir=None, *args): 18 | """Initialize the model 19 | 20 | Args: 21 | feat_dim (int): output dimension of the previous feature extractor 22 | num_classes (int, optional): Number of classes. Defaults to 1000. 23 | 24 | Returns: 25 | Class: Model 26 | """ 27 | print("Loading Dot Product Classifier.") 28 | clf = DotProduct_Classifier(num_classes, feat_dim) 29 | 30 | if pretrain: 31 | if path.exists(pretrain_dir): 32 | print("===> Load Pretrain Initialization for DotProductClassfier") 33 | weights = torch.load(pretrain_dir)["state_dict_best"]["classifier"] 34 | 35 | weights = { 36 | k: weights["module." + k] 37 | if "module." + k in weights 38 | else clf.state_dict()[k] 39 | for k in clf.state_dict() 40 | } 41 | clf.load_state_dict(weights) 42 | else: 43 | raise Exception(f"Pretrain path doesn't exist!!--{pretrain_dir}") 44 | else: 45 | print("===> Train classifier from the scratch") 46 | 47 | return clf 48 | -------------------------------------------------------------------------------- /libs/models/ResNet32Feature.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | from os import path 3 | from collections import OrderedDict 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import torch.nn.init as init 8 | 9 | 10 | def _weights_init(m): 11 | classname = m.__class__.__name__ 12 | if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d): 13 | init.kaiming_normal_(m.weight) 14 | 15 | 16 | class LambdaLayer(nn.Module): 17 | def __init__(self, lambd): 18 | super(LambdaLayer, self).__init__() 19 | self.lambd = lambd 20 | 21 | def forward(self, x): 22 | return self.lambd(x) 23 | 24 | 25 | class BasicBlock(nn.Module): 26 | expansion = 1 27 | 28 | def __init__(self, in_planes, planes, stride=1, option="A"): 29 | super(BasicBlock, self).__init__() 30 | self.conv1 = nn.Conv2d( 31 | in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False 32 | ) 33 | self.bn1 = nn.BatchNorm2d(planes) 34 | self.conv2 = nn.Conv2d( 35 | planes, planes, kernel_size=3, stride=1, padding=1, bias=False 36 | ) 37 | self.bn2 = nn.BatchNorm2d(planes) 38 | 39 | self.shortcut = nn.Sequential() 40 | if stride != 1 or in_planes != planes: 41 | if option == "A": 42 | """ 43 | For CIFAR10 ResNet paper uses option A. 44 | """ 45 | self.shortcut = LambdaLayer( 46 | lambda x: F.pad( 47 | x[:, :, ::2, ::2], 48 | (0, 0, 0, 0, planes // 4, planes // 4), 49 | "constant", 50 | 0, 51 | ) 52 | ) 53 | elif option == "B": 54 | self.shortcut = nn.Sequential( 55 | nn.Conv2d( 56 | in_planes, 57 | self.expansion * planes, 58 | kernel_size=1, 59 | stride=stride, 60 | bias=False, 61 | ), 62 | nn.BatchNorm2d(self.expansion * planes), 63 | ) 64 | 65 | def forward(self, x): 66 | out = F.relu(self.bn1(self.conv1(x))) 67 | out = self.bn2(self.conv2(out)) 68 | out += self.shortcut(x) 69 | out = F.relu(out) 70 | return out 71 | 72 | 73 | class BBN_ResNet_Cifar(nn.Module): 74 | """ResNet32 from the "BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition (CVPR 2020)" """ 75 | 76 | def __init__(self, block, num_blocks): 77 | """Initialize 78 | #FIXME 79 | Args: 80 | block ([type]): [description] 81 | num_blocks ([type]): [description] 82 | """ 83 | super(BBN_ResNet_Cifar, self).__init__() 84 | self.in_planes = 16 85 | 86 | self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False) 87 | self.bn1 = nn.BatchNorm2d(16) 88 | self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) 89 | self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) 90 | self.layer3 = self._make_layer(block, 64, num_blocks[2] - 1, stride=2) 91 | self.cb_block = block(self.in_planes, self.in_planes, stride=1) 92 | self.rb_block = block(self.in_planes, self.in_planes, stride=1) 93 | 94 | self.apply(_weights_init) 95 | 96 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 97 | 98 | def load_model(self, pretrain_dir): 99 | """Load a pre-trained model 100 | 101 | Args: 102 | pretrain_dir (str): path of pretrained model 103 | """ 104 | print(f"Loading Backbone pretrain model from {pretrain_dir}......") 105 | model_dict = self.state_dict() 106 | pretrain_dict = torch.load(pretrain_dir)["state_dict_best"]["feat_model"] 107 | 108 | new_dict = OrderedDict() 109 | 110 | # Removing FC and Classifier layers 111 | for k, v in pretrain_dict.items(): 112 | if k.startswith("module"): 113 | k = k[7:] 114 | if "fc" not in k and "classifier" not in k: 115 | new_dict[k] = v 116 | 117 | model_dict.update(new_dict) 118 | self.load_state_dict(model_dict) 119 | print("Backbone model has been loaded......") 120 | 121 | def _make_layer(self, block, planes, num_blocks, stride, add_flag=True): 122 | strides = [stride] + [1] * (num_blocks - 1) 123 | layers = [] 124 | for stride in strides: 125 | layers.append(block(self.in_planes, planes, stride)) 126 | self.in_planes = planes * block.expansion 127 | 128 | return nn.Sequential(*layers) 129 | 130 | def forward(self, x, **kwargs): 131 | out = F.relu(self.bn1(self.conv1(x))) 132 | out = self.layer1(out) 133 | out = self.layer2(out) 134 | out = self.layer3(out) 135 | if "feature_cb" in kwargs: 136 | out = self.cb_block(out) 137 | return out 138 | elif "feature_rb" in kwargs: 139 | out = self.rb_block(out) 140 | return out 141 | 142 | out1 = self.cb_block(out) 143 | out2 = self.rb_block(out) 144 | out = torch.cat((out1, out2), dim=1) 145 | 146 | out = self.avgpool(out) 147 | out = out.view(out.shape[0], -1) 148 | 149 | 150 | return out 151 | 152 | def create_model(pretrain=False, pretrain_dir=None, *args): 153 | """Initialize/load the model 154 | 155 | Args: 156 | pretrain (bool, optional): Use pre-trained model?. Defaults to False. 157 | pretrain_dir (str, optional): Directory of the pre-trained model. Defaults to None. 158 | 159 | Returns: 160 | class: Model 161 | """ 162 | 163 | print("Loading ResNet 32 Feature Model.") 164 | resnet32 = BBN_ResNet_Cifar(BasicBlock, [5, 5, 5]) 165 | 166 | if pretrain: 167 | if path.exists(pretrain_dir): 168 | print("===> Load Pretrain Initialization for ResNet32") 169 | resnet32.load_model(pretrain_dir=pretrain_dir) 170 | else: 171 | raise Exception(f"Pretrain path doesn't exist!!-{pretrain_dir}") 172 | else: 173 | print("===> Train backbone from the scratch") 174 | 175 | return resnet32 176 | -------------------------------------------------------------------------------- /libs/models/ResNext50Feature.py: -------------------------------------------------------------------------------- 1 | """Copyright (c) Facebook, Inc. and its affiliates. 2 | All rights reserved. 3 | 4 | This source code is licensed under the license found in the 5 | LICENSE file in the root directory of this source tree. 6 | """ 7 | 8 | # Imports 9 | from os import path 10 | from collections import OrderedDict 11 | import torch 12 | 13 | # Custom imports 14 | from libs.models.ResNextFeature import * 15 | from libs.utils.utils import * 16 | 17 | def create_model(pretrain=False, pretrain_dir=None, *args): 18 | """Initialize/load the model 19 | 20 | Args: 21 | pretrain (bool, optional): Use pre-trained model?. Defaults to False. 22 | pretrain_dir (str, optional): Directory of the pre-trained model. Defaults to None. 23 | 24 | Returns: 25 | class: Model 26 | """ 27 | 28 | print("Loading ResNext 50 Feature Model.") 29 | resnext50 = ResNext(Bottleneck, [3, 4, 6, 3], use_fc=False, dropout=None, 30 | groups=32, width_per_group=4, last_relu=True) 31 | 32 | if pretrain: 33 | if path.exists(pretrain_dir): 34 | print("===> Load Pretrain Initialization for ResNext50") 35 | model_dict = resnext50.state_dict() 36 | new_dict = load_model(pretrain_dir=pretrain_dir) 37 | model_dict.update(new_dict) 38 | resnext50.load_state_dict(model_dict) 39 | print("Backbone model has been loaded......") 40 | 41 | else: 42 | raise Exception(f"Pretrain path doesn't exist!!-{pretrain_dir}") 43 | else: 44 | print("===> Train backbone from the scratch") 45 | 46 | return resnext50 47 | 48 | def load_model(pretrain_dir): 49 | """Load a pre-trained model 50 | 51 | Args: 52 | pretrain_dir (str): path of pretrained model 53 | """ 54 | print(f"Loading Backbone pretrain model from {pretrain_dir}......") 55 | pretrain_dict = torch.load(pretrain_dir)["state_dict_best"]["feat_model"] 56 | 57 | new_dict = OrderedDict() 58 | 59 | # Removing FC and Classifier layers 60 | for k, v in pretrain_dict.items(): 61 | if k.startswith("module"): 62 | k = k[7:] 63 | if "fc" not in k and "classifier" not in k: 64 | new_dict[k] = v 65 | 66 | return new_dict 67 | -------------------------------------------------------------------------------- /libs/models/ResNextFeature.py: -------------------------------------------------------------------------------- 1 | """Copyright (c) Facebook, Inc. and its affiliates. 2 | All rights reserved. 3 | 4 | This source code is licensed under the license found in the 5 | LICENSE file in the root directory of this source tree. 6 | """ 7 | 8 | 9 | import math 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | from libs.utils.utils import * 13 | 14 | def conv3x3(in_planes, out_planes, stride=1): 15 | """3x3 convolution with padding""" 16 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 17 | padding=1, bias=False) 18 | 19 | class BasicBlock(nn.Module): 20 | expansion = 1 21 | 22 | def __init__(self, inplanes, planes, stride=1, downsample=None): 23 | super(BasicBlock, self).__init__() 24 | self.conv1 = conv3x3(inplanes, planes, stride) 25 | self.bn1 = nn.BatchNorm2d(planes) 26 | self.relu = nn.ReLU(inplace=True) 27 | self.conv2 = conv3x3(planes, planes) 28 | self.bn2 = nn.BatchNorm2d(planes) 29 | self.downsample = downsample 30 | self.stride = stride 31 | 32 | def forward(self, x): 33 | residual = x 34 | 35 | out = self.conv1(x) 36 | out = self.bn1(out) 37 | out = self.relu(out) 38 | 39 | out = self.conv2(out) 40 | out = self.bn2(out) 41 | 42 | if self.downsample is not None: 43 | residual = self.downsample(x) 44 | 45 | out += residual 46 | out = self.relu(out) 47 | 48 | return out 49 | 50 | class Bottleneck(nn.Module): 51 | expansion = 4 52 | 53 | def __init__(self, inplanes, planes, stride=1, downsample=None, 54 | groups=1, base_width=64, is_last=False, last_relu=True): 55 | super(Bottleneck, self).__init__() 56 | width = int(planes * (base_width / 64.)) * groups 57 | self.conv1 = nn.Conv2d(inplanes, width, kernel_size=1, bias=False) 58 | self.bn1 = nn.BatchNorm2d(width) 59 | self.conv2 = nn.Conv2d(width, width, kernel_size=3, stride=stride, 60 | groups=groups, padding=1, bias=False) 61 | self.bn2 = nn.BatchNorm2d(width) 62 | self.conv3 = nn.Conv2d(width, planes * 4, kernel_size=1, bias=False) 63 | self.bn3 = nn.BatchNorm2d(planes * 4) 64 | self.relu = nn.ReLU(inplace=True) 65 | self.downsample = downsample 66 | self.stride = stride 67 | self.is_last = is_last 68 | self.last_relu = last_relu 69 | 70 | def forward(self, x): 71 | residual = x 72 | 73 | out = self.conv1(x) 74 | out = self.bn1(out) 75 | out = self.relu(out) 76 | 77 | out = self.conv2(out) 78 | out = self.bn2(out) 79 | out = self.relu(out) 80 | 81 | out = self.conv3(out) 82 | out = self.bn3(out) 83 | 84 | if self.downsample is not None: 85 | residual = self.downsample(x) 86 | 87 | out += residual 88 | if self.is_last and (not self.last_relu): 89 | pass 90 | else: 91 | out = self.relu(out) 92 | 93 | return out 94 | 95 | class ResNext(nn.Module): 96 | 97 | def __init__(self, block, layers, groups=1, width_per_group=64, use_fc=False, dropout=None, 98 | use_glore=False, use_gem=False, last_relu=True): 99 | self.inplanes = 64 100 | super(ResNext, self).__init__() 101 | 102 | self.groups = groups 103 | self.base_width = width_per_group 104 | 105 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 106 | bias=False) 107 | self.bn1 = nn.BatchNorm2d(64) 108 | self.relu = nn.ReLU(inplace=True) 109 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 110 | self.layer1 = self._make_layer(block, 64, layers[0]) 111 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 112 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 113 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, is_last=True, last_relu=last_relu) 114 | self.avgpool = nn.AvgPool2d(7, stride=1) 115 | 116 | self.use_fc = use_fc 117 | self.use_dropout = True if dropout else False 118 | 119 | if self.use_fc: 120 | print('Using fc.') 121 | self.fc_add = nn.Linear(512*block.expansion, 512) 122 | 123 | if self.use_dropout: 124 | print('Using dropout.') 125 | self.dropout = nn.Dropout(p=dropout) 126 | 127 | for m in self.modules(): 128 | if isinstance(m, nn.Conv2d): 129 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 130 | m.weight.data.normal_(0, math.sqrt(2. / n)) 131 | elif isinstance(m, nn.BatchNorm2d): 132 | m.weight.data.fill_(1) 133 | m.bias.data.zero_() 134 | 135 | 136 | def _make_layer(self, block, planes, blocks, stride=1, is_last=False, last_relu=True): 137 | downsample = None 138 | if stride != 1 or self.inplanes != planes * block.expansion: 139 | downsample = nn.Sequential( 140 | nn.Conv2d(self.inplanes, planes * block.expansion, 141 | kernel_size=1, stride=stride, bias=False), 142 | nn.BatchNorm2d(planes * block.expansion), 143 | ) 144 | 145 | layers = [] 146 | layers.append(block(self.inplanes, planes, stride, downsample, 147 | groups=self.groups, base_width=self.base_width)) 148 | self.inplanes = planes * block.expansion 149 | for i in range(1, blocks): 150 | layers.append(block(self.inplanes, planes, 151 | groups=self.groups, base_width=self.base_width, 152 | is_last=(is_last and i == blocks-1), last_relu=last_relu)) 153 | 154 | return nn.Sequential(*layers) 155 | 156 | def forward(self, x, *args): 157 | x = self.conv1(x) 158 | x = self.bn1(x) 159 | x = self.relu(x) 160 | x = self.maxpool(x) 161 | 162 | x = self.layer1(x) 163 | x = self.layer2(x) 164 | x = self.layer3(x) 165 | x = self.layer4(x) 166 | 167 | x = self.avgpool(x) 168 | 169 | x = x.view(x.size(0), -1) 170 | 171 | if self.use_fc: 172 | x = F.relu(self.fc_add(x)) 173 | 174 | if self.use_dropout: 175 | x = self.dropout(x) 176 | 177 | return x -------------------------------------------------------------------------------- /libs/models/ecbd_converter.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | import torch.nn as nn 3 | from os import path 4 | import torch 5 | import torch.nn.functional as F 6 | 7 | # Custom imports 8 | from libs.utils.utils import * 9 | 10 | class ecbd_converter(nn.Module): 11 | """A linear layer that converts the student's feature to the same dimension as the concatenated teachers' dimensions. 12 | 13 | Args: 14 | nn ([type]): [description] 15 | """ 16 | def __init__(self, feat_in, feat_out, *args): 17 | super(ecbd_converter, self).__init__() 18 | self.fc = nn.Linear(feat_in, feat_out) 19 | 20 | def forward(self, x, *args): 21 | return self.fc(x) 22 | 23 | 24 | def create_model(feat_in, feat_out, pretrain=False, pretrain_dir=None, *args): 25 | """Initialize the model 26 | 27 | Args: 28 | feat_dim (int): output dimension of the previous feature extractor 29 | num_classes (int, optional): Number of classes. Defaults to 1000. 30 | 31 | Returns: 32 | Class: Model 33 | """ 34 | print("ECBD Converter.") 35 | clf = ecbd_converter(feat_in, feat_out) 36 | 37 | if pretrain: 38 | if path.exists(pretrain_dir): 39 | print("===> Load Pretrain Initialization for CosineDotProductClassfier") 40 | weights = torch.load(pretrain_dir)["state_dict_best"]["classifier"] 41 | 42 | weights = { 43 | k: weights["module." + k] 44 | if "module." + k in weights 45 | else clf.state_dict()[k] 46 | for k in clf.state_dict() 47 | } 48 | clf.load_state_dict(weights) 49 | else: 50 | raise Exception("Pretrain path doesn't exist!!") 51 | else: 52 | print("===> Train classifier from the scratch") 53 | 54 | return clf 55 | -------------------------------------------------------------------------------- /libs/samplers/ClassAwareSampler.py: -------------------------------------------------------------------------------- 1 | """Copyright (c) Facebook, Inc. and its affiliates. 2 | All rights reserved. 3 | 4 | This source code is licensed under the license found in the 5 | LICENSE file in the root directory of this source tree. 6 | 7 | Portions of the source code are from the OLTR project which 8 | notice below and in LICENSE in the root directory of 9 | this source tree. 10 | 11 | Copyright (c) 2019, Zhongqi Miao 12 | All rights reserved. 13 | """ 14 | 15 | import random 16 | import numpy as np 17 | from torch.utils.data.sampler import Sampler 18 | import pdb 19 | 20 | ################################## 21 | ## Class-aware sampling, partly implemented by frombeijingwithlove 22 | ################################## 23 | 24 | class RandomCycleIter: 25 | 26 | def __init__ (self, data, test_mode=False): 27 | self.data_list = list(data) 28 | self.length = len(self.data_list) 29 | self.i = self.length - 1 30 | self.test_mode = test_mode 31 | 32 | def __iter__ (self): 33 | return self 34 | 35 | def __next__ (self): 36 | self.i += 1 37 | 38 | if self.i == self.length: 39 | self.i = 0 40 | if not self.test_mode: 41 | random.shuffle(self.data_list) 42 | 43 | return self.data_list[self.i] 44 | 45 | def class_aware_sample_generator (cls_iter, data_iter_list, n, num_samples_cls=1): 46 | 47 | i = 0 48 | j = 0 49 | while i < n: 50 | 51 | # yield next(data_iter_list[next(cls_iter)]) 52 | 53 | if j >= num_samples_cls: 54 | j = 0 55 | 56 | if j == 0: 57 | temp_tuple = next(zip(*[data_iter_list[next(cls_iter)]]*num_samples_cls)) 58 | yield temp_tuple[j] 59 | else: 60 | yield temp_tuple[j] 61 | 62 | i += 1 63 | j += 1 64 | 65 | class ClassAwareSampler (Sampler): 66 | 67 | def __init__(self, data_source, num_samples_cls=1,): 68 | num_classes = len(np.unique(data_source.labels)) 69 | self.class_iter = RandomCycleIter(range(num_classes)) 70 | cls_data_list = [list() for _ in range(num_classes)] 71 | for i, label in enumerate(data_source.labels): 72 | cls_data_list[label].append(i) 73 | self.data_iter_list = [RandomCycleIter(x) for x in cls_data_list] 74 | self.num_samples = max([len(x) for x in cls_data_list]) * len(cls_data_list) 75 | self.num_samples_cls = num_samples_cls 76 | 77 | def __iter__ (self): 78 | return class_aware_sample_generator(self.class_iter, self.data_iter_list, 79 | self.num_samples, self.num_samples_cls) 80 | 81 | def __len__ (self): 82 | return self.num_samples 83 | 84 | def get_sampler(): 85 | return ClassAwareSampler 86 | 87 | ################################## -------------------------------------------------------------------------------- /libs/utils/default_config.yaml: -------------------------------------------------------------------------------- 1 | criterions: 2 | ClassifierLoss: 3 | def_file: ./loss/SoftmaxLoss.py 4 | loss_params: {} 5 | optim_params: null 6 | weight: 1.0 7 | 8 | endlr: 0.0 9 | networks: 10 | classifier: 11 | def_file: ./models/DotProductClassifier.py 12 | fix: false 13 | optim_params: {} 14 | scheduler_params: {} 15 | params: {} 16 | feat_model: 17 | def_file: ./models/ResNet32Feature.py 18 | fix: false 19 | optim_params: {lr: 0.2, momentum: 0.9, weight_decay: 0.0005} 20 | scheduler_params: {coslr: true} 21 | params: {pretrain: False, pretrain_dir: None} 22 | 23 | shuffle: false 24 | training_opt: 25 | backbone: 26 | batch_size: 27 | accumulation_step: 1 28 | dataset: 29 | display_step: 10 30 | num_classes: 31 | cifar_imb_ratio: # 0.01, 0.02, 0.1 for 100, 50, 10 32 | num_epochs: 33 | num_workers: 5 34 | open_threshold: 0.1 35 | sampler: null 36 | stage: 37 | log_dir: 38 | 39 | wandb_tags: ["",""] 40 | 41 | pg: 42 | generate: False 43 | 44 | 45 | 46 | -------------------------------------------------------------------------------- /libs/utils/globals.py: -------------------------------------------------------------------------------- 1 | # Loggers 2 | wandb_log = False 3 | log_offline = False 4 | log_dir = None #For offline logging 5 | 6 | # x-axis 7 | epoch_global = 0 8 | step_global = 0 9 | 10 | # seed 11 | seed = 1 12 | 13 | # resume 14 | resume = False 15 | 16 | -------------------------------------------------------------------------------- /libs/utils/logger.py: -------------------------------------------------------------------------------- 1 | """Copyright (c) Facebook, Inc. and its affiliates. 2 | All rights reserved. 3 | 4 | This source code is licensed under the license found in the 5 | LICENSE file in the root directory of this source tree. 6 | """ 7 | 8 | import os 9 | import yaml 10 | import csv 11 | import h5py 12 | import torch 13 | 14 | import libs.utils.globals as g 15 | 16 | 17 | 18 | class Logger(object): 19 | def __init__(self, logdir): 20 | if g.log_offline: 21 | self.logdir = logdir 22 | if not os.path.isdir(logdir): 23 | os.makedirs(logdir) 24 | self.cfg_file = os.path.join(self.logdir, 'cfg.yaml') 25 | self.acc_file = os.path.join(self.logdir, 'acc.csv') 26 | self.loss_file = os.path.join(self.logdir, 'loss.csv') 27 | self.ws_file = os.path.join(self.logdir, 'ws.h5') 28 | self.acc_keys = None 29 | self.loss_keys = None 30 | self.logging_ws = False 31 | 32 | def log_cfg(self, cfg): 33 | if g.log_offline: 34 | print('===> Saving cfg parameters to: ', self.cfg_file) 35 | with open(self.cfg_file, 'w') as f: 36 | yaml.dump(cfg, f) 37 | 38 | def log_acc(self, accs): 39 | if g.log_offline: 40 | if self.acc_keys is None and not(g.resume): 41 | self.acc_keys = [k for k in accs.keys()] 42 | with open(self.acc_file, 'w') as f: 43 | writer = csv.DictWriter(f, fieldnames=self.acc_keys) 44 | writer.writeheader() 45 | writer.writerow(accs) 46 | else: 47 | if self.acc_keys is None: 48 | self.acc_keys = [k for k in accs.keys()] 49 | with open(self.acc_file, 'a') as f: 50 | writer = csv.DictWriter(f, fieldnames=self.acc_keys) 51 | writer.writerow(accs) 52 | 53 | 54 | def log_loss(self, losses): 55 | if g.log_offline: 56 | valid_losses = losses 57 | if self.loss_keys is None and not(g.resume): 58 | self.loss_keys = [k for k in valid_losses.keys()] 59 | with open(self.loss_file, 'w') as f: 60 | writer = csv.DictWriter(f, fieldnames=self.loss_keys) 61 | writer.writeheader() 62 | writer.writerow(valid_losses) 63 | else: 64 | if self.loss_keys is None: 65 | self.loss_keys = [k for k in valid_losses.keys()] 66 | with open(self.loss_file, 'a') as f: 67 | writer = csv.DictWriter(f, fieldnames=self.loss_keys) 68 | writer.writerow(valid_losses) 69 | 70 | def log_ws(self, e, ws): 71 | if g.log_offline: 72 | mode = 'a' if self.logging_ws else 'w' 73 | self.logging_ws = True 74 | 75 | key = 'Epoch{:02d}'.format(e) 76 | with h5py.File(self.ws_file, mode) as f: 77 | g = f.create_group(key) 78 | for k, v in ws.items(): 79 | g.create_dataset(k, data=v) 80 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # Imports 2 | import os 3 | import argparse 4 | import yaml 5 | import resource 6 | import torch 7 | import random 8 | import numpy as np 9 | import yaml 10 | 11 | # Increase resource limit 12 | rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) 13 | resource.setrlimit(resource.RLIMIT_NOFILE, (4048, rlimit[1])) 14 | 15 | # argparsing 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument("--seed", default=1, type=int, help="Select seed for fixing it.") 18 | parser.add_argument("--gpu", default="0,1,2,3", type=str, help="Select the GPUs to be used.") 19 | 20 | parser.add_argument("--experiment", default=0.1, type=float, help="Experiment number (Check 'libs/utils/experiment_maker.py').") 21 | parser.add_argument("--dataset", default=0, type=int, help="Dataset number. Choice: 0 - CIFAR100, 1 - mini-imagenet.") 22 | parser.add_argument("--imbalance", default=1, type=int, help="Select Imbalance factor. Choice: 0: 1, 1: 100, 2: 50, 3: 10.") 23 | parser.add_argument("--type_of_val", type=str, default="vit", help="Choose which dataset split to use. Choice: vt: val_from_test, vtr: val_from_train, vit: val_is_test") 24 | 25 | parser.add_argument("--cv1", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 26 | parser.add_argument("--cv2", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 27 | parser.add_argument("--cv3", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 28 | parser.add_argument("--cv4", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 29 | parser.add_argument("--cv5", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 30 | parser.add_argument("--cv6", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 31 | parser.add_argument("--cv7", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 32 | parser.add_argument("--cv8", type=str, default="0.9", help="Custom variable to use in experiments - purpose changes according to the experiment.") 33 | parser.add_argument("--cv9", type=str, default="1", help="Custom variable to use in experiments - purpose changes according to the experiment.") 34 | 35 | parser.add_argument("--train", default=False, action="store_true", help="Run training sequence?") 36 | parser.add_argument("--generate", default=False, action="store_true", help="Run generation sequence?") 37 | parser.add_argument("--retraining", default=False, action="store_true", help="Run retraining sequence?") 38 | parser.add_argument("--resume", default=False, action="store_true", help="Will resume from the 'latest_model_checkpoint.pth' and wandb if applicable.") 39 | 40 | parser.add_argument("--save_features", default=False, action="store_true", help="Collect feature representations.") 41 | parser.add_argument("--save_features_phase", type=str, default="train", help="Dataset split of representations to collect.") 42 | 43 | parser.add_argument("--config", type=str, default=None, help="If you have a yaml file with appropriate config, provide the path here. Will override the 'experiment_maker'.") 44 | 45 | args = parser.parse_args() 46 | 47 | # CUDA devices used 48 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 49 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 50 | 51 | # custom imports 52 | from libs.utils.experiments_maker import experiment_maker 53 | from libs.data import dataloader 54 | from libs.utils.utils import * 55 | import libs.utils.globals as g 56 | 57 | # global configs 58 | g.wandb_log = True 59 | g.epoch_global = 0 60 | g.log_offline = True 61 | 62 | # Fixing random seed 63 | print(f"=======> Using seed: {args.seed} <========") 64 | random.seed(args.seed) 65 | torch.manual_seed(args.seed) 66 | torch.cuda.manual_seed(args.seed) 67 | torch.cuda.manual_seed_all(args.seed) 68 | np.random.seed(args.seed) 69 | torch.backends.cudnn.deterministic = True 70 | torch.backends.cudnn.benchmark = False 71 | g.seed = args.seed 72 | g.resume = args.resume 73 | 74 | # Roots of datasets (Change this according to your current directory) 75 | data_root = { 76 | "CIFAR100": "./datasets/CIFAR100", 77 | "mini-imagenet": "/home/rahul_intern/Imagenet/mini_imagenet", 78 | } 79 | 80 | # Either use an existing yaml file or use "experiment_maker" 81 | if args.config == None : 82 | config = experiment_maker(args.experiment, args.dataset, args.imbalance, data_root, args.seed, args.type_of_val, args.cv1, args.cv2, args.cv3, args.cv4, args.cv5, args.cv6, args.cv7, args.cv8, args.cv9) 83 | else: 84 | config = yaml.load(args.config) 85 | 86 | # wandb inits/resume 87 | if g.wandb_log: 88 | import wandb 89 | config_dictionary = config 90 | if args.resume: 91 | id = torch.load(config["training_opt"]["log_dir"]+"/latest_model_checkpoint.pth")['wandb_id'] 92 | print(f"\nResuming wandb id: {id}!\n") 93 | else: 94 | id = wandb.util.generate_id() 95 | print(f"\nStarting wandb id: {id}!\n") 96 | wandb.init( 97 | project="long-tail", 98 | entity="long-tail", 99 | reinit=True, 100 | name=f"{config['training_opt']['stage']}", 101 | allow_val_change=True, 102 | save_code=True, 103 | config=config_dictionary, 104 | tags=config["wandb_tags"], 105 | id=id, 106 | resume="allow", 107 | ) 108 | wandb.config.update(args, allow_val_change=True) 109 | config["wandb_id"] = id 110 | else: 111 | config["wandb_id"] = None 112 | 113 | # Create necessary directories for logging 114 | if (args.train or args.generate) and not(args.resume): 115 | if not os.path.isdir(config["training_opt"]["log_dir"]): 116 | os.makedirs(config["training_opt"]["log_dir"]) 117 | else: 118 | raise Exception("Directory already exists!!") 119 | g.log_dir = config["training_opt"]["log_dir"] 120 | if g.log_offline: 121 | if not os.path.isdir(f"{g.log_dir}/metrics"): 122 | os.makedirs(f"{g.log_dir}/metrics") 123 | 124 | # Save the config as yaml file 125 | ff = open(f'{config["training_opt"]["log_dir"]}/config.yaml', 'w+') 126 | yaml.dump(config, ff, default_flow_style=False, allow_unicode=True) 127 | 128 | # Splits 129 | splits = ["train", "val", "test"] 130 | 131 | # Generate dataloader for all the splits 132 | data = { 133 | x: dataloader.load_data( 134 | data_root=data_root[config["training_opt"]["dataset"].rstrip("_LT")], 135 | dataset=config["training_opt"]["dataset"], 136 | phase=x, 137 | batch_size=config["training_opt"]["batch_size"], #Use 512 above to shave 1 min off the FreeLunch accumulate 138 | sampler_dic=get_sampler_dict(config["training_opt"]["sampler"]), 139 | num_workers=config["training_opt"]["num_workers"], 140 | top_k_class=config["training_opt"]["top_k"] if "top_k" in config["training_opt"] else None, 141 | cifar_imb_ratio=config["training_opt"]["cifar_imb_ratio"] if "cifar_imb_ratio" in config["training_opt"] else None, 142 | imb_type=config["training_opt"]["imb_type"] if "imb_type" in config["training_opt"] else None, 143 | class_order=config["training_opt"]["class_order"] if "class_order" in config["training_opt"] else None, 144 | balanced=config["training_opt"]["balanced"] if "balanced" in config["training_opt"] else None, 145 | special_aug=config["training_opt"]["special_aug"] if "special_aug" in config["training_opt"] else False, 146 | seed=args.seed, 147 | jitter=config["training_opt"]["jitter"] if "jitter" in config["training_opt"] else True, 148 | type_of_val=args.type_of_val 149 | ) 150 | for x in splits 151 | } 152 | 153 | # Number of samples in each class 154 | config["training_opt"]["data_count"] = data["train"].dataset.img_num_list 155 | print(config["training_opt"]["data_count"]) 156 | 157 | # import appropriate core 158 | training_model = source_import(config["core"]).get_core(config, data) 159 | 160 | if args.train: 161 | # training sequence 162 | print("\nInitiating training sequence!") 163 | if args.resume: 164 | training_model.resume_run(config["training_opt"]["log_dir"]+"/latest_model_checkpoint.pth") 165 | training_model.train() 166 | 167 | if not(args.generate) and args.save_features: 168 | # Accumulate and save features alone 169 | training_model.reset_model(torch.load(config["training_opt"]["log_dir"]+"/final_model_checkpoint.pth")['state_dict_best']) 170 | print("Model reset to best model!") 171 | training_model.accumulate(phase=args.save_features_phase, save=args.save_features) 172 | 173 | if args.generate: 174 | # Point generation sequence 175 | print("\nInitiating point generation sequence!") 176 | print("Model reset to best model!") 177 | training_model.accumulate(phase="train", save=args.save_features) 178 | training_model.generate_points() 179 | 180 | if args.retraining: 181 | # Retraining sequence 182 | if not(args.generate) and not(args.save_features): 183 | training_model.reset_model(torch.load(config["training_opt"]["log_dir"]+"/final_model_checkpoint.pth")['state_dict_best']) 184 | print("Model reset to best model!") 185 | training_model.accumulate(phase="train", save=args.save_features) 186 | print("\nInitiating retraining sequence!") 187 | training_model.retrain() 188 | 189 | print("=" * 25, " ALL COMPLETED ", "=" * 25) 190 | -------------------------------------------------------------------------------- /readme_assets/long_tail-buzz_ss.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rahulvigneswaran/TailCalibX/b836fd93925efb48f610c08bd56bb0c5a3fca9b0/readme_assets/long_tail-buzz_ss.png -------------------------------------------------------------------------------- /run_TailCalibX_CIFAR100-LT.sh: -------------------------------------------------------------------------------- 1 | # CIFAR100_LT 2 | 3 | # Ours 4 | #---------------------------------------------------------- 5 | # -[x] CosineCE + TailCalib (1.2) # Pre-requisite: Experiment "0.2" - Check "run_all_CIFAR100-LT.sh" 6 | # -[x] CosineCE + TailCalibX (2.2) # Pre-requisite: Experiment "0.2" - Check "run_all_CIFAR100-LT.sh" 7 | # -[x] CBD + TailCalibX (2.4) # Pre-requisite: Experiment "0.4" - Check "run_all_CIFAR100-LT.sh" 8 | #---------------------------------------------------------- 9 | 10 | #----Cifar100_LT 11 | actual_dataset=1 12 | 13 | #----CosineCE + TailCalib 14 | experiment_no=1.2 15 | for imb in 1 16 | do 17 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 18 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 19 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 20 | done 21 | wait 22 | for imb in 2 23 | do 24 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 25 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 26 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 27 | done 28 | wait 29 | for imb in 3 30 | do 31 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 32 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 33 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 34 | done 35 | 36 | wait 37 | 38 | #----CosineCE + TailCalibX 39 | experiment_no=2.2 40 | for imb in 1 41 | do 42 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 43 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 44 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 45 | done 46 | wait 47 | for imb in 2 48 | do 49 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 50 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 51 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 52 | done 53 | wait 54 | for imb in 3 55 | do 56 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 57 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 58 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 59 | done 60 | 61 | wait 62 | 63 | #----CBD + TailCalibX 64 | experiment_no=2.4 65 | for imb in 1 66 | do 67 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 68 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 69 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 70 | done 71 | wait 72 | for imb in 2 73 | do 74 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 75 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 76 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 77 | done 78 | wait 79 | for imb in 3 80 | do 81 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 82 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 83 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 84 | done 85 | 86 | #--------------------------------------------------- Paper experiments end here 87 | -------------------------------------------------------------------------------- /run_TailCalibX_mini-imagenet-LT.sh: -------------------------------------------------------------------------------- 1 | # mini-ImageNet-LT 2 | 3 | # Ours 4 | #---------------------------------------------------------- 5 | # -[x] CosineCE + TailCalib (1.2) # Pre-requisite: Experiment "0.2" - Check "run_all_mini-ImageNet-LT.sh" 6 | # -[x] CosineCE + TailCalibX (2.2) # Pre-requisite: Experiment "0.2" - Check "run_all_mini-ImageNet-LT.sh" 7 | #---------------------------------------------------------- 8 | 9 | #----mini-ImageNet_LT 10 | actual_dataset=2 11 | seeds=1 12 | 13 | #----CosineCE + TailCalib 14 | experiment_no=1.2 15 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining 16 | 17 | wait 18 | 19 | #----CosineCE + TailCalibX 20 | experiment_no=2.2 21 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train 22 | 23 | wait 24 | #--------------------------------------------------- Paper experiments end here 25 | -------------------------------------------------------------------------------- /run_all_CIFAR100-LT.sh: -------------------------------------------------------------------------------- 1 | # CIFAR100_LT 2 | 3 | # Table 4 4 | #---------------------------------------------------------- 5 | # Baseline 6 | # -[x] CE (0.1) # Pre-requisite: None 7 | # -[x] CosineCE (0.2) # Pre-requisite: None 8 | #---------------------------------------------------------- 9 | # Decouple 10 | # -[x] cRT (0.3) # Pre-requisite: Experiment "0.1" 11 | #---------------------------------------------------------- 12 | # Distillation 13 | # -[x] CBD (0.4) # Pre-requisite: Experiment "0.2" but with seeds 10, 20, 30 14 | # For repo specific to CBD paper with much more detailed instructions, check https://github.com/rahulvigneswaran/Class-Balanced-Distillation-for-Long-Tailed-Visual-Recognition.pytorch 15 | #---------------------------------------------------------- 16 | # Generation 17 | # -[x] MODALS (0.5) # Pre-requisite: Experiment "0.1" 18 | #---------------------------------------------------------- 19 | # Ours 20 | # -[x] CosineCE + TailCalib (1.2) # Pre-requisite: Experiment "0.2" 21 | # -[x] CosineCE + TailCalibX (2.2) # Pre-requisite: Experiment "0.2" 22 | # -[x] CBD + TailCalibX (2.4) # Pre-requisite: Experiment "0.4" 23 | #---------------------------------------------------------- 24 | 25 | #----Cifar100_LT 26 | actual_dataset=1 27 | 28 | #----CE 29 | experiment_no=0.1 30 | for seeds in 1 2 3 31 | do 32 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=1 --seed=$seeds --train & 33 | python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=2 --seed=$seeds --train & 34 | python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=3 --seed=$seeds --train & 35 | done 36 | 37 | wait 38 | 39 | # ----CosineCE 40 | experiment_no=0.2 41 | for seeds in 1 2 3 10 20 30 42 | do 43 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=1 --seed=$seeds --train & 44 | python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=2 --seed=$seeds --train & 45 | python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=3 --seed=$seeds --train & 46 | done 47 | 48 | wait 49 | 50 | #----cRT 51 | experiment_no=0.3 52 | for seeds in 1 2 3 53 | do 54 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=1 --seed=$seeds --train & 55 | python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=2 --seed=$seeds --train & 56 | python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=3 --seed=$seeds --train & 57 | done 58 | 59 | wait 60 | 61 | #----CBD 62 | experiment_no=0.4 63 | for seeds in 1 2 3 64 | do 65 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=1 --seed=$seeds --cv1=0.8 --cv2=100 --train & 66 | python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=2 --seed=$seeds --cv1=0.8 --cv2=200 --train & 67 | python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=3 --seed=$seeds --cv1=0.8 --cv2=100 --train & 68 | done 69 | 70 | wait 71 | 72 | #----MODALS 73 | experiment_no=0.5 74 | for seeds in 1 2 3 75 | do 76 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=1 --seed=$seeds --cv1=0.01 --generate --retraining & 77 | python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=2 --seed=$seeds --cv1=0.01 --generate --retraining & 78 | python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=3 --seed=$seeds --cv1=0.01 --generate --retraining & 79 | done 80 | 81 | wait 82 | 83 | #----CosineCE + TailCalib 84 | experiment_no=1.2 85 | for imb in 1 86 | do 87 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 88 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 89 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining & 90 | done 91 | wait 92 | for imb in 2 93 | do 94 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 95 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 96 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --generate --retraining & 97 | done 98 | wait 99 | for imb in 3 100 | do 101 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 102 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 103 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --generate --retraining & 104 | done 105 | 106 | wait 107 | 108 | #----CosineCE + TailCalibX 109 | experiment_no=2.2 110 | for imb in 1 111 | do 112 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 113 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 114 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 115 | done 116 | wait 117 | for imb in 2 118 | do 119 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 120 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 121 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 122 | done 123 | wait 124 | for imb in 3 125 | do 126 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 127 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 128 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 129 | done 130 | 131 | wait 132 | 133 | #----CBD + TailCalibX 134 | experiment_no=2.4 135 | for imb in 1 136 | do 137 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 138 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 139 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train & 140 | done 141 | wait 142 | for imb in 2 143 | do 144 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 145 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 146 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.0 --cv4=0.2 --cv5=3 --train & 147 | done 148 | wait 149 | for imb in 3 150 | do 151 | taskset --cpu-list 10-19 python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --imbalance=$imb --seed=1 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 152 | taskset --cpu-list 20-29 python main.py --experiment=$experiment_no --gpu="2" --dataset=$actual_dataset --imbalance=$imb --seed=2 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 153 | taskset --cpu-list 30-39 python main.py --experiment=$experiment_no --gpu="3" --dataset=$actual_dataset --imbalance=$imb --seed=3 --cv1=0.9 --cv2=0.01 --cv3=0.9 --cv4=0.0 --cv5=2 --train & 154 | done 155 | 156 | #--------------------------------------------------- Paper experiments end here 157 | -------------------------------------------------------------------------------- /run_all_mini-ImageNet-LT.sh: -------------------------------------------------------------------------------- 1 | # mini-ImageNet-LT 2 | 3 | # Table 4 4 | #---------------------------------------------------------- 5 | # Baseline 6 | # -[x] CE (0.1) # Pre-requisite: None 7 | # -[x] CosineCE (0.2) # Pre-requisite: None 8 | #---------------------------------------------------------- 9 | # Decouple 10 | # -[x] cRT (0.3) # Pre-requisite: Experiment "0.1" 11 | #---------------------------------------------------------- 12 | # Distillation 13 | # -[x] CBD (0.4) # Pre-requisite: Experiment "0.2" but with seeds 10 14 | # For repo specific to CBD paper with much more detailed instructions, check https://github.com/rahulvigneswaran/Class-Balanced-Distillation-for-Long-Tailed-Visual-Recognition.pytorch 15 | #---------------------------------------------------------- 16 | # Generation 17 | # -[x] MODALS (0.5) # Pre-requisite: Experiment "0.1" 18 | #---------------------------------------------------------- 19 | # Ours 20 | # -[x] CosineCE + TailCalib (1.2) # Pre-requisite: Experiment "0.2" 21 | # -[x] CosineCE + TailCalibX (2.2) # Pre-requisite: Experiment "0.2" 22 | #---------------------------------------------------------- 23 | 24 | 25 | #----mini-ImageNet_LT 26 | actual_dataset=2 27 | seeds=1 28 | 29 | #----CE 30 | experiment_no=0.1 31 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=$seeds --train 32 | 33 | wait 34 | 35 | # ----CosineCE 36 | experiment_no=0.2 37 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=$seeds --train 38 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=10 --train #For CBD 39 | 40 | wait 41 | 42 | #----cRT 43 | experiment_no=0.3 44 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=$seeds --train 45 | 46 | wait 47 | 48 | #----CBD 49 | experiment_no=0.4 50 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=$seeds --cv1=0.4 --cv2=100 --train 51 | 52 | wait 53 | 54 | #----MODALS 55 | experiment_no=0.5 56 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --seed=$seeds --generate --cv1=0.01 --train 57 | 58 | wait 59 | 60 | #----CosineCE + TailCalib 61 | experiment_no=1.2 62 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --generate --retraining 63 | 64 | wait 65 | 66 | #----CosineCE + TailCalibX 67 | experiment_no=2.2 68 | python main.py --experiment=$experiment_no --gpu="1" --dataset=$actual_dataset --cv1=1.0 --cv2=0.01 --cv3=0.7 --cv4=0.0 --cv5=3 --train 69 | 70 | wait 71 | #--------------------------------------------------- Paper experiments end here 72 | -------------------------------------------------------------------------------- /tailcalib_pip/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Rahul Vigneswaran 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /tailcalib_pip/README.md: -------------------------------------------------------------------------------- 1 | # tailcalib 2 | 3 | __tailcalib__ is a Python library for balancing a __long-tailed / imbalanced__ dataset by generating synthetic datapoints which will inturn increase the class-wise and overall test accuracy on the original dataset. 4 | 5 | 6 | This package is based on the paper [Feature Generation for Long-tail Classification](https://github.com/rahulvigneswaran/TailCalibX) by [Rahul Vigneswaran](https://rahulvigneswaran.github.io/), [Marc T. Law](http://www.cs.toronto.edu/~law/), [Vineeth N. Balasubramanian](https://lab1055.github.io/), [Makarand Tapaswi](https://makarandtapaswi.github.io/). 7 | 8 | For much more detailed experiments, code and instructions, check [rahulvigneswaran/TailCalibX](https://github.com/rahulvigneswaran/TailCalibX) [![Star on GitHub](https://img.shields.io/github/stars/rahulvigneswaran/TailCalibX.svg?style=social)](https://github.com/rahulvigneswaran/TailCalibX/stargazers) 9 | . 10 | ## 💻 Installation 11 | 12 | Use the package manager [pip](https://pip.pypa.io/en/stable/) to install __tailcalib__. 13 | 14 | ```bash 15 | pip install tailcalib 16 | ``` 17 | 18 | ## 👨‍💻 Basic Usage 19 | 20 | ```python 21 | # Import 22 | from tailcalib import tailcalib 23 | 24 | # Initialize 25 | a = tailcalib(base_engine="numpy") # Options: "numpy", "pytorch" 26 | 27 | # Imbalanced random fake data 28 | import numpy as np 29 | X = np.random.rand(200,100) 30 | y = np.random.randint(0,10, (200,)) 31 | 32 | # Balancing the data using "tailcalib" 33 | feat, lab, gen = a.generate(X=X, y=y) 34 | 35 | # Output comparison 36 | print(f"Before: {np.unique(y, return_counts=True)}") 37 | print(f"After: {np.unique(lab, return_counts=True)}") 38 | ``` 39 | 40 | ## 🧪 Advanced Usage 41 | ### 🧩 Sample code 42 | 43 | ```python 44 | # Import 45 | from tailcalib import tailcalib 46 | 47 | # Initialize 48 | a = tailcalib(base_engine="numpy") # Options: "numpy", "pytorch" 49 | 50 | # Imbalanced random fake data 51 | import numpy as np 52 | # Train data 53 | X_train = np.random.rand(200,100) 54 | y_train = np.random.randint(0,10, (200,)) 55 | # Test data 56 | X_test = np.random.rand(20,100) 57 | y_test = np.random.randint(0,10, (20,)) 58 | 59 | # Balancing the data using "tailcalib". 60 | # Try to play with the other hyperparameters to get a better generated datapoint. 61 | feat, lab, gen = a.generate(X=X_train, y=y_train, tukey_value=1.0, alpha=0.0, topk=1, extra_points=0, shuffle=True) 62 | 63 | # Always remember to convert the val/test data before doing validation/testing. 64 | X_test, y_test = a.convert_others(X=X_test, y=y_test) 65 | 66 | # Output comparison 67 | print(f"Before: {np.unique(y_train, return_counts=True)}") 68 | print(f"After: {np.unique(lab, return_counts=True)}") 69 | ``` 70 | 71 | ### ⚙ Arguments 72 | - `X` : Features 73 | - `y` : Corresponding labels 74 | - `tukey_value` : Value to convert any distrubution of data into a normal distribution. Defaults to 1.0. 75 | - `alpha` : Decides how spread out the generated data is. Defaults to 0.0. 76 | - `topk` : Decides how many nearby classes should be taken into consideration for the mean and std of the newly generated data. Defaults to 1. 77 | - `extra_points` : By default the number of datapoints to be generated is decided based on the class with the maximum datapoints. This variable decides how many more extra datapoints should be generated on top of that. Defaults to 0. 78 | - `shuffle` : Shuffles the generated and original datapoints together. Defaults to True. 79 | 80 | ### 📤 Returns: 81 | - `feat_all` : Tukey transformed train data + generated datapoints 82 | - `labs_all` : Corresponding labels to feat_all 83 | - `generated_points` : Dict that consists of just the generated points with class label as keys. 84 | 85 | 86 | 87 | ## 🪀 Results on a Toy Dataset [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1Yj2qymSm3NgCBqvKn5r_cOiEFl9wGp3J?usp=sharing) 88 | 89 | The higher the `Imb ratio`, the more imbalanced the dataset is. 90 | `Imb ratio = maximum_sample_count/minimum_sample_count`. 91 | 92 | Check [this notebook](https://colab.research.google.com/drive/1Yj2qymSm3NgCBqvKn5r_cOiEFl9wGp3J?usp=sharing) to play with the toy example from which the plot below was generated. 93 | ![](toy_example_output.svg) 94 | 95 | ## 📃 Citation 96 | If you use this package in any of your work, cite as, 97 | ``` 98 | @inproceedings{rahul2021tailcalibX, 99 | title = {{Feature Generation for Long-tail Classification}}, 100 | author = {Rahul Vigneswaran, Marc T. Law, Vineeth N. Balasubramanian, Makarand Tapaswi}, 101 | booktitle = {ICVGIP}, 102 | year = {2021} 103 | } 104 | ``` 105 | 106 | ## 👁 Contributing 107 | Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change. 108 | 109 | ## ❤ About me 110 | [Rahul Vigneswaran](https://rahulvigneswaran.github.io/) 111 | 112 | ## ✨ Extras 113 | [🐝 Long-tail buzz](https://rahulvigneswaran.github.io/longtail-buzz/) : If you are interested in deep learning research which involves __long-tailed / imbalanced__ dataset, take a look at [Long-tail buzz](https://rahulvigneswaran.github.io/longtail-buzz/) to learn about the recent trending papers in this field. 114 | 115 | 116 | 120 | 121 | ## 📝 License 122 | [MIT](LICENSE) 123 | -------------------------------------------------------------------------------- /tailcalib_pip/setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r", encoding="utf8") as fh: 4 | long_description = fh.read() 5 | 6 | setuptools.setup( 7 | name="tailcalib", 8 | version="0.0.1", 9 | author="Rahul Vigneswaran", 10 | description="tailcalib is a Python library for balancing a long-tailed / imbalanced dataset by generating synthetic datapoints which will inturn increase the class-wise and overall test accuracy on the original dataset.", 11 | long_description=long_description, 12 | long_description_content_type="text/markdown", 13 | packages=setuptools.find_packages(), 14 | classifiers=[ 15 | "Programming Language :: Python :: 3", 16 | "License :: OSI Approved :: MIT License", 17 | "Operating System :: OS Independent", 18 | ], 19 | python_requires='>=3.6', 20 | py_modules=["tailcalib"], 21 | package_dir={'':'src'}, 22 | install_requires=["numpy", "scikit-learn"] 23 | ) -------------------------------------------------------------------------------- /tailcalib_pip/src/tailcalib.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | from dataclasses import dataclass 5 | from tqdm import tqdm 6 | import numpy as np 7 | 8 | @dataclass 9 | class tailcalib: 10 | base_engine: str = "numpy" # Options: NumPy, PyTorch 11 | 12 | def generate(self, X, y, tukey_value=0.9, alpha=0.0, topk=1, extra_points=0, shuffle=True): 13 | """Generate new datapoints 14 | 15 | Args: 16 | X : Features 17 | y : Corresponding labels 18 | tukey_value [Hyperparameter]: Value to convert any distrubution of data into a normal distribution. Defaults to 1.0. 19 | alpha [Hyperparameter]: Decides how spread out the generated data is. Defaults to 0.0. 20 | topk [Hyperparameter]: Decides how many nearby classes should be taken into consideration for the mean and std of the newly generated data. Defaults to 1. 21 | extra_points [Hyperparameter]: By default the number of datapoints to be generated is decided based on the class with the maximum datapoints. This variable decides how many more extra datapoints should be generated on top of that. Defaults to 0. 22 | shuffle (bool, optional): Shuffles the generated and original datapoints together. Defaults to True. 23 | 24 | Returns: 25 | feat_all: Tukey transformed train data + generated datapoints 26 | labs_all: Corresponding labels to feat_all 27 | generated_points: Dict that consists of just the generated points with class label as keys. 28 | """ 29 | 30 | self.X = X 31 | self.y = y 32 | self.tukey_value = tukey_value 33 | self.alpha = alpha 34 | self.topk = topk 35 | self.extra_points = extra_points 36 | self.shuffle = shuffle 37 | 38 | self.sanity_checks() 39 | 40 | if self.base_engine == "numpy": 41 | import numpy as np 42 | return self.generate_usingNumPy() 43 | elif self.base_engine == "pytorch": 44 | import torch 45 | return self.generate_usingPyTorch() 46 | else: 47 | raise Exception(f"Invalid base_engine choice - '{self.base_engine}' | Choose from: 'numpy', 'pytorch'.") 48 | 49 | def sanity_checks(self,): 50 | """Checks whether the type of X, y inputs matches with the chosen base_engine 51 | """ 52 | if self.base_engine == "numpy": 53 | import numpy as np 54 | assert isinstance(self.X, np.ndarray),"Base Engine is set to NumPy, so the X must be a NumPy instance!" 55 | assert isinstance(self.y, np.ndarray),"Base Engine is set to NumPy, so the y must be a NumPy instance!" 56 | else: 57 | import torch 58 | assert torch.is_tensor(self.X),"Base Engine is set to PyTorch, so the X must be a PyTorch instance!" 59 | assert torch.is_tensor(self.y),"Base Engine is set to PyTorch, so the y must be a PyTorch instance!" 60 | 61 | def generate_usingNumPy(self,): 62 | """Generate new datapoints using Numpy 63 | """ 64 | import scipy.spatial as sp 65 | 66 | feat = {} 67 | labs = {} 68 | y_unique, y_count = np.unique(self.y, return_counts=True) 69 | assert len(y_unique) >= self.topk, "The 'topk' is greater than the number of uniquely available classes. Try a lesser value." 70 | for i in y_unique: 71 | feat[i] = self.X[self.y == i] 72 | labs[i] = np.full((feat[i].shape[0],), i) 73 | 74 | # Class statistics 75 | base_means = [] 76 | base_covs = [] 77 | 78 | for i in feat.keys(): 79 | base_means.append(feat[i].mean(axis=0)) 80 | base_covs.append(np.expand_dims(self.get_cov(feat[i]),axis=0)) 81 | 82 | base_means = np.vstack(base_means) 83 | base_covs = np.vstack(base_covs) 84 | 85 | # Tukey's transform 86 | for i in feat.keys(): 87 | feat[i] = self.tukey_transform(feat[i], self.tukey_value) 88 | 89 | # Distribution calibration and feature sampling 90 | sample_from_each = self.get_sample_count(y_count,feat.keys(), self.extra_points) 91 | 92 | generated_points = {} 93 | 94 | for i in tqdm(feat.keys()): 95 | if np.sum(sample_from_each[i]) == 0 and self.extra_points == 0 : 96 | continue 97 | 98 | generated_points[i] = [] 99 | for k, x_ij in zip(sample_from_each[i], feat[i]): 100 | if k == 0: 101 | continue 102 | # Getting the top k nearest classes based on l2 distance 103 | distances = sp.distance.cdist(base_means, np.expand_dims(x_ij, axis=0)).squeeze() 104 | topk_idx = np.argsort(-distances)[::-1][:self.topk] 105 | 106 | # Calibrating mean and covariance 107 | calibrated_mean, calibrated_cov = self.calibrate_distribution(base_means[topk_idx], base_covs[topk_idx], self.topk, x_ij, self.alpha) 108 | 109 | # Trick to avoid cholesky decomposition from failing. Look at https://juanitorduz.github.io/multivariate_normal/ 110 | EPS = 1e-4 111 | calibrated_cov += (np.eye(calibrated_cov.shape[0])*EPS) 112 | 113 | gen = np.random.multivariate_normal(calibrated_mean, calibrated_cov,(int(k),)) 114 | 115 | generated_points[i].append(gen) 116 | 117 | generated_points[i] = np.vstack(generated_points[i]) 118 | 119 | print("Point Generation Completed!") 120 | print("Don't forget to use '.convert_others()' to apply tukey transformation on validation/test data before validation/testing. Use the same 'tukey_value' as the train data.\n") 121 | feat_all = [] 122 | labs_all = [] 123 | 124 | for i in labs.keys(): 125 | feat_all.append(feat[i]) 126 | labs_all.append(labs[i]) 127 | 128 | feat_all = np.vstack(feat_all) 129 | labs_all = np.hstack(labs_all) 130 | 131 | for i in generated_points.keys(): 132 | feat_all = np.concatenate((feat_all, generated_points[i])) 133 | labs_all = np.concatenate((labs_all, np.full((generated_points[i].shape[0],), int(i)))) 134 | 135 | feat_all, labs_all = self.shuffle_all(feat_all, labs_all) 136 | return feat_all, labs_all, generated_points 137 | 138 | def generate_usingPyTorch(self,): 139 | """Generate new datapoints using PyTorch 140 | """ 141 | import torch 142 | device = self.X.device 143 | 144 | feat = {} 145 | labs = {} 146 | y_unique, y_count = np.unique(self.y.cpu().numpy(), return_counts=True) 147 | assert len(y_unique) >= self.topk, "The 'topk' is greater than the number of uniquely available classes. Try a lesser value." 148 | for i in y_unique: 149 | feat[i] = self.X[self.y == i] 150 | labs[i] = torch.full((feat[i].size()[0],), i).to(device) 151 | 152 | # Class statistics 153 | base_means = [] 154 | base_covs = [] 155 | 156 | for i in feat.keys(): 157 | base_means.append(feat[i].mean(dim=0)) 158 | base_covs.append(self.get_cov(feat[i]).unsqueeze(dim=0)) 159 | 160 | base_means = torch.vstack(base_means) 161 | base_covs = torch.vstack(base_covs) 162 | 163 | # Tukey's transform 164 | for i in feat.keys(): 165 | feat[i] = self.tukey_transform(feat[i], self.tukey_value) 166 | 167 | # Distribution calibration and feature sampling 168 | sample_from_each = self.get_sample_count(y_count,feat.keys(), self.extra_points) 169 | 170 | generated_points = {} 171 | 172 | for i in tqdm(feat.keys()): 173 | if np.sum(sample_from_each[i]) == 0 and self.extra_points == 0 : 174 | continue 175 | 176 | generated_points[i] = [] 177 | for k, x_ij in zip(sample_from_each[i], feat[i]): 178 | if k == 0: 179 | continue 180 | # Getting the top k nearest classes based on l2 distance 181 | distances = torch.cdist(base_means, x_ij.unsqueeze(0)).squeeze() 182 | topk_idx = torch.topk(-distances, k=self.topk)[1][:self.topk] 183 | 184 | # Calibrating mean and covariance 185 | calibrated_mean, calibrated_cov = self.calibrate_distribution(base_means[topk_idx], base_covs[topk_idx], self.topk, x_ij, self.alpha) 186 | 187 | # Trick to avoid cholesky decomposition from failing. Look at https://juanitorduz.github.io/multivariate_normal/ 188 | EPS = 1e-4 189 | calibrated_cov += (torch.eye(calibrated_cov.shape[0])*EPS).to(device) 190 | 191 | new_dist = torch.distributions.multivariate_normal.MultivariateNormal(calibrated_mean, calibrated_cov) 192 | 193 | gen = new_dist.sample((int(k),)) 194 | 195 | generated_points[i].append(gen) 196 | 197 | generated_points[i] = torch.vstack(generated_points[i]) 198 | torch.cuda.empty_cache() 199 | 200 | print("Point Generation Completed!") 201 | print("Don't forget to use '.convert_others()' to apply tukey transformation on validation/test data before validation/testing. Use the same 'tukey_value' as the train data.\n") 202 | feat_all = [] 203 | labs_all = [] 204 | 205 | for i in labs.keys(): 206 | feat_all.append(feat[i]) 207 | labs_all.append(labs[i]) 208 | 209 | feat_all = torch.vstack(feat_all) 210 | labs_all = torch.hstack(labs_all).to(device) 211 | 212 | for i in generated_points.keys(): 213 | feat_all = torch.cat((feat_all, generated_points[i].to(device))) 214 | labs_all = torch.cat((labs_all, torch.full((generated_points[i].size()[0],), int(i)).to(device))) 215 | 216 | feat_all, labs_all = self.shuffle_all(feat_all, labs_all) 217 | return feat_all, labs_all, generated_points 218 | 219 | def convert_others(self, X, tukey_value=1.0): 220 | """Use for applying tukey transformation on validation/test data before validation/testing! 221 | 222 | Args: 223 | X : Features 224 | tukey_value [Hyperparameter]: Value to convert any distrubution of data into a normal distribution. Defaults to 1.0. 225 | 226 | Returns: 227 | Tukey transformed data 228 | """ 229 | # Sanity check 230 | if self.base_engine == "numpy": 231 | import numpy as np 232 | assert isinstance(X, np.ndarray),"Base Engine is set to NumPy, so the X must be a NumPy instance!" 233 | else: 234 | import torch 235 | assert torch.is_tensor(X),"Base Engine is set to PyTorch, so the X must be a PyTorch instance!" 236 | 237 | return self.tukey_transform(X, tukey_value) 238 | 239 | def shuffle_all(self, x, y): 240 | """Force shuffle data 241 | 242 | Args: 243 | x (float Tensor): Datapoints 244 | y (int): Labels 245 | 246 | Returns: 247 | floatTensor, int: Return shuffled datapoints and corresponding labels 248 | """ 249 | if self.base_engine == "numpy": 250 | index = np.random.permutation(x.shape[0]) 251 | else: 252 | index = torch.randperm(x.size(0)) 253 | x = x[index] 254 | y = y[index] 255 | return x, y 256 | 257 | def get_cov(self, X): 258 | """Calculate the covariance matrix for X 259 | 260 | Args: 261 | X (torch.tensor): Features 262 | 263 | Returns: 264 | [torch.tensor]: Covariance matrix of X 265 | """ 266 | n = X.shape[0] 267 | if self.base_engine == "numpy": 268 | mu = X.mean(axis=0) 269 | else: 270 | mu = X.mean(dim=0) 271 | X = (X - mu) 272 | 273 | return 1/(n-1) * (X.transpose(1, 0) @ X) # X^TX -> feat_size x num_of_samples @ num_of_samples x feat_size -> feat_size x feat_size 274 | 275 | def get_sample_count(self, count, keys, extra_points): 276 | """Decides how many samples must be generated based on each existing train datapoints. 277 | 278 | Args: 279 | count (list): Number of samples in each class 280 | keys (dict.keys): Class keys 281 | extra_points (int): The number of samples to be generated is based on the class with maximum sample count. "extra_points" decide how much more samples must be generated on top of that. 282 | 283 | Returns: 284 | dict: dict consists that has the info as to how many samples must be generated based on each existing train datapoints. 285 | """ 286 | sample_count_dict = {} 287 | for i in keys: 288 | current = count[i] 289 | head = max(count) 290 | # head is the sample count that we must match after the generation. This can be offset by "config["pg"]["extra_points"]". In our experiments this is set to 0 as it worked better. 291 | num_sample = head - current + extra_points 292 | ratio = num_sample / current 293 | # Makes sure each datapoint is being used atleast once 294 | new_sample_from_each = [np.floor(ratio)] * current 295 | 296 | # Rest of the datapoints used for generation are decided randomly 297 | while True: 298 | if sum(new_sample_from_each) == num_sample: 299 | break 300 | idx = np.random.randint(0, current) 301 | new_sample_from_each[idx] += 1 302 | 303 | # Sanity checks 304 | assert sum(new_sample_from_each) == num_sample 305 | assert len(new_sample_from_each) == current 306 | 307 | sample_count_dict[i] = new_sample_from_each 308 | return sample_count_dict 309 | 310 | 311 | def calibrate_distribution(self, base_means, base_cov, k, x_ij, alpha=0.0): 312 | """Calibration of the distribution for generation. Check equation 7 and 8 from our paper - Feature Generation for Long-tail Classification. 313 | 314 | Args: 315 | base_means (torch.tensor): List of all the means that are used for calibration. 316 | base_cov (torch.tensor): List of all the covariance matrices used for calibraton. 317 | k (int): Number of classes used for calibration. 318 | x_ij (torch.tensor): Datapoint chosen to be used for generation. 319 | alpha (float, optional): Decides the spread of the generated samples. Defaults to 0.0. 320 | 321 | Returns: 322 | torch.tensor : Calibrated mean and covariance matrix 323 | """ 324 | if self.base_engine == "numpy": 325 | calibrated_mean = (base_means.sum(axis=0) + x_ij)/(k+1) 326 | calibrated_cov = base_cov.sum(axis=0)/k + alpha 327 | else: 328 | calibrated_mean = (base_means.sum(dim=0) + x_ij)/(k+1) 329 | calibrated_cov = base_cov.sum(dim=0)/k + alpha 330 | 331 | return calibrated_mean, calibrated_cov 332 | 333 | def tukey_transform(self, x, lam=0.2): 334 | """Transforms any distribution into normal-distribution like. 335 | 336 | Args: 337 | x (torch.tensor): Features 338 | lam (float, optional): Adjusts how close the transformed features will be to the origin. Defaults to 0.2. 339 | 340 | Returns: 341 | torch.tensor: Normal distribution like features. 342 | """ 343 | if lam == 0: 344 | EPS = 1e-6 345 | x = x + EPS 346 | return x.log() 347 | else : 348 | return x**lam 349 | 350 | def sample_test(self): 351 | """[For internal testing only] Uses a randomly generated sample set to check if both the engines work error free. 352 | """ 353 | if self.base_engine == "numpy": 354 | import numpy as np 355 | X = np.random.rand(200,100) 356 | y = np.random.randint(0,10, (200,)) 357 | feat, lab, gen = self.generate(X=X, y=y) 358 | print(f"Before: {np.unique(y, return_counts=True)}") 359 | print(f"After: {np.unique(lab, return_counts=True)}") 360 | else: 361 | import torch 362 | X = torch.rand((200,100)) 363 | y = torch.randint(0,10, (200,)) 364 | feat, lab, gen = self.generate(X=X, y=y) 365 | print(f"Before: {torch.unique(y, return_counts=True)}") 366 | print(f"After: {torch.unique(lab, return_counts=True)}") --------------------------------------------------------------------------------