{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Import model code and read metadata"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport json\nimport glob\nimport random\nimport logging\nimport numpy as np\nfrom tqdm import tqdm\n\nimport glob\nimport torch\nimport torch.utils.data as data\nfrom torch import nn\nfrom shutil import copyfile\n\n\ncopyfile(src = \"../input/resnetpy/resnet.py\", dst = \"../working/resnet.py\")\nfrom resnet import *\n        \nmetafiles = sorted(glob.glob('/kaggle/input/dfdc-video-faces/part*/*/*.json'))\ntrain_metafiles = metafiles[2:]\neval_metafiles = metafiles[:2]\nprint(train_metafiles)\nprint(eval_metafiles)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataloading"},{"metadata":{"trusted":true},"cell_type":"code","source":"class VideoDataset(data.Dataset):\n\n    def __init__(self, dataset_type, meta_files, batch_size):\n        self.dataset_type = dataset_type\n        self.batch_size = batch_size\n        self.clip_shape = (16, 112, 112)\n        self.load_meta(meta_files)\n\n        print('%s dataset: %d real and %d fake samples' % (\n              dataset_type, self.n_reals, self.n_fakes))\n\n    def load_clip(self, filename, target):\n        depth, height, width = self.clip_shape\n        clip = np.zeros((3, depth, height, width), dtype=np.uint8)\n        reader = cv2.VideoCapture(filename)\n        if not reader.isOpened():\n            logging.warn('could not open %s' % filename)\n            return torch.from_numpy(clip).float()\n\n        # If training, use the same cropping parameters for an entire set\n        # of video clips\n        if target == 0 or self.dataset_type == 'eval':\n            nframes = int(reader.get(cv2.CAP_PROP_FRAME_COUNT))\n            frame_height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n            frame_width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n            self.start_frame = random.randint(0, nframes - self.clip_shape[0])\n            self.start_row = random.randint(0, frame_height - height)\n            self.start_col = random.randint(0, frame_width - width)\n\n        for i in range(self.start_frame):\n            reader.grab()\n\n        for i in range(depth):\n            reader.grab()\n            success, frame = reader.retrieve()\n            if not success:\n                logging.warn('could not load frame %d in %s' % (\n                             start_frame + i, filename))\n                break\n            frame = frame[self.start_row:self.start_row + height,\n                          self.start_col:self.start_col + width]\n            clip[:, i] = frame.transpose((2, 0, 1))\n\n        reader.release()\n        return torch.from_numpy(clip).float()\n\n    def load_meta(self, meta_files):\n        meta = []\n        for meta_file in meta_files:\n            dirname = os.path.dirname(meta_file)\n            with open(meta_file) as meta_fd:\n                meta_dict = json.load(meta_fd)\n                new_dict = {}\n                # Expand filenames to their paths\n                for real in meta_dict:\n                    fakes = meta_dict[real]\n                    fakes = [os.path.join(dirname, fake) for fake in fakes]\n                    new_dict[os.path.join(dirname, real)] = fakes\n                meta_list = list(new_dict.items())\n                meta.extend(meta_list)\n\n        random.shuffle(meta)\n        self.clips = []\n        self.targets = []\n        for item in meta:\n            real = item[0]\n            fakes = item[1]\n            random.shuffle(fakes)\n            for fake in fakes:\n                # Oversample from non-fake videos\n                self.clips.append(real)\n                self.targets.append(0)\n                self.clips.append(fake)\n                self.targets.append(1)\n                # Use a small subset for evaluation\n                if self.dataset_type == 'eval':\n                    break\n\n        if self.dataset_type == 'eval':\n            # Make 4 copies to get random crops from\n            for _ in range(2):\n                self.clips.extend(self.clips)\n                self.targets.extend(self.targets)\n\n        self.len = len(self.clips)\n        self.n_fakes = np.sum(self.targets)\n        self.n_reals = self.len -  self.n_fakes\n\n    def __getitem__(self, index):\n        filename = self.clips[index]\n        target = self.targets[index]\n        clip = self.load_clip(filename, target)\n\n        return clip, target\n\n    def __len__(self):\n        return self.len\n\n\nepochs = 8\nbatch_size = 16\n\ntrain_dataset = VideoDataset('train', train_metafiles, batch_size)\n\ntrain_loader = torch.utils.data.DataLoader(\n    train_dataset, batch_size=batch_size, shuffle=False,\n    num_workers=1, pin_memory=True, sampler=None)\n\neval_dataset = VideoDataset('eval', eval_metafiles, batch_size)\n\neval_loader = torch.utils.data.DataLoader(\n    eval_dataset, batch_size=batch_size, shuffle=False,\n    num_workers=1, pin_memory=True, sampler=None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create model"},{"metadata":{"trusted":true},"cell_type":"code","source":"device = 'cuda'\nmodel = resnet3d18(num_classes=2)\nmodel = model.to(device)\n\ncrit = nn.CrossEntropyLoss().to(device)\nopt = torch.optim.Adam(model.parameters())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Train and validate"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(loader, model, crit, optimizer, epoch):\n    model.train()\n\n    loss_sum = 0\n    for clips, targets in tqdm(loader):\n        clips = clips.to(device)\n        targets = targets.to(device)\n\n        logits = model(clips)\n        loss = crit(logits, targets)\n        loss_sum += loss.data.cpu().numpy()\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n    return (loss_sum / len(loader))\n\ndef bce(probs, labels):\n    safelog =  lambda x: np.log(np.maximum(x, np.exp(-50.)))\n    return np.mean(-labels * safelog(probs) - (1 - labels) * safelog(1 - probs))\n\ndef validate(loader, model, crit):\n    model.eval()\n    sm = nn.Softmax(dim=1)\n    labels = np.zeros((len(loader.dataset)), dtype=np.float32)\n    probs = np.zeros((len(loader.dataset), 2), dtype=np.float32)\n    with torch.no_grad():\n        for i, (clips, targets) in enumerate(tqdm(loader)):\n            start = i*batch_size\n            end = start + clips.shape[0]\n            labels[start:end] = targets\n            clips = clips.to(device)\n\n            logits = model(clips)\n            probs[start:end] = sm(logits).cpu().numpy()\n\n    probs = probs.reshape(4, -1, 2).mean(axis=0)\n    labels = labels.reshape(4, -1).mean(axis=0)\n\n    preds = probs.argmax(axis=1)\n    correct = (preds == labels).sum()\n    acc = correct*100//preds.shape[0]\n    loss = bce(probs[:, 1], labels)\n    print('validation accuracy %d%%' % acc)\n    return loss\n\nmodel_file = 'model.pth'\nif os.path.exists(model_file):\n    checkpoint = torch.load(model_file)\n    model.load_state_dict(checkpoint['state_dict'])\n    print('loaded %s' % model_file)\n\ntry:\n    for epoch in range(epochs):\n        # Train for one epoch\n        train_loss = train(train_loader, model, crit, opt, epoch)\n\n        # Evaluate on validation set\n        val_loss = validate(eval_loader, model, crit)\n        print('epoch %d training loss %.2f validation loss %.2f\\n' % (\n              epoch, train_loss, val_loss))\nfinally:\n    torch.save({'state_dict': model.state_dict()}, model_file)\nprint('done')","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}