{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 🌱🌿🌾Sorghum PyTorch DDP baseline🚀\n\nThis notbook is ***based on*** the [PyTorch DDP](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html).  \n\n>If you want to use **TPU** DDP, only a little code you need to change, check this [pytorch-xla](https://github.com/pytorch/xla). \n\nThanks to the small jpegs Sorghum images from https://www.kaggle.com/datasets/mithilsalunkhe/small-jpegs-fgvc\n\nIf you have any ***question*** about my baseline, please *feel free* to ***make a comment***. I will reply as soon as possible! If you like it, please **upvote**👏👏👏\n\n\n## Inference result:\n>##### The first fold in 5-Fold model, without TTA, get LB **Acc@1: 82.9%**\n>##### The 1,2 fold mix up(Avg), without TTA, get LB **Acc@1: 85.5%**\n>##### All train image(instead of k-fold) model, without TTA, get LB **Acc@1: 85.6%**\n>##### All train image(instead of k-fold) model, use flip and crop TTA, get LB **Acc@1: 86.3%**\n>##### The 1,2,3 fold in 5-Fold and all-enrolled model mixup, use flip and crop TTA, get LB **Acc@1: 87.7%**\n\n## main idea\n+ *pre-process*\nthe images(the Visualization part in my notebook will explain the reason)\n+ visualize and check the pre-process\n+ use the pretrained model(import timm)\n+ images augmentation\n+ train the model with DDP (and mix precise)\n+ [Inference](https://www.kaggle.com/code/leoooo333/lb-0-85-sorghum-higer-accuracy)\n\n## tricks\n#### Pre-process\n+ **CLAHE** : To pre-process images, use CLAHE in opencv to make image brighter and a higer constrast. [opencv official tutorial here](https://docs.opencv.org/4.x/d5/daf/tutorial_py_histogram_equalization.html)\n\n#### Model\n+ **Arcface Loss** : It helps a little in classification accuracy.[README](https://arxiv.org/pdf/1801.07698.pdf) to know more about Arcface Loss! \n\n+ **Multiple Dropout** : To train a general model, try it.[README](https://arxiv.org/pdf/1905.09788.pdf) to learn about multiple dropout!\n\n+ **Concat global pooling** : use both avgpool and maxpool and concat them in the globalpool layer.\n\n+ **pseudo label** : use efficientnet-noisy student pretrained model, which is small in size but high in accuracy.\n\n#### Training\n+ **mix precisison** : efficent and secure! Save the GPU memory and speed up your training.\n\n+ **Ranger** : better than AdamW, most of time\n\n+ **Exponential Warmup** : really help on the models with attention machnism. But you should be awared of **local optimum**, which could be a result of small lr. \n\n+ **Consine scheduler** : after warmup, descend the lr as cosine function do.\n\n+ **Big lr makes surprise** : Though we take transfer learning method, using big lr(like lr=4e-3) in proper time will speed up, expecially when you find the loss do not change for a long time.\n\n#### Inference\n+ **Inference Tutorial is now available! [click here](https://www.kaggle.com/code/leoooo333/sorghum-higer-accuracy/notebook)**\n\n## DDP\n**The most important thing**\n>To enable DDP, please do not forget to change this notebook to a .py file,and run it in command line.\n>You can also download my [.py version](https://drive.google.com/file/d/1McrTZxIxmvn8L72yzYhW1Sx36OgPB5YE/view?usp=sharing) of this baseline.","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:37:49.001353Z","iopub.execute_input":"2022-04-20T11:37:49.001656Z","iopub.status.idle":"2022-04-20T11:37:49.011915Z","shell.execute_reply.started":"2022-04-20T11:37:49.001624Z","shell.execute_reply":"2022-04-20T11:37:49.010557Z"}}},{"cell_type":"code","source":"!pip install seaborn","metadata":{"execution":{"iopub.status.busy":"2022-04-21T02:07:41.705839Z","iopub.execute_input":"2022-04-21T02:07:41.706245Z","iopub.status.idle":"2022-04-21T02:07:54.986645Z","shell.execute_reply.started":"2022-04-21T02:07:41.706138Z","shell.execute_reply":"2022-04-21T02:07:54.985225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install timm","metadata":{"execution":{"iopub.status.busy":"2022-04-20T10:57:37.523622Z","iopub.execute_input":"2022-04-20T10:57:37.523935Z","iopub.status.idle":"2022-04-20T10:57:51.377693Z","shell.execute_reply.started":"2022-04-20T10:57:37.523899Z","shell.execute_reply":"2022-04-20T10:57:51.376556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchvision import transforms, models\nimport pandas as pd\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom PIL import Image as Img\nfrom tqdm import tqdm\nimport os\nfrom sklearn.metrics import accuracy_score\nimport timm\nfrom tqdm import tqdm  \nfrom torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR\n\nimport torchvision\nimport matplotlib.pyplot as plt\nimport re\nfrom sklearn.model_selection import train_test_split, StratifiedKFold\nimport pytorch_lightning as pl\nimport seaborn as sns\nimport cv2 as cv\nimport numpy as np\nimport torch.nn.functional as F\n\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.tensorboard import SummaryWriter\n\nimport numpy as np\nimport math","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:27:32.644679Z","iopub.execute_input":"2022-04-20T11:27:32.645015Z","iopub.status.idle":"2022-04-20T11:27:34.598079Z","shell.execute_reply.started":"2022-04-20T11:27:32.644971Z","shell.execute_reply":"2022-04-20T11:27:34.597142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"MODEL_NAME = 'tf_efficientnet_b5_ns'\nLR = 1e-3\nLR_MIN = 1e-5\nBATCH_SIZE = 24\nIMAGE_SIZE = 900\nEPOCH = 20\nWARM_UP=5\nWEIGHT_DECAY = 1e-3\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nFOLD = 5\nLOSS = 'CrossEntropy'\nOPTIM = 'adamW'\nSCHEDULER = 'Cosine'\nUSE_AMP = True\nINIT = False\n\n'''Last two parameter depends on devices'''\nWORLD_SIZE = torch.cuda.device_count() # DistributedDataParallel\nNUM_WORKERS = 18\n\nroot_in = '../input/small-jpegs-fgvc' #Folder with input (image, lable)\nroot_out = './' #Folder with output (csv, pth) \nhave_index = False # If the breed label have been map to a index\nSEED = 42\nFOLD = 5\n\n'''ArcFace parameter'''\nNUM_CLASSES = 100\nEMBEDDING_SIZE = 1024\nS, M = 30.0, 0.5 # S:consine scale in arcloss. M:arg penalty\nEASY_MERGING, LS_EPS = False, 0.0","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:28:32.593422Z","iopub.execute_input":"2022-04-20T11:28:32.593761Z","iopub.status.idle":"2022-04-20T11:28:32.603284Z","shell.execute_reply.started":"2022-04-20T11:28:32.593725Z","shell.execute_reply":"2022-04-20T11:28:32.60256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils\n","metadata":{}},{"cell_type":"markdown","source":"### train and test utils","metadata":{}},{"cell_type":"code","source":"class Accumulator():\n    '''A counter util, which count the float value of the input'''\n    def __init__(self, nums):\n        self.metric = list(torch.zeros((nums,)).numpy())\n        \n    def __getitem__(self, index):\n        return self.metric[index]\n    \n    def add(self, *args):\n        for i, item in enumerate(args):\n            self.metric[i] += float(item)","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.837844Z","iopub.status.idle":"2022-04-20T11:21:32.838191Z","shell.execute_reply.started":"2022-04-20T11:21:32.83802Z","shell.execute_reply":"2022-04-20T11:21:32.838038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def accuracy(y_hat, y):\n    '''used to count the right type'''\n    y_hat = y_hat.exp().argmax(dim=1)\n    y_hat.reshape((-1))\n    y.reshape((-1))\n    return accuracy_score(y.cpu().numpy(), y_hat.cpu().numpy(), normalize=False)","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.839251Z","iopub.status.idle":"2022-04-20T11:21:32.839562Z","shell.execute_reply.started":"2022-04-20T11:21:32.83939Z","shell.execute_reply":"2022-04-20T11:21:32.839406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def evaluate_accuracy(net, data_iter, device=None):\n    '''Evalue the valid dataset'''\n    if isinstance(net, nn.Module):\n        net.eval()\n        if not device:\n            device = next(iter(net.parameters())).device\n    \n    metric = Accumulator(2)\n    with torch.no_grad():\n        for X, y in data_iter:\n            if isinstance(X, list):\n                X = [x.to(device) for x in X]\n            else:\n                X = X.to(device)\n            y.to(device)\n            with torch.cuda.amp.autocast(enabled=True):\n                metric.add(accuracy(net(X, y), y), y.numel())\n    return metric[0] / metric[1]","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.840766Z","iopub.status.idle":"2022-04-20T11:21:32.841114Z","shell.execute_reply.started":"2022-04-20T11:21:32.840914Z","shell.execute_reply":"2022-04-20T11:21:32.84093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict_test(net, test_iter, device=None):\n    '''Inference'''\n    if isinstance(net, nn.Module):\n        net.eval()\n        if not device:\n            device = next(iter(net.parameters())).device\n    y = []\n    net.to(device)\n    #softmax = nn.Softmax(dim=1)\n    with torch.no_grad():\n        for X in test_iter:\n            if isinstance(X, list):\n                X = [x.to(device) for x in X]\n            else:\n                X = X.to(device)\n            with torch.cuda.amp.autocast(enabled=True):\n                #y += softmax(net(X).cpu())\n                _, indice = net(X).exp().sort(dim=1,descending=True)\n                y += indice[:, 0:5].cpu()\n\n    return list(Y.numpy() for Y in y)","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.843927Z","iopub.status.idle":"2022-04-20T11:21:32.844418Z","shell.execute_reply.started":"2022-04-20T11:21:32.844229Z","shell.execute_reply":"2022-04-20T11:21:32.844253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    pl.utilities.seed.seed_everything(seed)\n    return seed","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.845356Z","iopub.status.idle":"2022-04-20T11:21:32.845671Z","shell.execute_reply.started":"2022-04-20T11:21:32.845504Z","shell.execute_reply":"2022-04-20T11:21:32.84552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model utils\n","metadata":{"tags":[]}},{"cell_type":"code","source":"def freeze_pretrained_layers(model):\n    '''Freeze all layers except the last layer(fc or classifier)'''\n    for param in model.parameters():\n            param.requires_grad = False\n    #nn.init.xavier_normal_(model.fc.weight)\n    #nn.init.zeros_(model.fc.bias)\n    model.classifier.weight.requires_grad = True\n    model.classifier.bias.requires_grad = True","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.846864Z","iopub.status.idle":"2022-04-20T11:21:32.847194Z","shell.execute_reply.started":"2022-04-20T11:21:32.847033Z","shell.execute_reply":"2022-04-20T11:21:32.84705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def debarcle_layers(model, num_debarcle=0, db_all=False):\n    '''Debarcle From the last [-1]layer to the [-num_debarcle] layers, \n    approximately(for there is Conv2d which has only weight parameter),\n    if db_all == True, debarcle all layers'''\n    num_debarcle *= 2\n    param_debarcle = param_name[-num_debarcle:]\n    if param_debarcle[0].split('.')[-1] == 'bias':\n        param_debarcle = param_name[-(num_debarcle + 1):]\n    if db_all:\n        for name, param in model.named_parameters():\n            param.requires_grad = True\n    else:\n        for name, param in model.named_parameters():\n            param.requires_grad = True if name in param_debarcle else False","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.848336Z","iopub.status.idle":"2022-04-20T11:21:32.848653Z","shell.execute_reply.started":"2022-04-20T11:21:32.848482Z","shell.execute_reply":"2022-04-20T11:21:32.848498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### DDP utils","metadata":{}},{"cell_type":"code","source":"def setup(rank, world_size):\n    os.environ['MASTER_ADDR'] = 'localhost'\n    os.environ['MASTER_PORT'] = '12355'\n\n    # initialize the process group\n    dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.850007Z","iopub.status.idle":"2022-04-20T11:21:32.850341Z","shell.execute_reply.started":"2022-04-20T11:21:32.850179Z","shell.execute_reply":"2022-04-20T11:21:32.850197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def cleanup():\n    dist.destroy_process_group()","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.85147Z","iopub.status.idle":"2022-04-20T11:21:32.851781Z","shell.execute_reply.started":"2022-04-20T11:21:32.851624Z","shell.execute_reply":"2022-04-20T11:21:32.85164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare(dataset, rank, world_size, batch_size=BATCH_SIZE, pin_memory=False, num_workers=0):\n    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False)\n    dataloader = DataLoader(dataset, batch_size=batch_size, pin_memory=pin_memory, num_workers=num_workers, drop_last=False, sampler=sampler)\n    return dataloader","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:21:32.852787Z","iopub.status.idle":"2022-04-20T11:21:32.853095Z","shell.execute_reply.started":"2022-04-20T11:21:32.852925Z","shell.execute_reply":"2022-04-20T11:21:32.852941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Build & Check label index","metadata":{}},{"cell_type":"code","source":"def data_pre_access(file, output):\n    '''transfer train label into index'''\n    labels = pd.read_csv(file, index_col='image')\n    labels_map = dict()\n    labels['label_index'] = torch.zeros((labels.shape[0])).type(torch.int32).numpy()\n    for i, label in enumerate(labels.cultivar.unique()):\n        labels_map[i] = label\n        labels.loc[labels.cultivar == label, 'label_index'] = i\n    labels.to_csv(output)\n    \n    return labels_map","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:22:53.687041Z","iopub.execute_input":"2022-04-20T11:22:53.68794Z","iopub.status.idle":"2022-04-20T11:22:53.693914Z","shell.execute_reply.started":"2022-04-20T11:22:53.687882Z","shell.execute_reply":"2022-04-20T11:22:53.693271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if have_index:\n    labels_map = {}\n    train_df = pd.read_csv(os.path.join(root_out,'labels_index.csv'))\n    def label_f(m):\n        labels_map[int(m.label_index)] = m.cultivar\n    train_df.apply(label_f,axis=1)\nelse:\n    labels_map = data_pre_access(os.path.join(root_in,'train_cultivar_mapping.csv'), output=os.path.join(root_out,'labels_index.csv'))\n    train_df = train_df = pd.read_csv(os.path.join(root_out,'labels_index.csv'))","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:34:11.175253Z","iopub.execute_input":"2022-04-20T11:34:11.175573Z","iopub.status.idle":"2022-04-20T11:34:11.827322Z","shell.execute_reply.started":"2022-04-20T11:34:11.175539Z","shell.execute_reply":"2022-04-20T11:34:11.826578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"check_sum = 0\nfor key, val in tqdm(labels_map.items()):\n    train_df[train_df.label_index == key].cultivar.unique() == val\n    check_sum += 1","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:34:12.996281Z","iopub.execute_input":"2022-04-20T11:34:12.996571Z","iopub.status.idle":"2022-04-20T11:34:13.084829Z","shell.execute_reply.started":"2022-04-20T11:34:12.996543Z","shell.execute_reply":"2022-04-20T11:34:13.083994Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"check_sum, check_sum==len(labels_map)","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:34:14.168291Z","iopub.execute_input":"2022-04-20T11:34:14.168585Z","iopub.status.idle":"2022-04-20T11:34:14.175736Z","shell.execute_reply.started":"2022-04-20T11:34:14.168554Z","shell.execute_reply":"2022-04-20T11:34:14.174659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class Sorghum_Train_Dataset(Dataset):\n    '''Train Dataset'''\n    def __init__(self, img_path_csv='', df=None, transform=None):\n        if df is not None:\n            self.df = df\n        else:\n            self.df = pd.read_csv(img_path_csv)\n        self.transform = transform\n        \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, index):\n        img = Img.open(os.path.join(root_in, 'train', self.df.iloc[index, 0]))\n        label_index = self.df.iloc[index, 4]\n        if self.transform is not None:\n            img = self.transform(img)\n        return img, label_index","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Sorghum_Test_Dataset(Sorghum_Train_Dataset):\n    '''Test Dataset'''\n    def __getitem__(self, index):\n        img = Image.open(os.path.join(root_in, 'test', self.df.iloc[index, 0]))\n        if self.transform:\n            img = self.transform(img)\n        return img","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Look Inside Sorghum breed","metadata":{"tags":[]}},{"cell_type":"code","source":"labels = pd.read_csv(os.path.join(root_out, 'labels_index.csv'))","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:24:30.568844Z","iopub.execute_input":"2022-04-20T11:24:30.569448Z","iopub.status.idle":"2022-04-20T11:24:30.625431Z","shell.execute_reply.started":"2022-04-20T11:24:30.569413Z","shell.execute_reply":"2022-04-20T11:24:30.624769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sns.catplot(y=\"cultivar\", kind=\"count\", data=labels, height=20)","metadata":{"tags":[],"execution":{"iopub.status.busy":"2022-04-20T11:28:37.303847Z","iopub.execute_input":"2022-04-20T11:28:37.304426Z","iopub.status.idle":"2022-04-20T11:28:39.463502Z","shell.execute_reply.started":"2022-04-20T11:28:37.304391Z","shell.execute_reply":"2022-04-20T11:28:39.462812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def first(ls):\n    for i, flag in enumerate(ls):\n        if flag == True:\n            return i\n        \nsamples = []\nfor cultivar in labels.cultivar.unique():\n    img, label = labels.iloc[first(labels.cultivar == cultivar), [0, 1]]\n    samples += [(img, label)]\n\nf, axarr = plt.subplots(3,3,figsize=(50,50))\nfor i in range(3):\n    for j in range(3):\n        axarr[i, j].imshow(Img.open(os.path.join(root_in, 'train', samples[3*i + j][0])))\n        axarr[i, j].set_title(samples[3*i + j][1])","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:28:41.738917Z","iopub.execute_input":"2022-04-20T11:28:41.74004Z","iopub.status.idle":"2022-04-20T11:28:50.761842Z","shell.execute_reply.started":"2022-04-20T11:28:41.739981Z","shell.execute_reply":"2022-04-20T11:28:50.760496Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CLAHE","metadata":{}},{"cell_type":"code","source":"def CLAHE_Convert(origin_input):\n    clahe = cv.createCLAHE(clipLimit=40, tileGridSize=(10,10))\n    t = np.asarray(origin_input)\n    t = cv.cvtColor(t, cv.COLOR_BGR2HSV)\n    t[:,:,-1] = clahe.apply(t[:,:,-1])\n    t = cv.cvtColor(t, cv.COLOR_HSV2BGR)\n    t = Img.fromarray(t)\n    return t","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:29:05.148354Z","iopub.execute_input":"2022-04-20T11:29:05.14925Z","iopub.status.idle":"2022-04-20T11:29:05.156482Z","shell.execute_reply.started":"2022-04-20T11:29:05.149198Z","shell.execute_reply":"2022-04-20T11:29:05.155505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"p0 = Img.open(os.path.join(root_in, 'train', samples[10*9 + 3][0]))\ng0 = transforms.Grayscale(num_output_channels=1)(p0)\nt0 = CLAHE_Convert(p0)\nn0 = transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])(transforms.ToTensor()(t0))\nnn0 = transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])(transforms.ToTensor()(p0))\n\np1 = Img.open(os.path.join(root_in, 'train', samples[10*7 + 9][0]))\ng1 = transforms.Grayscale(num_output_channels=1)(p1)\nt1 = CLAHE_Convert(p1)\nn1 = transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])(transforms.ToTensor()(t1))\nnn1 = transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])(transforms.ToTensor()(p1))\n\nf, axarr = plt.subplots(2,5,figsize=(25,10))\naxarr[0,0].imshow(g0, 'gray')\naxarr[0,0].set_title(\"GRAY\",fontsize=26)\naxarr[0,1].imshow(p0)\naxarr[0,1].set_title(\"ORIGINAL\",fontsize=26)\naxarr[0,2].imshow(t0)\naxarr[0,2].set_title(\"CLAHE\",fontsize=26)\naxarr[0,3].imshow(n0.permute(1, 2, 0))\naxarr[0,3].set_title(\"Normalize(CLAHE)\",fontsize=26)\naxarr[0,4].imshow(nn0.permute(1, 2, 0))\naxarr[0,4].set_title(\"Normalize(ORIGINAL)\",fontsize=26)\n\naxarr[1,0].imshow(g1, 'gray')\naxarr[1,1].imshow(p1)\naxarr[1,2].imshow(t1)\naxarr[1,3].imshow(n1.permute(1, 2, 0))\naxarr[1,4].imshow(nn1.permute(1, 2, 0))","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:29:06.348067Z","iopub.execute_input":"2022-04-20T11:29:06.348587Z","iopub.status.idle":"2022-04-20T11:29:10.4244Z","shell.execute_reply.started":"2022-04-20T11:29:06.348535Z","shell.execute_reply":"2022-04-20T11:29:10.423392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{"tags":[]}},{"cell_type":"code","source":"class ArcMarginProduct(nn.Module):\n    r\"\"\"Implement of large margin arc distance: :\n    Args:\n        in_features: size of each input sample\n        out_features: size of each output sample\n        s: norm of input feature\n        m: margin\n        cos(theta + m)\n    \"\"\"\n\n    def __init__(\n        self,\n        in_features: int,\n        out_features: int,\n        s: float,\n        m: float,\n        easy_margin: bool,\n        ls_eps: float,\n        rank\n    ):\n        super(ArcMarginProduct, self).__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.s = s\n        self.m = m\n        self.ls_eps = ls_eps  # label smoothing\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n\n        self.easy_margin = easy_margin\n        self.cos_m = math.cos(m)\n        self.sin_m = math.sin(m)\n        self.th = math.cos(math.pi - m)\n        self.mm = math.sin(math.pi - m) * m\n        self.rank = rank\n\n    def forward(self, input: torch.Tensor, label: torch.Tensor, device = 'cuda') -> torch.Tensor:\n        # --------------------------- cos(theta) & phi(theta) ---------------------\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        # Enable 16 bit precision\n        cosine = cosine.to(torch.float32)\n\n        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\n        phi = cosine * self.cos_m - sine * self.sin_m\n        if self.easy_margin:\n            phi = torch.where(cosine > 0, phi, cosine)\n        else:\n            phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n        # --------------------------- convert label to one-hot ---------------------\n        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')\n        one_hot = torch.zeros(cosine.size(), device=self.rank)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        if self.ls_eps > 0:\n            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n        # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n\n        return output","metadata":{"execution":{"iopub.status.busy":"2022-04-20T11:29:36.964157Z","iopub.execute_input":"2022-04-20T11:29:36.964498Z","iopub.status.idle":"2022-04-20T11:29:36.980996Z","shell.execute_reply.started":"2022-04-20T11:29:36.964461Z","shell.execute_reply":"2022-04-20T11:29:36.98019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SorghumModel(nn.Module):\n    def __init__(self, model_name, embedding_size, map_location, k_fold, rank, pretrained=True):\n        super(SorghumModel, self).__init__()       \n        \n        #model_effecient_b6 = timm.create_model(model_name, pretrained=pretrained, num_classes=NUM_CLASSES)\n        #global param_name\n        #param_name = [name for name,_ in model_effecient_b6.named_parameters()] # All parameters name\n        #del model_effecient_b6\n            \n        self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=NUM_CLASSES)\n        \n        #freeze_pretrained_layers(self.model)\n        #debarcle_layers(self.model, db_all=True) # Debarcle all layers()\n        \n        print('load Start!!!')\n        in_features = self.model.classifier.in_features\n        self.model.classifier = nn.Identity()\n        self.pooling = self.model.global_pool\n        self.model.global_pool = nn.Identity()\n        #self.pooling = GeM()\n        self.rank = rank\n        self.multiple_dropout = [nn.Dropout(0.25) for i in range(8)]\n        self.embedding = nn.Linear(in_features * 2, embedding_size)\n        self.fc = ArcMarginProduct(embedding_size, \n                                   NUM_CLASSES,\n                                   S, \n                                   M, \n                                   EASY_MERGING, \n                                   LS_EPS,\n                                  self.rank)\n\n    def forward(self, images, labels):\n        features = self.model(images)\n        pooled_features_avg = self.pooling(features).flatten(1)\n        pooled_features_max = nn.AdaptiveMaxPool2d((1,1))(features).flatten(1)\n        pooled_features = torch.cat((pooled_features_avg, pooled_features_max), dim=1)\n        pooled_features_dropout = torch.zeros((pooled_features.shape),device=self.rank)\n        for i in range(8):\n            pooled_features_dropout += self.multiple_dropout[i](pooled_features)\n        pooled_features_dropout /= 8\n        embedding = self.embedding(pooled_features_dropout)\n        #pooled_features = nn.Dropout(0.5)(pooled_features)\n        #embedding = self.embedding(pooled_features)\n        output = self.fc(embedding, labels)\n        return output\n    \n    def extract(self, images):\n        features = self.model(images)\n        pooled_features_avg = self.pooling(features).flatten(1)\n        pooled_features_max = nn.AdaptiveMaxPool2d((1,1))(features).flatten(1)\n        pooled_features = torch.cat((pooled_features_avg, pooled_features_max), dim=1)\n        embedding = self.embedding(pooled_features)\n        return embedding","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train ","metadata":{}},{"cell_type":"code","source":"def train_model_dist(rank,  world_size, model_name, num_epochs, loss_name, lr, lr_min, weight_decay, optim, use_amp, init, scheduler_type, warm_up):\n    '''Parameters:\n        lr(float): the begining learning rate\n        lr_min(float): min learning rate\n        optim(String): the optimizer type\n        use_amp(Boolean): Use mixed precision on GPU or not\n        init(Boolean): Need initial the layers parameter or not\n        scheduler_type(String): Learning rate scheduler\n        \n       Detail:\n        The train process will save the model's parameter every 10 epochs.\n        Every epoch, scheduler update once, and evaluate the train_dataset's accuracy and print it to std 5 times, \n        print the valid_dataset's accuracy once.\n        If the valid accuracy >= 0.75, save the model's parameters as well.\n    '''\n    seed_everything(SEED)\n    \n    if have_index:\n        labels_map = {}\n        train_df = pd.read_csv(os.path.join(root_out,'labels_index.csv'))\n        def label_f(m):\n            labels_map[int(m.label_index)] = m.cultivar\n        train_df.apply(label_f,axis=1)\n    else:\n        labels_map = data_pre_access(os.path.join(root_in,'train_cultivar_mapping.csv'), output=os.path.join(root_out,'labels_index.csv'))\n    \n    train_transform = transforms.Compose([\n        transforms.Resize(IMAGE_SIZE),\n        transforms.ColorJitter(brightness=0.2, contrast=0.05, saturation=0.1),\n        transforms.RandomHorizontalFlip(p=0.5),\n        transforms.RandomVerticalFlip(p=0.5),\n        transforms.RandomApply(transforms=\n                      [transforms.RandomResizedCrop(size=IMAGE_SIZE, scale=(0.3,0.4), \n                                                    ratio=(1/3,3),interpolation=\n                                                    transforms.InterpolationMode.BICUBIC)],p=0.2),\n        transforms.ToTensor(),\n        # Normalize to fit pretrained model\n        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])\n\n    val_test_transform = transforms.Compose([\n        transforms.Resize(IMAGE_SIZE),\n        transforms.ToTensor(),\n        # Normalize to fit pretrained model\n        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])\n\n    sfolder = StratifiedKFold(n_splits=FOLD,random_state=SEED,shuffle=True)\n    train_folds = []\n    val_folds = []\n    for train_idx, val_idx in sfolder.split(train_df.image, train_df.label_index):\n        train_folds.append(train_idx)\n        val_folds.append(val_idx)\n        print(len(train_folds), len(val_folds))\n        \n\n    \n    def init_xavier(m):\n        if type(m) == nn.Linear or type(m) == nn.Conv2d:\n            nn.init.xavier_normal_(m.weight)\n            \n    \n\n    if rank == 0:\n        os.makedirs('/root/tf-logs/' + model_name, exist_ok=True)\n        writer = SummaryWriter(log_dir='/root/tf-logs/' + model_name)\n    \n    \n    print('training on', rank) \n    setup(rank, world_size)\n    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}\n    \n    for k_fold in range(FOLD):\n        if rank == 0:  \n            Value_train_l = list()\n            Value_train_acc = list()\n            Value_test_acc = list()\n            Time = list()\n            print('\\n ********** Fold %d **********\\n'%k_fold)\n            sub_fold = model_name + '_F_' + str(k_fold)\n            os.makedirs(os.path.join(root_out, sub_fold), exist_ok=True)\n\n        train_dataset = Sorghum_Train_Dataset(df=train_df.iloc[train_folds[k_fold]],\n                                            transform=train_transform)\n\n        val_dataset = Sorghum_Train_Dataset(df=train_df.iloc[val_folds[k_fold]],\n                                          transform=val_test_transform)\n        \n        net = SorghumModel(model_name, EMBEDDING_SIZE, map_location, k_fold, rank)\n\n        net = net.to(rank)\n        net = DDP(net, device_ids=[rank], output_device=rank)\n\n        \n        #net.load_state_dict(torch.load(os.path.join(root_out, \n        #                                            model_name+ '_F_' + str(k_fold),\n        #                                            'Sorghum3.params'),\n        #                               map_location=map_location))\n        print('load Finish!!!')\n\n        train_loader = prepare(train_dataset, rank, world_size, num_workers=NUM_WORKERS)\n        val_loader = prepare(val_dataset, rank, world_size, num_workers=NUM_WORKERS)\n     \n        if optim == 'sgd':\n            optimizer = torch.optim.SGD((param for param in net.parameters() if param.requires_grad), lr=lr, weight_decay=weight_decay)\n        elif optim == 'adam':\n            optimizer = torch.optim.Adam((param for param in net.parameters() if param.requires_grad), lr=lr, weight_decay=weight_decay)\n        elif optim =='adamW':\n            optimizer = torch.optim.AdamW((param for param in net.parameters() if param.requires_grad), lr=lr, weight_decay=weight_decay)\n        elif optim == 'ranger':\n            optimizer = Ranger((param for param in net.parameters() if param.requires_grad), lr=lr, weight_decay=weight_decay)\n\n        scaler = torch.cuda.amp.GradScaler(enabled=use_amp) # mixed_precison\n\n        if scheduler_type == 'Cosine':\n            scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=lr_min)\n                \n        def warm_up_scheduler(epoch):\n            return (1 / 2) ** (warm_up-epoch)\n        \n        if loss_name == 'CrossEntropy':\n            loss = nn.CrossEntropyLoss(label_smoothing=0.1)\n            \n        if init:\n            net.apply(init_xavier)\n            \n        num_batches = len(train_loader)\n        best_accuracy = 0\n        \n        scheduler = LambdaLR(optimizer, lr_lambda=warm_up_scheduler)\n        \n        for epoch in range(num_epochs):\n            if epoch == warm_up:\n                if scheduler_type == 'Cosine':\n                    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs-warm_up, eta_min=lr_min)\n                \n            net.train()\n            train_loader.sampler.set_epoch(epoch)\n            val_loader.sampler.set_epoch(epoch)\n\n            metric = Accumulator(3)\n\n            for i, (X, y) in enumerate(tqdm(train_loader)):\n                X = X.to(rank)\n                y = y.to(rank)\n                with torch.cuda.amp.autocast(enabled=use_amp):\n                    y_hat = net(X, y)\n                    l = loss(y_hat, y)\n\n                scaler.scale(l).backward()\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad(set_to_none=True)\n\n                with torch.no_grad():\n                    metric.add(l * X.shape[0], accuracy(y_hat, y), X.shape[0])\n\n                if rank == 0:    \n                    train_l = metric[0] / metric[2]\n                    train_acc = metric[1] / metric[2]\n                    if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:\n                        print(epoch + (i + 1) / num_batches,\n                                     'train_l train_acc\\t',(train_l, train_acc,None))\n                        writer.add_scalars('Loss/Accuracy/train/Fold-' + str(k_fold), \n                                           {'train_accuracy':np.array(train_acc), 'train_loss':np.array(train_l)}, \n                                           5 * np.array(epoch + (i + 1) / num_batches))\n                        Value_train_l.append(train_l)\n                        Value_train_acc.append(train_acc)\n                        Value_test_acc.append(None)\n                        Time.append(epoch + (i + 1) / num_batches)\n\n            scheduler.step()\n            \n            test_acc = evaluate_accuracy(net, val_loader, device=rank)\n            \n            if rank == 0: \n                print('lr = ', optimizer.param_groups[0]['lr'])\n                print(epoch + 1,'test_acc\\t', (None, None, test_acc))\n                writer.add_scalars('Loss/Accuracy/test/Fold-' + str(k_fold), \n                                   {'val_accuracy':np.array(test_acc)}, \n                                   5 * np.array(epoch + 1))\n                Value_train_l.append(None)\n                Value_train_acc.append(None)\n                Value_test_acc.append(test_acc)\n                Time.append(epoch + 1)\n\n                if epoch % 10 == 0 or test_acc >= best_accuracy:\n                    best_accuracy = test_acc\n                    torch.save(net.state_dict(),os.path.join(root_out, sub_fold, 'Sorghum' + str(epoch + 1) + '.params'))\n                record_data = pd.DataFrame(zip(Value_train_l, Value_train_acc, Value_test_acc, Time))    \n                record_data.to_csv(os.path.join(root_out, sub_fold, 'Record_Sorghum.csv')) \n\n        if rank == 0:\n            torch.save(net.state_dict(),os.path.join(root_out, sub_fold, 'Sorghum.params'))\n\n            print(f'loss {train_l:.3f}, train acc {train_acc:.3f}, '\n                  f'test acc {test_acc:.3f}')\n            print(f'on {str(rank)}')\n        torch.cuda.empty_cache()\n    cleanup()\n    writer.close()","metadata":{"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train with DDP","metadata":{"tags":[]}},{"cell_type":"code","source":"if __name__ == '__main__':\n    print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')    \n    mp.spawn(\n        train_model_dist,\n        args=(WORLD_SIZE, MODEL_NAME, EPOCH, LOSS, LR, LR_MIN, WEIGHT_DECAY, OPTIM, USE_AMP, INIT, SCHEDULER, WARM_UP),\n        nprocs=WORLD_SIZE\n    )","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Convert notebook to .py file ","metadata":{"tags":[]}},{"cell_type":"code","source":"!jupyter nbconvert --to script Sorghum_DDP.ipynb","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Run your .py file in command line","metadata":{"tags":[]}},{"cell_type":"code","source":"'''Do remember run in command line, instead of in jupyter notebook\n   And delete the convert code above in .py file'''","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> python Sorghum_DDP.py","metadata":{}}]}