{"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":"code","source":"import argparse\nimport copy\nimport os\nimport pathlib\n\nimport numpy as np\nimport sklearn.model_selection\nimport timm\nimport torch\nimport torchvision\nimport torchvision.transforms.functional\nfrom torchvision import transforms\nfrom tqdm.notebook import tqdm\n\nCLIP_THRESHOLD = 0.0125\n\n\n# コード引用あり＠5節\ndef setup_train_val_split(labels, dryrun=False, seed=0):\n    x = np.arange(len(labels))\n    y = np.array(labels)\n    splitter = sklearn.model_selection.StratifiedShuffleSplit(\n        n_splits=1, train_size=0.8, random_state=seed\n    )\n    train_indices, val_indices = next(splitter.split(x, y))\n\n    if dryrun:\n        train_indices = np.random.choice(train_indices, 100, replace=False)\n        val_indices = np.random.choice(val_indices, 100, replace=False)\n\n    return train_indices, val_indices\n\n\n# ↓コード引用あり@8節\ndef setup_cv_split(labels, n_folds, fold, dryrun=False, seed=0):\n    x = np.arange(len(labels))\n    y = np.array(labels)\n    splitter = sklearn.model_selection.StratifiedKFold(\n        n_splits=n_folds, shuffle=True, random_state=seed\n    )\n    splits = list(splitter.split(x, y))\n    train_indices, val_indices = splits[fold]\n\n    if dryrun:\n        train_indices = np.random.choice(train_indices, 100, replace=False)\n        val_indices = np.random.choice(val_indices, 100, replace=False)\n\n    return train_indices, val_indices\n\n\ndef get_labels(dataset):\n    if isinstance(dataset, torch.utils.data.Subset):\n        return get_labels(dataset.dataset)[dataset.indices]\n    else:\n        return np.array([img[1] for img in dataset.imgs])\n\n\n# コード引用あり＠5節\ndef setup_train_val_datasets(data_dir, dryrun=False):\n    dataset = torchvision.datasets.ImageFolder(\n        os.path.join(data_dir, \"train\"),\n        transform=setup_center_crop_transform(),\n    )\n    labels = get_labels(dataset)\n    train_indices, val_indices = setup_train_val_split(labels, dryrun)\n\n    train_dataset = torch.utils.data.Subset(dataset, train_indices)\n    val_dataset = torch.utils.data.Subset(dataset, val_indices)\n\n    return train_dataset, val_dataset\n\n\n# コード引用あり＠8節\ndef setup_cv_datasets(data_dir, n_folds, fold, dryrun=False):\n    dataset = torchvision.datasets.ImageFolder(\n        os.path.join(data_dir, \"train\"),\n        transform=setup_center_crop_transform(),\n    )\n    labels = get_labels(dataset)\n    train_indices, val_indices = setup_cv_split(labels, n_folds, fold, dryrun)\n\n    train_dataset = torch.utils.data.Subset(dataset, train_indices)\n    val_dataset = torch.utils.data.Subset(dataset, val_indices)\n\n    return train_dataset, val_dataset\n\n\ndef set_transform(dataset, transform):\n    if isinstance(dataset, torch.utils.data.Subset):\n        set_transform(dataset.dataset, transform)\n    else:\n        dataset.transform = transform\n\n\n# コード引用あり＠5節\ndef setup_train_val_loaders(data_dir, batch_size, dryrun=False):\n    train_dataset, val_dataset = setup_train_val_datasets(\n        data_dir, dryrun=dryrun\n    )\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n        num_workers=8,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=batch_size, num_workers=8\n    )\n    return train_loader, val_loader\n\n\n########################################################################################################################\n# transform\n########################################################################################################################\n\n# コード引用あり＠5節\ndef setup_center_crop_transform():\n    return transforms.Compose(\n        [\n            transforms.Resize(256),\n            transforms.CenterCrop(224),\n            transforms.ToTensor(),\n            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n        ]\n    )\n\n\n# コード引用あり＠7節\ndef setup_crop_flip_transform():\n    return transforms.Compose(\n        [\n            transforms.RandomResizedCrop(224),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n        ]\n    )\n\n\n# コード引用あり＠8-2\ndef setup_tta_transforms():\n    return [\n        transforms.Compose(\n            [\n                transforms.Resize(256),\n                transforms.CenterCrop(224),\n                transforms.ToTensor(),\n                transforms.Normalize(\n                    [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]\n                ),\n            ]\n        ),\n        transforms.Compose(\n            [\n                transforms.Resize(256),\n                transforms.CenterCrop(224),\n                torchvision.transforms.functional.hflip,\n                transforms.ToTensor(),\n                transforms.Normalize(\n                    [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]\n                ),\n            ]\n        ),\n    ]\n\n\n########################################################################################################################\n# train loop 5節（lr schedulerなし）\n########################################################################################################################\n\n# コード引用あり＠5節\ndef train_1epoch(model, train_loader, lossfun, optimizer, device):\n    model.train()\n    total_loss, total_acc = 0.0, 0.0\n\n    for x, y in tqdm(train_loader):\n        x = x.to(device)\n        y = y.to(device)\n\n        optimizer.zero_grad()\n        out = model(x)\n        loss = lossfun(out, y)\n        _, pred = torch.max(out.detach(), 1)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item() * x.size(0)\n        total_acc += torch.sum(pred == y)\n\n    avg_loss = total_loss / len(train_loader.dataset)\n    avg_acc = total_acc / len(train_loader.dataset)\n    return avg_acc, avg_loss\n\n\n# コード引用あり＠5節\ndef validate_1epoch(model, val_loader, lossfun, device):\n    model.eval()\n    total_loss, total_acc = 0.0, 0.0\n\n    with torch.no_grad():\n        for x, y in tqdm(val_loader):\n            x = x.to(device)\n            y = y.to(device)\n\n            out = model(x)\n            loss = lossfun(out.detach(), y)\n            _, pred = torch.max(out, 1)\n\n            total_loss += loss.item() * x.size(0)\n            total_acc += torch.sum(pred == y)\n\n    avg_loss = total_loss / len(val_loader.dataset)\n    avg_acc = total_acc / len(val_loader.dataset)\n    return avg_acc, avg_loss\n\n\n# コード引用あり＠5節\ndef train(model, optimizer, train_loader, val_loader, n_epochs, device):\n    lossfun = torch.nn.CrossEntropyLoss()\n\n    for epoch in tqdm(range(n_epochs)):\n        train_acc, train_loss = train_1epoch(\n            model, train_loader, lossfun, optimizer, device\n        )\n        val_acc, val_loss = validate_1epoch(model, val_loader, lossfun, device)\n        print(\n            f\"epoch={epoch}, train loss={train_loss}, train accuracy={train_acc}, val loss={val_loss}, val accuracy={val_acc}\"\n        )\n\n\n########################################################################################################################\n# train loop 6節以降（lr schedulerあり）\n########################################################################################################################\n\n# コード引用あり＠5節\ndef train_1epoch2(\n    model, train_loader, lossfun, optimizer, lr_scheduler, device\n):\n    model.train()\n    total_loss, total_acc = 0.0, 0.0\n\n    for x, y in tqdm(train_loader):\n        x = x.to(device)\n        y = y.to(device)\n\n        optimizer.zero_grad()\n        out = model(x)\n        _, pred = torch.max(out.detach(), 1)\n        loss = lossfun(model(x), y)\n        loss.backward()\n        optimizer.step()\n        lr_scheduler.step()\n\n        total_loss += loss.item() * x.size(0)\n        total_acc += torch.sum(pred == y)\n\n    avg_loss = total_loss / len(train_loader.dataset)\n    avg_acc = total_acc / len(train_loader.dataset)\n    return avg_acc, avg_loss\n\n\ndef train2(\n    model, optimizer, lr_scheduler, train_loader, val_loader, n_epochs, device\n):\n    lossfun = torch.nn.CrossEntropyLoss()\n\n    for epoch in tqdm(range(n_epochs)):\n        train_acc, train_loss = train_1epoch2(\n            model, train_loader, lossfun, optimizer, lr_scheduler, device\n        )\n        val_acc, val_loss = validate_1epoch(model, val_loader, lossfun, device)\n\n        lr = optimizer.param_groups[0][\"lr\"]\n        print(\n            f\"epoch={epoch}, train loss={train_loss}, train accuracy={train_acc}, val loss={val_loss}, val accuracy={val_acc}, lr={lr}\"\n        )\n\n\n########################################################################################################################\n# mixup\n########################################################################################################################\n\n\n# コード引用あり@7-3節\ndef train_1epoch_mixup(\n    model, train_loader, lossfun, optimizer, lr_scheduler, mixup_alpha, device\n):\n    model.train()\n    total_loss, total_acc = 0.0, 0.0\n\n    for x, y in tqdm(train_loader):\n        x = x.to(device)\n        y = y.to(device)\n\n        lmd = np.random.beta(mixup_alpha, mixup_alpha)\n        perm = torch.randperm(x.shape[0]).to(device)\n        x2 = x[perm, :]\n        y2 = y[perm]\n\n        optimizer.zero_grad()\n        out = model(lmd * x + (1.0 - lmd) * x2)\n        loss = lmd * lossfun(out, y) + (1.0 - lmd) * lossfun(out, y2)\n        loss.backward()\n        optimizer.step()\n        lr_scheduler.step()\n\n        _, pred = torch.max(out.detach(), 1)\n        total_loss += loss.item() * x.size(0)\n        total_acc += lmd * torch.sum(pred == y) + (1.0 - lmd) * torch.sum(\n            pred == y2\n        )\n\n    avg_loss = total_loss / len(train_loader.dataset)\n    avg_acc = total_acc / len(train_loader.dataset)\n    return avg_acc, avg_loss\n\n\ndef train3_mixup(\n    model,\n    optimizer,\n    lr_scheduler,\n    train_loader,\n    val_loader,\n    n_epochs,\n    n_mixup_epochs,\n    mixup_alpha,\n    device,\n):\n    lossfun = torch.nn.CrossEntropyLoss()\n    last_val_loss = 0.0\n\n    for epoch in tqdm(range(n_epochs)):\n        if epoch < n_mixup_epochs:\n            train_acc, train_loss = train_1epoch_mixup(\n                model,\n                train_loader,\n                lossfun,\n                optimizer,\n                lr_scheduler,\n                mixup_alpha,\n                device,\n            )\n        else:\n            train_acc, train_loss = train_1epoch2(\n                model, train_loader, lossfun, optimizer, lr_scheduler, device\n            )\n\n        val_acc, val_loss = validate_1epoch(model, val_loader, lossfun, device)\n        last_val_loss = val_loss\n\n        lr = optimizer.param_groups[0][\"lr\"]\n        print(\n            f\"epoch={epoch}, train loss={train_loss}, train accuracy={train_acc}, val loss={val_loss}, val accuracy={val_acc}, lr={lr}\"\n        )\n\n    return last_val_loss\n\n\n########################################################################################################################\n# predict部分\n########################################################################################################################\n\n# コード引用あり＠5節\ndef setup_test_loader(data_dir, batch_size, dryrun):\n    dataset = torchvision.datasets.ImageFolder(\n        os.path.join(data_dir, \"test\"), transform=setup_center_crop_transform()\n    )\n    image_ids = [\n        os.path.splitext(os.path.basename(path))[0] for path, _ in dataset.imgs\n    ]\n\n    if dryrun:\n        dataset = torch.utils.data.Subset(dataset, range(0, 100))\n        image_ids = image_ids[:100]\n\n    loader = torch.utils.data.DataLoader(\n        dataset, batch_size=batch_size, num_workers=8\n    )\n    return loader, image_ids\n\n\n# コード引用あり＠5節\ndef predict(model, loader, device):\n    pred_fun = torch.nn.Softmax(dim=1)\n    preds = []\n    for x, _ in tqdm(loader):\n        with torch.set_grad_enabled(False):\n            x = x.to(device)\n            y = pred_fun(model(x))\n        y = y.cpu().numpy()\n        y = y[:, 1]  # cat:0, dog: 1\n        preds.append(y)\n    preds = np.concatenate(preds)\n    return preds\n\n\n# コード引用あり＠5節\ndef write_prediction(image_ids, prediction, out_path):\n    with open(out_path, \"w\") as f:\n        f.write(\"id,label\\n\")\n        for i, p in zip(image_ids, prediction):\n            f.write(\"{},{}\\n\".format(i, p))\n\n\n# コード引用あり＠5節\ndef write_prediction_with_clip(\n    image_ids, prediction, clip_threshold, out_path\n):\n    with open(out_path, \"w\") as f:\n        f.write(\"id,label\\n\")\n        for i, p in zip(image_ids, prediction):\n            p = np.clip(p, clip_threshold, 1.0 - clip_threshold)\n            f.write(\"{},{}\\n\".format(i, p))\n\n\n########################################################################################################################\n# 各種実行設定\n########################################################################################################################\n\n#\n# 5: First try\n#\n\n# コード引用あり＠5節\ndef train_subsec5(data_dir, batch_size, dryrun=False, device=\"cuda:0\"):\n    model = torchvision.models.resnet50(pretrained=True)\n    model.fc = torch.nn.Linear(model.fc.in_features, 2)\n    model.to(device)\n\n    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)\n    train_loader, val_loader = setup_train_val_loaders(\n        data_dir, batch_size, dryrun\n    )\n    train(\n        model, optimizer, train_loader, val_loader, n_epochs=1, device=device\n    )\n\n    return model\n\n\n# コード引用あり＠5節\ndef predict_subsec5(\n    data_dir, out_dir, model, batch_size, dryrun=False, device=\"cuda:0\"\n):\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    preds = predict(model, test_loader, device)\n    write_prediction(image_ids, preds, out_dir / \"out.csv\")\n\n\ndef run_5(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    model = train_subsec5(data_dir, batch_size, dryrun, device)\n\n    # clip無しの推論\n    predict_subsec5(data_dir, out_dir, model, batch_size, dryrun, device)\n\n    # clip有りの推論\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    preds = predict(model, test_loader, device)\n    write_prediction_with_clip(\n        image_ids, preds, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n#\n# 6: Weight decay, cosine annealing\n#\n\n\ndef run_6(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n\n    model = torchvision.models.resnet50(pretrained=True)\n    model.fc = torch.nn.Linear(model.fc.in_features, 2)\n    model.to(device)\n\n    train_dataset, val_dataset = setup_train_val_datasets(\n        data_dir, dryrun=dryrun\n    )\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n        num_workers=8,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=batch_size, num_workers=8\n    )\n\n    optimizer = torch.optim.SGD(\n        model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n    )\n    n_iterations = len(train_loader) * n_epochs\n    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, n_iterations\n    )\n\n    train2(\n        model,\n        optimizer,\n        lr_scheduler,\n        train_loader,\n        val_loader,\n        n_epochs=n_epochs,\n        device=device,\n    )\n\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    preds = predict(model, test_loader, device)\n    write_prediction_with_clip(\n        image_ids, preds, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n#\n# 7-1: Random crop&flip\n#\n\n\ndef run_7_1(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n\n    model = torchvision.models.resnet50(pretrained=True)\n    model.fc = torch.nn.Linear(model.fc.in_features, 2)\n    model.to(device)\n\n    train_dataset, val_dataset = setup_train_val_datasets(\n        data_dir, dryrun=dryrun\n    )\n    train_dataset = copy.deepcopy(\n        train_dataset\n    )  # transformを設定した際にval_datasetに影響したくない\n    train_transform = setup_crop_flip_transform()\n    set_transform(train_dataset, train_transform)\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n        num_workers=8,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=batch_size, num_workers=8\n    )\n\n    optimizer = torch.optim.SGD(\n        model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n    )\n    n_iterations = len(train_loader) * n_epochs\n    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, n_iterations\n    )\n\n    train2(\n        model,\n        optimizer,\n        lr_scheduler,\n        train_loader,\n        val_loader,\n        n_epochs=n_epochs,\n        device=device,\n    )\n\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    preds = predict(model, test_loader, device)\n    write_prediction_with_clip(\n        image_ids, preds, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n#\n# 7-3: Mixup\n#\n\n\ndef run_7_3(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n    n_mixup_epochs = 1 if dryrun else 7\n    mixup_alpha = 0.4\n\n    model = torchvision.models.resnet50(pretrained=True)\n    model.fc = torch.nn.Linear(model.fc.in_features, 2)\n    model.to(device)\n\n    train_dataset, val_dataset = setup_train_val_datasets(\n        data_dir, dryrun=dryrun\n    )\n    train_dataset = copy.deepcopy(\n        train_dataset\n    )  # transformを設定した際にval_datasetに影響したくない\n    train_transform = setup_crop_flip_transform()\n    set_transform(train_dataset, train_transform)\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n        num_workers=8,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=batch_size, num_workers=8\n    )\n\n    optimizer = torch.optim.SGD(\n        model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n    )\n    n_iterations = len(train_loader) * n_epochs\n    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, n_iterations\n    )\n\n    train3_mixup(\n        model,\n        optimizer,\n        lr_scheduler,\n        train_loader,\n        val_loader,\n        n_epochs=n_epochs,\n        n_mixup_epochs=n_mixup_epochs,\n        mixup_alpha=mixup_alpha,\n        device=device,\n    )\n\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    preds = predict(model, test_loader, device)\n    write_prediction_with_clip(\n        image_ids, preds, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n#\n# 8-1: Averaging\n#\n\n# コード引用はしないけど引用されてるコードから呼ばれてる@8-1\ndef train_predict_1fold(\n    data_dir,\n    fold,\n    n_folds,\n    n_epochs,\n    n_mixup_epochs,\n    mixup_alpha,\n    test_loader,\n    batch_size,\n    dryrun=False,\n    device=\"cuda:0\",\n):\n    print(f\"Fold {fold}\")\n    model = torchvision.models.resnet50(pretrained=True)\n    model.fc = torch.nn.Linear(model.fc.in_features, 2)\n    model.to(device)\n    train_dataset, val_dataset = setup_cv_datasets(\n        data_dir, n_folds, fold, dryrun=dryrun\n    )\n    train_dataset = copy.deepcopy(train_dataset)\n    train_transform = setup_crop_flip_transform()\n    set_transform(train_dataset, train_transform)\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n        num_workers=8,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=batch_size, num_workers=8\n    )\n    optimizer = torch.optim.SGD(\n        model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n    )\n    n_iterations = len(train_loader) * n_epochs\n    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, n_iterations\n    )\n\n    val_loss = train3_mixup(\n        model,\n        optimizer,\n        lr_scheduler,\n        train_loader,\n        val_loader,\n        n_epochs=n_epochs,\n        n_mixup_epochs=n_mixup_epochs,\n        mixup_alpha=mixup_alpha,\n        device=device,\n    )\n\n    test_pred = predict(model, test_loader, device)\n    return val_loss, test_pred\n\n\n# コード引用あり@8-1\ndef train_predict_subsec81(\n    data_dir,\n    out_dir,\n    n_folds,\n    n_epochs,\n    n_mixup_epochs,\n    mixup_alpha,\n    batch_size,\n    dryrun=False,\n    device=\"cuda:0\",\n):\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n\n    val_losses = []\n    test_preds = []\n    for fold in range(n_folds):\n        val_loss, test_pred = train_predict_1fold(\n            data_dir,\n            fold,\n            n_folds,\n            n_epochs,\n            n_mixup_epochs,\n            mixup_alpha,\n            test_loader,\n            batch_size,\n            dryrun,\n            device,\n        )\n        val_losses.append(val_loss)\n        test_preds.append(test_pred)\n\n    # 全foldでのvalidation scoreの平均を最終的なvalidation scoreとして出力\n    val_loss = np.mean(val_losses)\n    print(f\"val loss={val_loss}\")\n\n    # 全foldでのモデルのtest dataに対する予測結果を平均したものを最終的な予測として出力\n    test_pred = np.mean(test_preds, axis=0)\n    write_prediction_with_clip(\n        image_ids, test_pred, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\ndef run_8_1(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n    n_mixup_epochs = 1 if dryrun else 7\n    n_folds = 2 if dryrun else 5\n    mixup_alpha = 0.4\n\n    train_predict_subsec81(\n        data_dir,\n        out_dir,\n        n_folds,\n        n_epochs,\n        n_mixup_epochs,\n        mixup_alpha,\n        batch_size,\n        dryrun,\n        device,\n    )\n\n\n#\n# 8-2: TTA\n#\n\n\n# コード引用あり＠8−2\ndef predict_tta(model, loader, device):\n    tta_transforms = setup_tta_transforms()\n    preds = []\n    for transform in tta_transforms:\n        set_transform(loader.dataset, transform)\n        val_pred = predict(model, loader, device)\n        preds.append(val_pred)\n\n    pred = np.mean(preds, axis=0)\n    return pred\n\n\n# TTA\ndef run_8_2(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n    n_mixup_epochs = 1 if dryrun else 7\n    n_folds = 2 if dryrun else 5\n    mixup_alpha = 0.4\n\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    val_losses = []\n    test_preds = []\n\n    for fold in range(n_folds):\n        print(f\"Fold {fold}\")\n        model = torchvision.models.resnet50(pretrained=True)\n        model.fc = torch.nn.Linear(model.fc.in_features, 2)\n        model.to(device)\n\n        train_dataset, val_dataset = setup_cv_datasets(\n            data_dir, n_folds, fold, dryrun=dryrun\n        )\n        train_dataset = copy.deepcopy(train_dataset)\n        train_transform = setup_crop_flip_transform()\n        set_transform(train_dataset, train_transform)\n        train_loader = torch.utils.data.DataLoader(\n            train_dataset,\n            batch_size=batch_size,\n            shuffle=True,\n            drop_last=True,\n            num_workers=8,\n        )\n        val_loader = torch.utils.data.DataLoader(\n            val_dataset, batch_size=batch_size, num_workers=8\n        )\n\n        optimizer = torch.optim.SGD(\n            model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n        )\n        n_iterations = len(train_loader) * n_epochs\n        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            optimizer, n_iterations\n        )\n\n        train3_mixup(\n            model,\n            optimizer,\n            lr_scheduler,\n            train_loader,\n            val_loader,\n            n_epochs=n_epochs,\n            n_mixup_epochs=n_mixup_epochs,\n            mixup_alpha=mixup_alpha,\n            device=device,\n        )\n\n        val_pred = predict_tta(model, val_loader, device)\n        val_labels = get_labels(val_loader.dataset)\n        val_loss = sklearn.metrics.log_loss(val_labels, val_pred)\n        val_losses.append(val_loss)\n        print(f\"fold={fold}, val loss with TTA={val_loss}\")\n\n        test_pred = predict_tta(model, test_loader, device)\n        test_preds.append(test_pred)\n\n    val_loss = np.mean(val_losses)\n    print(f\"val loss={val_loss}\")\n\n    test_pred = np.mean(test_preds, axis=0)\n    write_prediction_with_clip(\n        image_ids, test_pred, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n# TTA\ndef run_9_1(data_dir, out_dir, dryrun, device):\n    batch_size = 32\n    n_epochs = 2 if dryrun else 10\n    n_mixup_epochs = 1 if dryrun else 7\n    n_folds = 2 if dryrun else 5\n    mixup_alpha = 0.4\n\n    test_loader, image_ids = setup_test_loader(\n        data_dir, batch_size, dryrun=dryrun\n    )\n    val_losses = []\n    test_preds = []\n\n    for fold in range(n_folds):\n        print(f\"Fold {fold}\")\n        model = timm.create_model(\n            \"regnety_080\", pretrained=True, num_classes=2\n        )\n        model.to(device)\n\n        train_dataset, val_dataset = setup_cv_datasets(\n            data_dir, n_folds, fold, dryrun=dryrun\n        )\n        train_dataset = copy.deepcopy(train_dataset)\n        train_transform = setup_crop_flip_transform()\n        set_transform(train_dataset, train_transform)\n        train_loader = torch.utils.data.DataLoader(\n            train_dataset,\n            batch_size=batch_size,\n            shuffle=True,\n            drop_last=True,\n            num_workers=8,\n        )\n        val_loader = torch.utils.data.DataLoader(\n            val_dataset, batch_size=batch_size, num_workers=8\n        )\n\n        optimizer = torch.optim.SGD(\n            model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001\n        )\n        n_iterations = len(train_loader) * n_epochs\n        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            optimizer, n_iterations\n        )\n\n        train3_mixup(\n            model,\n            optimizer,\n            lr_scheduler,\n            train_loader,\n            val_loader,\n            n_epochs=n_epochs,\n            n_mixup_epochs=n_mixup_epochs,\n            mixup_alpha=mixup_alpha,\n            device=device,\n        )\n\n        val_pred = predict_tta(model, val_loader, device)\n        val_labels = get_labels(val_loader.dataset)\n        val_loss = sklearn.metrics.log_loss(val_labels, val_pred)\n        val_losses.append(val_loss)\n        print(f\"fold={fold}, val loss with TTA={val_loss}\")\n\n        test_pred = predict_tta(model, test_loader, device)\n        test_preds.append(test_pred)\n\n    val_loss = np.mean(val_losses)\n    print(f\"val loss={val_loss}\")\n\n    test_pred = np.mean(test_preds, axis=0)\n    write_prediction_with_clip(\n        image_ids, test_pred, CLIP_THRESHOLD, out_dir / \"out_clip.csv\"\n    )\n\n\n# def main():\n#     parser = argparse.ArgumentParser()\n#     parser.add_argument(\"--config\", required=True)\n#     parser.add_argument(\"--data_dir\", required=True)\n#     parser.add_argument(\"--out_dir\", default=\"./out\")\n#     parser.add_argument(\"--device\", default=\"cuda:0\")\n#     parser.add_argument(\"--dryrun\", action=\"store_true\")\n#     args = parser.parse_args()\n\n#     data_dir = pathlib.Path(args.data_dir)\n#     out_dir = pathlib.Path(args.out_dir)\n#     device = args.device\n#     dryrun = args.dryrun\n\n#     out_dir.mkdir(parents=True, exist_ok=True)\n#     if args.config == \"5\":\n#         run_5(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"6\":\n#         run_6(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"7-1\":\n#         run_7_1(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"7-3\":\n#         run_7_3(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"8-1\":\n#         run_8_1(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"8-2\":\n#         run_8_2(data_dir, out_dir, dryrun, device)\n#     elif args.config == \"9-1\":\n#         run_9_1(data_dir, out_dir, dryrun, device)\n#     else:\n#         raise ValueError(f\"Unknown config: {args.config}\")\n\n\n# if __name__ == \"__main__\":\n#     main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-26T12:04:01.269236Z","iopub.execute_input":"2023-06-26T12:04:01.269622Z","iopub.status.idle":"2023-06-26T12:04:06.179270Z","shell.execute_reply.started":"2023-06-26T12:04:01.269587Z","shell.execute_reply":"2023-06-26T12:04:06.178326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_val_score(stdout_path):\n    with open(stdout_path) as f:\n        lines = f.readlines()\n\n    last_line = lines[-1]\n\n    if last_line.startswith(\"val loss=\"):\n        return {\n            \"val method\": \"CV\",\n            \"val score\": float(last_line.split(\"=\")[1]),\n        }\n    else:\n        tokens = last_line.split(\",\")\n        tokens = [token.strip() for token in tokens]\n        tokens = [token for token in tokens if token.startswith(\"val loss=\")]\n        assert len(tokens) == 1\n        token = tokens[0]\n        return {\"val method\": \"split\", \"val score\": float(token.split(\"=\")[1])}","metadata":{"execution":{"iopub.status.busy":"2023-06-26T12:04:06.181310Z","iopub.execute_input":"2023-06-26T12:04:06.181710Z","iopub.status.idle":"2023-06-26T12:04:06.189645Z","shell.execute_reply.started":"2023-06-26T12:04:06.181673Z","shell.execute_reply":"2023-06-26T12:04:06.187968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!unzip -q -n ../input/dogs-vs-cats-redux-kernels-edition/train.zip -d dogs_vs_cats\n!unzip -q -n ../input/dogs-vs-cats-redux-kernels-edition/test.zip -d dogs_vs_cats\n!mkdir -p dogs_vs_cats/train/cat\n!mkdir -p dogs_vs_cats/train/dog\n!mkdir -p dogs_vs_cats/test/unknown\n!mv dogs_vs_cats/train/cat*.jpg dogs_vs_cats/train/cat\n!mv dogs_vs_cats/train/dog*.jpg dogs_vs_cats/train/dog\n!mv dogs_vs_cats/test/*.jpg dogs_vs_cats/test/unknown","metadata":{"execution":{"iopub.status.busy":"2023-06-26T12:04:06.191069Z","iopub.execute_input":"2023-06-26T12:04:06.191630Z","iopub.status.idle":"2023-06-26T12:04:35.345400Z","shell.execute_reply.started":"2023-06-26T12:04:06.191598Z","shell.execute_reply":"2023-06-26T12:04:35.344098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from contextlib import redirect_stdout\nimport io\nimport json\nimport pathlib\nimport sys","metadata":{"execution":{"iopub.status.busy":"2023-06-26T12:04:35.349232Z","iopub.execute_input":"2023-06-26T12:04:35.349591Z","iopub.status.idle":"2023-06-26T12:04:35.354411Z","shell.execute_reply.started":"2023-06-26T12:04:35.349558Z","shell.execute_reply":"2023-06-26T12:04:35.353195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_dir = \"./dogs_vs_cats/\"\ndryrun = False\ndevice = \"cuda:0\"\nroot_out_dir = pathlib.Path(\"./out\")\nos.makedirs(\"./out\", exist_ok=True)\n\nresults = []","metadata":{"execution":{"iopub.status.busy":"2023-06-26T12:04:35.355927Z","iopub.execute_input":"2023-06-26T12:04:35.356561Z","iopub.status.idle":"2023-06-26T12:04:35.366396Z","shell.execute_reply.started":"2023-06-26T12:04:35.356527Z","shell.execute_reply":"2023-06-26T12:04:35.365436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"5\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_5(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-25T23:59:08.296485Z","iopub.execute_input":"2023-06-25T23:59:08.296836Z","iopub.status.idle":"2023-06-26T00:03:22.503544Z","shell.execute_reply.started":"2023-06-25T23:59:08.296804Z","shell.execute_reply":"2023-06-26T00:03:22.502329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"6\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_6(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T00:03:22.505668Z","iopub.execute_input":"2023-06-26T00:03:22.507180Z","iopub.status.idle":"2023-06-26T00:32:08.351973Z","shell.execute_reply.started":"2023-06-26T00:03:22.507138Z","shell.execute_reply":"2023-06-26T00:32:08.350483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"7-1\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_7_1(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T00:32:08.354190Z","iopub.execute_input":"2023-06-26T00:32:08.354601Z","iopub.status.idle":"2023-06-26T01:00:47.509268Z","shell.execute_reply.started":"2023-06-26T00:32:08.354560Z","shell.execute_reply":"2023-06-26T01:00:47.507887Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"7-3\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_7_3(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T01:00:47.511025Z","iopub.execute_input":"2023-06-26T01:00:47.511765Z","iopub.status.idle":"2023-06-26T01:25:55.699471Z","shell.execute_reply.started":"2023-06-26T01:00:47.511722Z","shell.execute_reply":"2023-06-26T01:25:55.698060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"8-1\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_8_1(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T01:25:55.701603Z","iopub.execute_input":"2023-06-26T01:25:55.702025Z","iopub.status.idle":"2023-06-26T03:32:28.014212Z","shell.execute_reply.started":"2023-06-26T01:25:55.701983Z","shell.execute_reply":"2023-06-26T03:32:28.012736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"8-2\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_8_2(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T03:32:28.016573Z","iopub.execute_input":"2023-06-26T03:32:28.017065Z","iopub.status.idle":"2023-06-26T05:47:54.158443Z","shell.execute_reply.started":"2023-06-26T03:32:28.016999Z","shell.execute_reply":"2023-06-26T05:47:54.157178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = \"9-1\"\nout_dir = root_out_dir / f\"config{config}\"\nout_dir.mkdir(parents=True, exist_ok=True)\nstdout_path = out_dir / \"stdout.txt\"\nwith open(stdout_path, 'w') as f:\n    with redirect_stdout(f):\n        run_9_1(data_dir, out_dir, dryrun, device)\n\nval_report = read_val_score(stdout_path)\nval_report[\"config\"] = config\nresults.append(val_report)\nprint(val_report)\nwith open(out_dir / \"report.json\", \"w\") as f:\n    json.dump(val_report, f)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T12:04:53.485854Z","iopub.execute_input":"2023-06-26T12:04:53.486257Z","iopub.status.idle":"2023-06-26T19:04:10.831987Z","shell.execute_reply.started":"2023-06-26T12:04:53.486227Z","shell.execute_reply":"2023-06-26T19:04:10.830783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}