{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Training a binary image classifier for deepfakes\n\nEarlier I published [this kernel](https://www.kaggle.com/humananalog/inference-demo) that shows how to do inference on the Deepfakes competition. It includes a ResNeXt50 checkpoint that I trained on my own machine and that gets 0.46788 on the leaderboard.\n\nI did not include code for training that ResNeXt50 model because my process involves running a bunch of different scripts to create and clean up a training set, amongst other things. It's not feasible to do all of this in a Kaggle kernel.\n\nHowever, Kaggler [dagnelies](https://www.kaggle.com/dagnelies) has made [this dataset of face crops](https://www.kaggle.com/dagnelies/deepfake-faces) available. Using this dataset, we can actually train our model using a Kaggle kernel!\n\n**Tip:** Enable GPU for this. It's *really slow* using CPU."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os, sys, random\nimport numpy as np\nimport pandas as pd\nimport cv2\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom tqdm.notebook import tqdm\n\n%matplotlib inline\nimport matplotlib.pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gpu = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\ngpu","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_size = 224\nbatch_size = 64","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The data\n\nYou need a bunch of face crops that are extracted from the full training dataset. Because we can't use the full training data from a kernel, I'm using [this dataset](https://www.kaggle.com/dagnelies/deepfake-faces), which is a nice starting point. You'll probably want to create your own dataset of face crops, as the quality of your training data makes all the difference."},{"metadata":{"_cell_guid":"","_uuid":"","trusted":true},"cell_type":"code","source":"crops_dir = \"../input/deepfake-faces/faces_155\"\n\nmetadata_df = pd.read_csv(\"../input/deepfake-faces/metadata.csv\")\nmetadata_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(metadata_df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"How many faces of each class do we have?"},{"metadata":{"trusted":true},"cell_type":"code","source":"len(metadata_df[metadata_df.label == \"REAL\"]), len(metadata_df[metadata_df.label == \"FAKE\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Look at a random face image:"},{"metadata":{"trusted":true},"cell_type":"code","source":"img_path = os.path.join(crops_dir, np.random.choice(os.listdir(crops_dir)))\nplt.imshow(cv2.imread(img_path)[..., ::-1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The dataset and data loaders"},{"metadata":{},"cell_type":"markdown","source":"Like most other torchvision models, the model we're using (ResNeXt50) requires that input images are normalized using mean and stddev. For making plots, we also define an \"unnormalize\" transform that can take a normalized image and turn it back into regular pixels."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.transforms import Normalize\n\nclass Unnormalize:\n    \"\"\"Converts an image tensor that was previously Normalize'd\n    back to an image with pixels in the range [0, 1].\"\"\"\n    def __init__(self, mean, std):\n        self.mean = mean\n        self.std = std\n\n    def __call__(self, tensor):\n        mean = torch.as_tensor(self.mean, dtype=tensor.dtype, device=tensor.device).view(3, 1, 1)\n        std = torch.as_tensor(self.std, dtype=tensor.dtype, device=tensor.device).view(3, 1, 1)\n        return torch.clamp(tensor*std + mean, 0., 1.)\n\n\nmean = [0.485, 0.456, 0.406]\nstd = [0.229, 0.224, 0.225]\nnormalize_transform = Normalize(mean, std)\nunnormalize_transform = Unnormalize(mean, std)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"During training, we'll apply data augmentation. In this kernel we just do random horizontal flips, but you can add other image transformations here too, such as rotation, zooming, etc. It's possible to use torchvision transforms for this, or a library such as [imgaug](https://www.github.com/aleju/imgaug), but I rolled my own using OpenCV2."},{"metadata":{"trusted":true},"cell_type":"code","source":"def random_hflip(img, p=0.5):\n    \"\"\"Random horizontal flip.\"\"\"\n    if random.random() < p:\n        return cv2.flip(img, 1)\n    else:\n        return img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Some helper code for loading a training image and its label:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_image_and_label(filename, cls, crops_dir, image_size, augment):\n    \"\"\"Loads an image into a tensor. Also returns its label.\"\"\"\n    img = cv2.imread(os.path.join(crops_dir, filename))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    if augment: \n        img = random_hflip(img)\n\n    img = cv2.resize(img, (image_size, image_size))\n\n    img = torch.tensor(img).permute((2, 0, 1)).float().div(255)\n    img = normalize_transform(img)\n\n    target = 1 if cls == \"FAKE\" else 0\n    return img, target","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"It's always smart to test that the code actually works. The following cell should return a normalized PyTorch tensor of shape (3, 224, 224) and the target 1 (for fake).\n\nNote that this dataset has 155x155 images but our model needs at least 224x224, so we resize them."},{"metadata":{"trusted":true},"cell_type":"code","source":"img, target = load_image_and_label(\"aabuyfvwrh.jpg\", \"FAKE\", crops_dir, 224, augment=True)\nimg.shape, target","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"To plot the image, we need to unnormalize it and also permute it from (3, 224, 224) to (224, 224, 3). "},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.imshow(unnormalize_transform(img).permute((1, 2, 0)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"To use the PyTorch data loader, we need to create a `Dataset` object.\n\nBecause of the class imbalance (many more fakes than real videos), we're using a dataset that samples a given number of REAL faces and the same number of FAKE faces, so it's always 50-50."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import Dataset\n\nclass VideoDataset(Dataset):\n    \"\"\"Face crops dataset.\n\n    Arguments:\n        crops_dir: base folder for face crops\n        df: Pandas DataFrame with metadata\n        split: if \"train\", applies data augmentation\n        image_size: resizes the image to a square of this size\n        sample_size: evenly samples this many videos from the REAL\n            and FAKE subfolders (None = use all videos)\n        seed: optional random seed for sampling\n    \"\"\"\n    def __init__(self, crops_dir, df, split, image_size, sample_size=None, seed=None):\n        self.crops_dir = crops_dir\n        self.split = split\n        self.image_size = image_size\n        \n        if sample_size is not None:\n            real_df = df[df[\"label\"] == \"REAL\"]\n            fake_df = df[df[\"label\"] == \"FAKE\"]\n            sample_size = np.min(np.array([sample_size, len(real_df), len(fake_df)]))\n            print(\"%s: sampling %d from %d real videos\" % (split, sample_size, len(real_df)))\n            print(\"%s: sampling %d from %d fake videos\" % (split, sample_size, len(fake_df)))\n            real_df = real_df.sample(sample_size, random_state=seed)\n            fake_df = fake_df.sample(sample_size, random_state=seed)\n            self.df = pd.concat([real_df, fake_df])\n        else:\n            self.df = df\n\n        num_real = len(self.df[self.df[\"label\"] == \"REAL\"])\n        num_fake = len(self.df[self.df[\"label\"] == \"FAKE\"])\n        print(\"%s dataset has %d real videos, %d fake videos\" % (split, num_real, num_fake))\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        filename = row[\"videoname\"][:-4] + \".jpg\"\n        cls = row[\"label\"]\n        return load_image_and_label(filename, cls, self.crops_dir, \n                                    self.image_size, self.split == \"train\")\n    def __len__(self):\n        return len(self.df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's test that the dataset actually works..."},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = VideoDataset(crops_dir, metadata_df, \"val\", image_size, sample_size=1000, seed=1234)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.imshow(unnormalize_transform(dataset[0][0]).permute(1, 2, 0))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Split up the data into train / validation. There are many different ways to do this. For this kernel, we're going to just grab a percentage of the REAL faces as well as their corresponding FAKEs. This way, a real video and all the fakes that are derived from it will be either completely in the training set or completely in the validation set. \n\n(This is still not ideal because the same person may appear in many different videos. Ideally we want a person to be either in train or in val, but not in both. But it will do for now.)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_splits(crops_dir, metadata_df, frac):\n    # Make a validation split. Sample a percentage of the real videos, \n    # and also grab the corresponding fake videos.\n    real_rows = metadata_df[metadata_df[\"label\"] == \"REAL\"]\n    real_df = real_rows.sample(frac=frac, random_state=666)\n    fake_df = metadata_df[metadata_df[\"original\"].isin(real_df[\"videoname\"])]\n    val_df = pd.concat([real_df, fake_df])\n\n    # The training split is the remaining videos.\n    train_df = metadata_df.loc[~metadata_df.index.isin(val_df.index)]\n\n    return train_df, val_df","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Sanity check:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df, val_df = make_splits(crops_dir, metadata_df, frac=0.05)\n\nassert(len(train_df) + len(val_df) == len(metadata_df))\nassert(len(train_df[train_df[\"videoname\"].isin(val_df[\"videoname\"])]) == 0)\n\ndel train_df, val_df","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Use all of the above building blocks to create `DataLoader` objects. Note that we use only a portion of the full amount of training data, for speed reasons. If you have more patience, increase the `sample_size`."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import DataLoader\n\ndef create_data_loaders(crops_dir, metadata_df, image_size, batch_size, num_workers):\n    train_df, val_df = make_splits(crops_dir, metadata_df, frac=0.05)\n\n    train_dataset = VideoDataset(crops_dir, train_df, \"train\", image_size, sample_size=10000)\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, \n                              num_workers=num_workers, pin_memory=True)\n\n    val_dataset = VideoDataset(crops_dir, val_df, \"val\", image_size, sample_size=500, seed=1234)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, \n                            num_workers=num_workers, pin_memory=True)\n\n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader, val_loader = create_data_loaders(crops_dir, metadata_df, image_size, \n                                               batch_size, num_workers=2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And, as usual, a check that it works... The `train_loader` should give a different set of examples each time you run it (because `shuffle=True`), while the `val_loader` always returns the examples in the same order."},{"metadata":{"trusted":true},"cell_type":"code","source":"X, y = next(iter(train_loader))\nplt.imshow(unnormalize_transform(X[0]).permute(1, 2, 0))\nprint(y[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"X, y = next(iter(val_loader))\nplt.imshow(unnormalize_transform(X[0]).permute(1, 2, 0))\nprint(y[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helper code for training"},{"metadata":{},"cell_type":"markdown","source":"Evaluation function for running the model on the validation set:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def evaluate(net, data_loader, device, silent=False):\n    net.train(False)\n\n    bce_loss = 0\n    total_examples = 0\n\n    with tqdm(total=len(data_loader), desc=\"Evaluation\", leave=False, disable=silent) as pbar:\n        for batch_idx, data in enumerate(data_loader):\n            with torch.no_grad():\n                batch_size = data[0].shape[0]\n                x = data[0].to(device)\n                y_true = data[1].to(device).float()\n\n                y_pred = net(x)\n                y_pred = y_pred.squeeze()\n\n                bce_loss += F.binary_cross_entropy_with_logits(y_pred, y_true).item() * batch_size\n\n            total_examples += batch_size\n            pbar.update()\n\n    bce_loss /= total_examples\n\n    if silent:\n        return bce_loss\n    else:\n        print(\"BCE: %.4f\" % (bce_loss))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Simple training loop. I prefer to write those myself from scratch each time, because then you can tweak it to do whatever you like."},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit(epochs):\n    global history, iteration, epochs_done, lr\n\n    with tqdm(total=len(train_loader), leave=False) as pbar:\n        for epoch in range(epochs):\n            pbar.reset()\n            pbar.set_description(\"Epoch %d\" % (epochs_done + 1))\n            \n            bce_loss = 0\n            total_examples = 0\n\n            net.train(True)\n\n            for batch_idx, data in enumerate(train_loader):\n                batch_size = data[0].shape[0]\n                x = data[0].to(gpu)\n                y_true = data[1].to(gpu).float()\n                \n                optimizer.zero_grad()\n\n                y_pred = net(x)\n                y_pred = y_pred.squeeze()\n                \n                loss = F.binary_cross_entropy_with_logits(y_pred, y_true)\n                loss.backward()\n                optimizer.step()\n                \n                batch_bce = loss.item()\n                bce_loss += batch_bce * batch_size\n                history[\"train_bce\"].append(batch_bce)\n\n                total_examples += batch_size\n                iteration += 1\n                pbar.update()\n\n            bce_loss /= total_examples\n            epochs_done += 1\n\n            print(\"Epoch: %3d, train BCE: %.4f\" % (epochs_done, bce_loss))\n\n            val_bce_loss = evaluate(net, val_loader, device=gpu, silent=True)\n            history[\"val_bce\"].append(val_bce_loss)\n            \n            print(\"              val BCE: %.4f\" % (val_bce_loss))\n\n            # TODO: can do LR annealing here\n            # TODO: can save checkpoint here\n\n            print(\"\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The model"},{"metadata":{},"cell_type":"markdown","source":"Need to load pretrained ImageNet weights into the model.\n\nYou can get these weights from `https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth`, or from [this dataset](https://www.kaggle.com/tony92151/pretrained-pytorch) by Kaggler [tonyguo](https://www.kaggle.com/tony92151)."},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint = torch.load(\"../input/pretrained-pytorch/resnext50_32x4d-7cdf4587.pth\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision.models as models\n\nclass MyResNeXt(models.resnet.ResNet):\n    def __init__(self, training=True):\n        super(MyResNeXt, self).__init__(block=models.resnet.Bottleneck,\n                                        layers=[3, 4, 6, 3], \n                                        groups=32, \n                                        width_per_group=4)\n\n        self.load_state_dict(checkpoint)\n\n        # Override the existing FC layer with a new one.\n        self.fc = nn.Linear(2048, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net = MyResNeXt().to(gpu)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del checkpoint","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Test the model on a small batch to see what its output shape is:"},{"metadata":{"trusted":true},"cell_type":"code","source":"out = net(torch.zeros((10, 3, image_size, image_size)).to(gpu))\nout.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Freeze the early layers of the model:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def freeze_until(net, param_name):\n    found_name = False\n    for name, params in net.named_parameters():\n        if name == param_name:\n            found_name = True\n        params.requires_grad = found_name","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"freeze_until(net, \"layer4.0.conv1.weight\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"These are the layers we will train:"},{"metadata":{"trusted":true},"cell_type":"code","source":"[k for k,v in net.named_parameters() if v.requires_grad]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Before we train, let's run the model on the validation set. This should give a logloss of about 0.6931."},{"metadata":{"trusted":true},"cell_type":"code","source":"evaluate(net, val_loader, device=gpu)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"lr = 0.01\nwd = 0.\n\nhistory = { \"train_bce\": [], \"val_bce\": [] }\niteration = 0\nepochs_done = 0\n\noptimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"At this point you can load the model from the previous checkpoint. If you do, also make sure to restore the optimizer state! Something like this:\n\n```python\ncheckpoint = torch.load(\"model-checkpoint.pth\")\nnet.load_state_dict(checkpoint)\n\ncheckpoint = torch.load(\"optimizer-checkpoint.pth\")\noptimizer.load_state_dict(checkpoint)\n```"},{"metadata":{},"cell_type":"markdown","source":"Let's start training!"},{"metadata":{"trusted":true},"cell_type":"code","source":"fit(5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Manual learning rate annealing:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_lr(optimizer, lr):\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lr /= 10\nset_lr(optimizer, lr)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fit(5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Plot training progress. It's nicer to use something like TensorBoard for this, but a simple plot also works. ;-)"},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(history[\"train_bce\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(history[\"val_bce\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**All done!** You can now use this checkpoint in the [inference kernel](https://www.kaggle.com/humananalog/inference-demo)."},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(net.state_dict(), \"checkpoint.pth\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}