{ "cells": [ { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import cv2\n", "import os\n", "from tqdm import tqdm,trange\n", "import sklearn.metrics\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "\n", "import warnings\n", "warnings.filterwarnings(\"ignore\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Initializing a project" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Project name: deepfake-shield-mayukh\n", "Artifacts path: /media/mayukh/Data/storage/repositories/repos/deep-shield-temp/notebooks/artifact_dir\n" ] } ], "source": [ "from os import path\n", "import mlrun\n", "\n", "# Set the base project name\n", "project_name_base = 'deepfake-shield'\n", "# Initialize the MLRun environment and save the project name and artifacts path\n", "project_name, artifact_path = mlrun.set_environment(project=project_name_base,\n", " user_project=True, api_path = './api_dir', artifact_path = './artifact_dir')\n", " \n", "# Display the current project name and artifacts path\n", "print(f'Project name: {project_name}')\n", "print(f'Artifacts path: {artifact_path}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Initialize training code with `mlrun`" ] }, { "cell_type": "code", "execution_count": 87, "metadata": {}, "outputs": [], "source": [ "# mlrun: start-code" ] }, { "cell_type": "code", "execution_count": 88, "metadata": {}, "outputs": [], "source": [ "import os\n", "import cv2\n", "import numpy as np\n", "import pandas as pd\n", "import sklearn.metrics\n", "from tqdm import tqdm,trange\n", "import matplotlib.pyplot as plt\n", "from mlrun.datastore import DataItem \n", "from mlrun.artifacts import PlotArtifact\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "\n", "import warnings\n", "warnings.filterwarnings(\"ignore\")\n", "\n", "import gc \n", "import random\n", "from torch.utils.data import Dataset, DataLoader\n", "from pytorchcv.model_provider import get_model as ptcv_get_model" ] }, { "cell_type": "code", "execution_count": 89, "metadata": {}, "outputs": [], "source": [ "\n", "class ImageDataset(Dataset):\n", " def __init__(self, x, y, training=True, transform=None):\n", " self.x = x\n", " self.y = y\n", " self.transform = transform\n", " self.training = training\n", "\n", " def __len__(self):\n", " return len(self.x)\n", " \n", " def __getitem__(self, idx):\n", " if torch.is_tensor(idx):\n", " idx = idx.tolist()\n", " \n", " img_path = self.x[idx]\n", " \n", " img_numpy = read_img(img_path)\n", "# print(img_numpy.max(), img_numpy.min())\n", "\n", " if self.transform is not None:\n", " res = self.transform(image=img_numpy)\n", " img_numpy = res['image']\n", " \n", " img_tensor = torch.tensor(img_numpy).permute(-1,0,1)\n", " \n", "\n", " labels = self.y[idx]\n", " return {\n", " 'image': img_tensor,\n", " 'label': labels\n", " }" ] }, { "cell_type": "code", "execution_count": 90, "metadata": {}, "outputs": [], "source": [ "class Head(torch.nn.Module):\n", " def __init__(self, in_f, out_f):\n", " super(Head, self).__init__()\n", "\n", " self.f = nn.Flatten()\n", " self.l = nn.Linear(in_f, 512)\n", " self.d = nn.Dropout(0.75)\n", " self.o = nn.Linear(512, out_f)\n", " self.b1 = nn.BatchNorm1d(in_f)\n", " self.b2 = nn.BatchNorm1d(512)\n", " self.r = nn.ReLU()\n", "\n", " def forward(self, x):\n", "# print(x.shape)\n", " x = x.reshape(x.size(0), -1)\n", " x = self.b1(x)\n", " x = self.d(x)\n", "\n", " x = self.l(x)\n", " x = self.r(x)\n", " x = self.b2(x)\n", " x = self.d(x)\n", "\n", " out = self.o(x)\n", " return out\n", " \n", "class FCN(torch.nn.Module):\n", " def __init__(self, base, in_f):\n", " super(FCN, self).__init__()\n", " self.base = base\n", " self.h1 = Head(in_f, 1)\n", " \n", " def forward(self, x):\n", " x = self.base(x)\n", " return self.h1(x)\n" ] }, { "cell_type": "code", "execution_count": 91, "metadata": {}, "outputs": [], "source": [ "def get_model():\n", " model = ptcv_get_model(\"efficientnet_b4\")\n", " model = nn.Sequential(*list(model.children())[:-1]) # Remove original output layer\n", " model[0].final_block.pool = nn.Sequential(nn.AdaptiveAvgPool2d(1))\n", " model = FCN(model, 1792)\n", " model = model.cuda()\n", " \n", " return model" ] }, { "cell_type": "code", "execution_count": 92, "metadata": {}, "outputs": [], "source": [ "def read_img(path):\n", " return cv2.cvtColor(cv2.imread(path),cv2.COLOR_BGR2RGB)\n", "\n", "def shuffle(X,y):\n", " new_train=[]\n", " for m,n in zip(X,y):\n", " new_train.append([m,n])\n", " random.shuffle(new_train)\n", " X,y=[],[]\n", " for x in new_train:\n", " X.append(x[0])\n", " y.append(x[1])\n", " return X,y\n", "\n", "def get_random_sampling(paths, y, val_paths, val_y, num_train_samples = None, num_val_samples = None):\n", " real=[]\n", " fake=[]\n", " for m,n in zip(paths,y):\n", " if n==0:\n", " real.append(m)\n", " else:\n", " fake.append(m)\n", " # fake=random.sample(fake,len(real))\n", " paths,y=[],[]\n", " for x in real:\n", " paths.append(x)\n", " y.append(0)\n", " for x in fake:\n", " paths.append(x)\n", " y.append(1)\n", "\n", " real=[]\n", " fake=[]\n", " \n", " for m,n in zip(val_paths,val_y):\n", " if n==0:\n", " real.append(m)\n", " else:\n", " fake.append(m)\n", " \n", " \n", " val_paths,val_y=[],[]\n", " for x in real:\n", " val_paths.append(x)\n", " val_y.append(0)\n", " for x in fake:\n", " val_paths.append(x)\n", " val_y.append(1)\n", "\n", " X=[]\n", " for img in tqdm(paths):\n", " X.append(img)\n", " val_X=[]\n", " for img in tqdm(val_paths):\n", " val_X.append(img)\n", "\n", " # Balance with ffhq dataset\n", " ffhq = os.listdir('training_data/ffhq/thumbnails128x128')\n", " X_ = []\n", " for file in tqdm(ffhq):\n", " path = f'training_data/ffhq/thumbnails128x128/{file}'\n", " X_.append(path)\n", " random.shuffle(X_)\n", "\n", " for i in range(64773 - 12130):\n", " X.append(X_[i])\n", " y.append(0)\n", " \n", " \n", " del X_[0:64773 - 12130]\n", "\n", " for i in range(6108 - 1258):\n", " val_X.append(X_[i])\n", " val_y.append(0)\n", "\n", " X, y = shuffle(X,y)\n", " val_X, val_y = shuffle(val_X,val_y)\n", " \n", " if num_train_samples is not None:\n", " X, y = X[:num_train_samples], y[:num_train_samples]\n", " \n", " if num_val_samples is not None:\n", " val_X, val_y = X[:num_val_samples], y[:num_val_samples]\n", "\n", " return X, val_X, y, val_y\n", "\n", "def train_n_epochs(model, num_epochs , train_loader, val_loader, checkpoint_name = 'model.pth', lr = 0.001, loss = None, context = None):\n", " \n", " best_acc = 0.\n", " accuracies = []\n", " losses = []\n", " \n", " optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n", " scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, mode='min', factor=0.7, verbose=True, min_lr=1e-5)\n", "\n", " for epoch in range(num_epochs):\n", " torch.cuda.empty_cache()\n", " gc.collect()\n", "\n", " model, loss = train_model(model, epoch, optimizer, scheduler, loss = loss, train_loader= train_loader)\n", " losses.append(loss)\n", " \n", " acc = evaluate_model(model, val_loader)\n", " accuracies.append(acc)\n", " \n", " if acc > best_acc:\n", " best_acc = acc\n", " print(f'Saving best model: {checkpoint_name}')\n", " torch.save(model.state_dict(), checkpoint_name)\n", " \n", " if context is not None:\n", " context.logger.info(f'epoch: {epoch+1} current acc: {round(acc, 4)} best acc: {round(best_acc, 4)}')\n", " \n", " \n", " fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (12,4))\n", " \n", " ax[0].plot(losses)\n", " ax[0].set_ylabel('Loss', fontsize = 18)\n", " ax[0].set_xlabel('Epochs', fontsize = 18)\n", " \n", " ax[1].plot(accuracies, c = 'g')\n", " ax[1].set_ylabel('Accuracy', fontsize = 18)\n", " ax[1].set_xlabel('Epochs', fontsize = 18)\n", " \n", " context.log_artifact(PlotArtifact('training logs', body=fig))\n", "\n", " return model" ] }, { "cell_type": "code", "execution_count": 93, "metadata": {}, "outputs": [], "source": [ "\n", "def evaluate_model(model, val_loader):\n", " model.eval()\n", " pred = []\n", " real = []\n", " with torch.no_grad():\n", " for batch in val_loader:\n", " \n", " img_batch, y_batch = batch['image'].cuda().float(), batch['label'].cuda().float()\n", "\n", " o1 = model(img_batch)\n", " \n", " for j in o1:\n", " pred.append(F.sigmoid(j))\n", " \n", " for i in y_batch:\n", " real.append(i.data.cpu())\n", " \n", " pred = [p.data.cpu().numpy() for p in pred]\n", " pred = [np.round(p) for p in pred]\n", " pred = np.array(pred)\n", " acc = sklearn.metrics.recall_score(real, pred, average='macro')\n", "\n", " return acc\n", "\n", "def criterion(pred1, targets):\n", " l1 = F.binary_cross_entropy(F.sigmoid(pred1).squeeze(-1), targets)\n", " return l1\n", "\n", "def train_model(model, epoch, optimizer, scheduler, loss , train_loader):\n", " model.train()\n", " total_loss = 0\n", "\n", " t = tqdm(train_loader, disable = True)\n", " for i, batch in enumerate(train_loader):\n", " \n", " img_batch = batch['image']\n", " img_batch = img_batch.cuda()\n", "\n", " y_batch = batch['label']\n", " y_batch = y_batch.cuda().float()\n", "\n", " optimizer.zero_grad()\n", "\n", " out = model(img_batch)\n", " loss = criterion(out, y_batch)\n", " \n", "\n", " total_loss += loss\n", " t.set_description(f'Epoch {epoch+1}, LR: %6f, Loss: %.4f'%(optimizer.state_dict()['param_groups'][0]['lr'],total_loss/(i+1)))\n", "\n", " loss.backward()\n", " optimizer.step()\n", " \n", " scheduler.step(1.)\n", "\n", " return model, total_loss" ] }, { "cell_type": "code", "execution_count": 94, "metadata": {}, "outputs": [], "source": [ "from mlrun.mlutils.data import get_sample\n", "import pandas as pd \n", "import albumentations\n", "from albumentations import (\n", " ShiftScaleRotate, \n", " HorizontalFlip, \n", " Normalize, \n", " GaussNoise, \n", ")\n", "\n", "def get_training_data(train_dataset: DataItem, val_dataset: DataItem, \n", " label_column = 'labels', df_train_path = 'df_train.csv',\n", " df_val_path = 'df_val.csv', ideal_params = None): \n", " \n", " df_train = pd.read_csv(df_train_path)\n", " df_val = pd.read_csv(df_val_path)\n", " \n", " mean = [0.485, 0.456, 0.406]\n", " std = [0.229, 0.224, 0.225]\n", "\n", " train_transform = albumentations.Compose([\n", " ShiftScaleRotate(p=0.3, scale_limit=0.25, border_mode=1, rotate_limit=25),\n", " HorizontalFlip(p=0.2),\n", " GaussNoise(p=.2),\n", " Normalize(mean = mean, std = std),\n", " albumentations.augmentations.geometric.resize.Resize(128,128)\n", " ])\n", " \n", " val_transform = albumentations.Compose([\n", " Normalize(mean = mean, std = std),\n", " albumentations.augmentations.geometric.resize.Resize(128,128)\n", " ])\n", "\n", " \n", " raw, labels, header = get_sample(train_dataset, sample=-1, label=label_column)\n", " raw_val, labels_val, header = get_sample(val_dataset, sample=-1, label=label_column) \n", " \n", " X, y = raw['paths'].values, labels.values\n", " val_X, val_y = raw_val['paths'].values, labels_val.values\n", "\n", " X, val_X, Y, val_Y = get_random_sampling(\n", " df_train['paths'].values, \n", " df_train['labels'].values, \n", " df_val['paths'].values,\n", " df_val['labels'].values, \n", " num_train_samples= None, # Set this to None if you want to train on the whole dataset\n", " num_val_samples= None # Set this to None if you want to train on the whole dataset\n", " )\n", "\n", " train_dataset = ImageDataset(X, Y, transform=train_transform)\n", " val_dataset = ImageDataset(val_X, val_Y, transform=val_transform)\n", "\n", " train_loader = DataLoader(dataset=train_dataset, batch_size=ideal_params['batch_size'], shuffle=True, num_workers=4)\n", " val_loader = DataLoader(dataset=val_dataset, batch_size= 10, shuffle=False, num_workers=0)\n", "\n", " return {\n", " 'train_loader': train_loader,\n", " 'val_loader': val_loader\n", " }\n", "\n", "\n", "\n", "def train_model_auto(context, train_dataset: DataItem, val_dataset: DataItem, \n", " num_epochs, hyperparams, checkpoint_name, df_train_path, df_val_path):\n", " \n", " print(\"started training :)\")\n", " \n", " data = get_training_data(\n", " train_dataset= train_dataset,\n", " val_dataset= val_dataset,\n", " df_train_path = df_train_path,\n", " df_val_path = df_val_path,\n", " ideal_params= hyperparams\n", " )\n", " \n", " model = get_model()\n", " \n", " train_n_epochs(\n", " model = model,\n", " num_epochs = num_epochs,\n", " train_loader= data['train_loader'],\n", " val_loader= data['val_loader'],\n", " lr = hyperparams['learning_rate'],\n", " loss = criterion,\n", " checkpoint_name = checkpoint_name,\n", " context = context ## context passed into fn\n", " )\n", " " ] }, { "cell_type": "code", "execution_count": 95, "metadata": {}, "outputs": [], "source": [ "# mlrun: end-code" ] }, { "cell_type": "code", "execution_count": 96, "metadata": {}, "outputs": [], "source": [ "train_func = mlrun.code_to_function(\n", " name='train_model_auto',\n", " kind='job',\n", " image='mlrun/mlrun'\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Training the model with `mlrun`" ] }, { "cell_type": "code", "execution_count": 97, "metadata": {}, "outputs": [], "source": [ "train_dataset = f'store://{project_name}/prep_data_deepfake_dataset_train'\n", "val_dataset = f'store://{project_name}/prep_data_deepfake_dataset_val'" ] }, { "cell_type": "code", "execution_count": 84, "metadata": {}, "outputs": [], "source": [ "results = pd.read_csv('artifact_dir/iteration_results.csv')\n", "best_iter = results[results['output.loss'] == results['output.loss'].min()]\n", "\n", "ideal_params = {\n", " 'learning_rate': best_iter['param.learning_rate'].values[0],\n", " 'batch_size': best_iter['param.batch_size'].values[0]\n", "}" ] }, { "cell_type": "code", "execution_count": 98, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'learning_rate': 0.002508, 'batch_size': 16}" ] }, "execution_count": 98, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ideal_params" ] }, { "cell_type": "code", "execution_count": 99, "metadata": { "scrolled": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "> 2021-07-03 13:17:02,549 [info] starting run train_model_auto uid=a0174e29d33a47d2adf827fe393c5e02 DB=./api_dir\n", "started training :)\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████| 104890/104890 [00:00<00:00, 4958585.11it/s]\n", "100%|██████████| 7366/7366 [00:00<00:00, 4893133.24it/s]\n", "100%|██████████| 70000/70000 [00:00<00:00, 3777388.26it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Saving best model: model.pt\n", "> 2021-07-03 13:17:07,229 [info] epoch: 1 current acc: 0.5 best acc: 0.5\n", "> 2021-07-03 13:17:08,558 [info] epoch: 2 current acc: 0.5 best acc: 0.5\n", "> 2021-07-03 13:17:09,903 [info] epoch: 3 current acc: 0.5 best acc: 0.5\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:11,580 [info] epoch: 4 current acc: 0.5244 best acc: 0.5244\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:13,270 [info] epoch: 5 current acc: 0.5366 best acc: 0.5366\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:15,008 [info] epoch: 6 current acc: 0.6463 best acc: 0.6463\n", "Epoch 7: reducing learning rate of group 0 to 1.7556e-03.\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:16,690 [info] epoch: 7 current acc: 0.7195 best acc: 0.7195\n", "> 2021-07-03 13:17:18,032 [info] epoch: 8 current acc: 0.6707 best acc: 0.7195\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:19,728 [info] epoch: 9 current acc: 0.7805 best acc: 0.7805\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:21,410 [info] epoch: 10 current acc: 0.8049 best acc: 0.8049\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:23,098 [info] epoch: 11 current acc: 0.878 best acc: 0.878\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:24,772 [info] epoch: 12 current acc: 0.9146 best acc: 0.9146\n", "Epoch 13: reducing learning rate of group 0 to 1.2289e-03.\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:26,451 [info] epoch: 13 current acc: 0.939 best acc: 0.939\n", "Saving best model: model.pt\n", "> 2021-07-03 13:17:28,121 [info] epoch: 14 current acc: 0.9512 best acc: 0.9512\n", "> 2021-07-03 13:17:29,573 [info] epoch: 15 current acc: 0.9268 best acc: 0.9512\n" ] }, { "data": { "text/html": [ "\n", "
project | \n", "uid | \n", "iter | \n", "start | \n", "state | \n", "name | \n", "labels | \n", "inputs | \n", "parameters | \n", "results | \n", "artifacts | \n", "
---|---|---|---|---|---|---|---|---|---|---|
deepfake-shield-mayukh | \n", "...3c5e02 | \n",
" 0 | \n", "Jul 03 07:47:02 | \n", "completed | \n", "train_model_auto | \n", "kind= owner=mayukh host=leopard | \n",
" train_dataset val_dataset | \n",
" num_epochs=15 hyperparams={'learning_rate': 0.002508, 'batch_size': 16} checkpoint_name=model.pt df_train_path=df_train.csv df_val_path=df_val.csv | \n",
" \n", " | training logs | \n",
"