{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"# !pip install torchsummary\n\nimport os, sys, glob, gc, copy\nimport time\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\n# from torchsummary import summary\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport cv2\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check GPU is available or not\nis_cuda = torch.cuda.is_available()\nif is_cuda:\n    device = torch.device(\"cuda\")\n    print(\"GPU is available\")\nelse:\n    device = torch.device(\"cpu\")\n    print(\"GPU not available, CPU used\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_path = '../input/cassava-leaf-disease-classification/train_images/'\ntotal_csv = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_img(path):\n    im_bgr = cv2.imread(path)\n    im_rgb = im_bgr[:, :, ::-1]\n    #print(im_rgb)\n    return im_rgb\n\n\nclass HDD_Dataset(Dataset):\n    def __init__(self, df, data_root, device, transform=None):\n        super().__init__()\n        self.df = df.reset_index(drop=True).copy()\n        self.data_root = data_root\n        self.device = device\n        self.transform = transform\n        self.N_class = 5  # number of classes\n    \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, index: int):\n        # read label and one-hot\n        label = torch.tensor(self.df.iloc[index]['label'], device=device)\n        label_onehot = nn.functional.one_hot(label, self.N_class)\n        # prepare image\n        path = \"{}/{}\".format(self.data_root, self.df.iloc[index]['image_id'])\n        img = torch.tensor(get_img(path).copy()).float().to(device)/255  # normalize to 1\n        img = img.permute(2,0,1)\n        if self.transform:\n            img = self.transform(img)\n#         img = (img-0.5)*2  # normalize images to [-1, 1] and use tanh\n        return {'image':img, 'label':label, 'label_onehot':label_onehot.float()}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## split data to `test` and `train`"},{"metadata":{"trusted":true},"cell_type":"code","source":"def frac_train_val(N:int, val_frac:float):\n    \"\"\"N: dataset length e.g. 2100\n        val_frac: validation fraction of total e.g. 0.2 \n        return indices of train, validation\"\"\"\n    perm = np.random.permutation(N)\n    thrshld = int(N*val_frac)\n    return perm[thrshld:].tolist(), perm[:thrshld].tolist()\n\ninds_train, inds_test = frac_train_val(total_csv.shape[0], 0.1)\ntrain_csv = total_csv.loc[inds_train]\ntest_csv = total_csv.loc[inds_test]\n#create datasets\ncrop_transform = transforms.RandomCrop(512)\ntrain_dataset = HDD_Dataset(train_csv, data_path, device, crop_transform)\ntest_dataset = HDD_Dataset(test_csv, data_path, device, crop_transform)\n# gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define CVAE architectur"},{"metadata":{"code_folding":[11,51,131],"trusted":true},"cell_type":"code","source":"class print_layer(nn.Module):\n    def __init__(self):\n        super(print_layer, self).__init__()\n    def forward(self, x):\n        print(x.shape)\n        return x\n\nclass encoder(nn.Module):\n    \"\"\"encoder for CVAE\n        `image` input shape:(N, 3, 512, 512)\n        `c` shape: (N, 5)\"\"\"\n    def __init__(self, zdim:int, ydim:int):\n        super(encoder, self).__init__()\n        self.x_conv = nn.Sequential( \n            nn.Conv2d(3, 16, kernel_size=5, stride=2, padding=2), nn.LeakyReLU(),  # (256,256)\n            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),\n            nn.BatchNorm2d(32), nn.LeakyReLU(),\n            nn.Conv2d(32, 32, kernel_size=5, stride=1, padding=2),\n            nn.BatchNorm2d(32), nn.LeakyReLU(),\n            nn.Conv2d(32, 64, kernel_size=5, stride=2, padding=2),  # (128,128)\n            nn.BatchNorm2d(64), nn.LeakyReLU(),\n            nn.Conv2d(64, 64, kernel_size=5, stride=1, padding=2), nn.LeakyReLU(),\n            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),  # (64, 64)\n            nn.BatchNorm2d(64), nn.LeakyReLU(),\n            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(),\n            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),  # (32, 32)\n            nn.BatchNorm2d(128), nn.LeakyReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),\n            nn.BatchNorm2d(128), nn.LeakyReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),  # (16, 16)\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(),\n            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),  # (8, 8)\n            nn.BatchNorm2d(256), nn.LeakyReLU(),\n            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(),\n            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),  # (4, 4)\n            nn.BatchNorm2d(512), nn.LeakyReLU(), \n            nn.Conv2d(512, 1024, kernel_size=4, stride=1, padding=0), nn.LeakyReLU()\n        )\n        self.x_fc = nn.Sequential(\n            nn.Linear(in_features=1024, out_features=1024),\n            nn.ReLU(),\n            nn.Dropout()\n        )\n        self.y_fc = nn.Sequential(\n            nn.Linear(in_features=ydim, out_features=zdim), nn.ReLU(),\n            nn.Linear(in_features=zdim, out_features=zdim),  #????????? maybe remove\n        )\n        self.mu_lin = nn.Sequential(nn.Linear(in_features=1024+zdim, out_features=zdim))\n        self.var_lin = nn.Sequential(nn.Linear(in_features=1024+zdim, out_features=zdim))\n\n    def forward(self, x, y):\n        x = self.x_conv(x)\n        x = x.view(-1, 1024)\n        x = self.x_fc(x)\n        y = self.y_fc(y)\n        cat = torch.cat([x, y],dim=1)\n        mean = self.mu_lin(cat)\n        log_sigma = self.var_lin(cat) + 1e-6\n        return mean, log_sigma\n\n\nclass decoder(nn.Module):\n    def __init__(self, zdim:int, ydim:int):\n        super(decoder, self).__init__()\n        self.y_decode = nn.Sequential(\n            nn.Linear(in_features=ydim, out_features=zdim*2),\n            nn.ReLU()\n        )\n        self.linearMix = nn.Sequential(\n            nn.Linear(in_features=zdim*3, out_features=256),\n            nn.ReLU()\n        )\n        self.to_conv = nn.Linear(in_features=256, out_features=256*8*8)\n        self.conv_decoder = nn.Sequential(\n            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1), # (8, 8)\n            nn.BatchNorm2d(256), nn.ReLU(),\n            nn.ConvTranspose2d(256, 256, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),  # (16, 16)\n            nn.BatchNorm2d(128), nn.ReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),\n            nn.BatchNorm2d(128), nn.ReLU(),\n            nn.ConvTranspose2d(128, 128, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),  # (32, 32)\n            nn.BatchNorm2d(128), nn.ReLU(),\n            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),\n            nn.BatchNorm2d(128), nn.ReLU(),\n            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),  # (64, 64)\n            nn.BatchNorm2d(64), nn.ReLU(),\n            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),\n            nn.BatchNorm2d(64), nn.ReLU(),\n            nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),  # (128, 128)\n            nn.BatchNorm2d(32), nn.ReLU(),\n            nn.Conv2d(32, 32, kernel_size=5, stride=1, padding=2),\n            nn.BatchNorm2d(32), nn.ReLU(),\n            nn.ConvTranspose2d(32, 32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(32, 16, kernel_size=5, stride=1, padding=2),  # (256, 256)\n            nn.BatchNorm2d(16), nn.ReLU(),\n            nn.Conv2d(16, 16, kernel_size=5, stride=1, padding=2),\n            nn.BatchNorm2d(16), nn.ReLU(),\n            nn.ConvTranspose2d(16, 16, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),\n            nn.Conv2d(16, 8, kernel_size=5, stride=1, padding=2), # (512, 512)\n            nn.BatchNorm2d(8), nn.ReLU(),\n            nn.Conv2d(8, 4, kernel_size=5, stride=1, padding=2), nn.ReLU(),\n            nn.Conv2d(4, 3, kernel_size=5, stride=1, padding=2),\n            nn.Sigmoid()\n        )\n    \n    def forward(self, z, y):\n        y = self.y_decode(y)\n#         z = copy.deepcopy(z)    #?????????????\n        z = torch.cat((y,z),dim=1)\n        z = self.linearMix(z)\n        z = self.to_conv(z)\n        z = z.view(-1, 256,8,8)\n        z = self.conv_decoder(z)\n        return z\n\n\nclass CVAE(nn.Module):\n    def __init__(self, zdim:int, ydim:int):\n        super(CVAE, self).__init__()\n        self.zdim = zdim\n        self.ydim = ydim\n        self.encoder = encoder(zdim, ydim)\n        self.decoder = decoder(zdim, ydim)\n        \n    def sampler(self, mu, log_sigma):\n        std = torch.exp(log_sigma / 2)\n        eps = torch.randn_like(std)\n        x_sample = eps.mul(std) + mu\n        return x_sample\n\n    def forward(self, x, y):\n        mu, log_sigma = self.encoder(x, y)\n        z = self.sampler(mu, log_sigma)\n        out_img = self.decoder(z, y)\n        return out_img, mu, log_sigma","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### test model output shape"},{"metadata":{"trusted":true},"cell_type":"code","source":"zdim = 100\nydim = 5\n####\nmodel = CVAE(zdim, ydim)\n\nx = torch.randn(2,3,512,512)\ny = F.one_hot(torch.randint(5,size=(2,)), 5).float()\nout, mu, log_sigma = model(x, y)\nprint(out.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"def calculate_loss(x, reconstructed_x, mean, log_var):\n    # reconstruction loss\n    Num = mean.shape[0]\n    RCL = F.binary_cross_entropy(reconstructed_x, x, reduction='sum')\n    # kl divergence loss\n    KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())\n    return (RCL + KLD)/Num","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define model and optimizers"},{"metadata":{"trusted":true},"cell_type":"code","source":"# hyper parameters\nzdim = 100\nydim = 5\n# Create model instance\nmodel = CVAE(zdim, ydim)\nmodel = model.to(device, non_blocking=True)  # to GPU or CPU\n# summary(model, input_size=(3, 512, 512))\n# Define hyperparameters\nlr = torch.tensor(0.00001).to(device)\nlr_decay = torch.tensor(0.99).to(device)  # per epoch\nlr_floor = torch.tensor(0.000005).to(device)\n# Define Loss, Optimizer\noptimizer = torch.optim.Adam(model.parameters(), lr=lr)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!nvidia-smi --gpu-reset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Hyper parameters\nbatch_size = 48\nn_epochs = 10\n# dataloader\ntrain_dataloader = DataLoader(train_dataset,batch_size=batch_size, shuffle=True)\ntot_batch = len(train_dataloader)\n# loss\ntrain_loss = []\n# Training RNN\nfor epoch in range(1, n_epochs + 1):\n    epoch_start = time.time()\n    batch_start = time.time()\n    for num_batch, data_batch in enumerate(train_dataloader):\n        \n        optimizer.zero_grad()\n        output, mu, log_sigma = model(data_batch['image']-1/2, data_batch['label_onehot'])  # zero-mean preprocessing ??????\n        loss = calculate_loss(data_batch['image'], output, mu, log_sigma)\n        train_loss.append(loss.item())\n#         loss.backward(torch.tensor(1/100,device=device))\n        loss.backward()\n        optimizer.step()\n        \n        batch_stop = time.time()\n        print(f\"*****Epoch {epoch}, minibatch: {num_batch+1}/{tot_batch} elapsed {batch_stop-batch_start:0.1f}s, loss:{loss:0.4f}*****\")  # write in replace\n        batch_start = time.time()\n    epoch_stop = time.time()\n    print(f\"Epoch {epoch}/{n_epochs} has finished in {epoch_stop-epoch_start:0.1f}s\")\n    lr *= lr_decay\n    if lr < lr_floor:\n        lr = lr_floor\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model, 'CVAE_model_v2.h5')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plot loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1,1)\nax.plot(range(1,1+len(train_loss)), train_loss)\nax.set(xlabel='minibatch')\nax.set_title('CVAE loss')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model.eval()\n\nN = 4\ntmp_dec = model.decoder\nz = torch.randn(N, zdim).to(device)\nclass_label = torch.randint(5,size=(N,))\ny = F.one_hot(class_label, 5).float().to(device)\nout = tmp_dec(z, y)\n\n# model.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax =plt.subplots(1, N)\nfor i in range(N):\n    iimg = out[i].to(torch.device('cpu')).permute(1,2,0)\n    ax[i].imshow(iimg.detach().numpy())\n    ax[i].axis('off')\n    ax[i].set_title(f'cls:{class_label[i]}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load model from last training"},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls -ltrh\nmodel = torch.load('CVAE_model_v1.h5')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}