[
  {
    "path": ".gitignore",
    "content": "# File extensions\n*.out\n*.mp4\n\n# Directories\n/__pychache__\n/moco/__pychache__\n\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2021 Emin Orhan\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "# Self-supervised learning through the eyes of a child\n\nThis repository contains code for reproducing the results reported in the following paper:\n\nOrhan AE, Gupta VV, Lake BM (2020) [Self-supervised learning through the eyes of a child.](https://arxiv.org/abs/2007.16189) *Advances in Neural Information Processing Systems 34 (NeurIPS 2020)*.\n\n## Requirements\n\n* pytorch == 1.5.1\n* torchvision == 0.6.1\n\nSlightly older or newer versions will probably work fine as well.\n\n## Datasets\n\nThis project uses the SAYCam dataset described in the following paper: \n\nSullivan J, Mei M, Perfors A, Wojcik EH, Frank MC (2020) [SAYCam: A large, longitudinal audiovisual dataset recorded from the infant’s perspective.](https://psyarxiv.com/fy8zx/) PsyArXiv.\n\nThe dataset is hosted on the [Databrary](https://nyu.databrary.org/) repository for behavioral science. Unfortunately, we are unable to publicly share the SAYCam dataset here due to the terms of use. However, interested researchers can apply for access to the dataset with approval from their institution's IRB. \n\nIn addition, this project also uses the Toybox dataset for evaluation purposes. The Toybox dataset is publicly available at [this address](https://aivaslab.github.io/toybox/).\n\n## Code description\n\n* [`temporal_classification.py`](https://github.com/eminorhan/baby-vision/blob/master/temporal_classification.py): trains temporal classification models as described in the paper. This file uses code recycled from the PyTorch ImageNet training [example](https://github.com/pytorch/examples/tree/master/imagenet).\n* [`read_saycam.py`](https://github.com/eminorhan/baby-vision/blob/master/read_saycam.py): SAYCam video-to-image reader.\n* [`moco`](https://github.com/eminorhan/baby-vision/tree/master/moco) directory contains helper files for training static and temporal MoCo models. The code here was modified from [Facebook's MoCo repository](https://github.com/facebookresearch/moco).\n* [`moco_img.py`](https://github.com/eminorhan/baby-vision/blob/master/moco_img.py): trains an image-based MoCo model as described in the paper. This code was modified from [Facebook's MoCo repository](https://github.com/facebookresearch/moco).\n* [`moco_temp.py`](https://github.com/eminorhan/baby-vision/blob/master/moco_temp.py): trains a temporal MoCo model as described in the paper. This code was also modified from [Facebook's MoCo repository](https://github.com/facebookresearch/moco).\n* [`moco_utils.py`](https://github.com/eminorhan/baby-vision/blob/master/moco_utils.py): some utility functions for MoCo training.\n* [`linear_decoding.py`](https://github.com/eminorhan/baby-vision/blob/master/linear_decoding.py): evaluates self-supervised models on downstream linear classification tasks.\n* [`linear_combination_maps.py`](https://github.com/eminorhan/baby-vision/blob/master/linear_combination_maps.py): plots spatial attention maps as in Figure 4b and Figure 6 in the paper.\n* [`highly_activating_imgs.py`](https://github.com/eminorhan/baby-vision/blob/master/highly_activating_imgs.py): finds highly activating images for a given feature as in Figure 7b in the paper.\n* [`selectivities.py`](https://github.com/eminorhan/baby-vision/blob/master/selectivities.py): measures the class selecitivity indices of all features in a given layer as in Figure 7a in the paper.\n* [`hog_baseline.py`](https://github.com/eminorhan/baby-vision/blob/master/hog_baseline.py): runs the HOG baseline model as described in the paper.\n* [`imagenet_finetuning.py`](https://github.com/eminorhan/baby-vision/blob/master/imagenet_finetuning.py): ImageNet evaluations.\n* [`feature_animation.py`](https://github.com/eminorhan/baby-vision/blob/master/feature_animation.py) and [`feature_animation_class.py`](https://github.com/eminorhan/baby-vision/blob/master/feature_animation_class.py): Some tools for visualizing the learned features.\n\nFor specific usage examples, please see the slurm scripts provided in the [`scripts`](https://github.com/eminorhan/baby-vision/tree/master/scripts) directory.\n\n## Pre-trained models\n\n### ResNeXt \nSince the publication of the paper, we have found that training larger capacity models for longer with the temporal classification objective significantly improves the evaluation results. Hence, we provide below pre-trained `resnext50_32x4d` type models that are currently our best models trained with the SAYCam data. We encourage people to use these new models instead of the `mobilenet_v2` type models reported in the paper (the pre-trained `mobilenet_v2` models reported in the paper are also provided below for the record). \n\nFour pre-trained `resnext50_32x4d` models are provided here: temporal classification models trained on data from the individual children in the SAYCam dataset (`TC-S-resnext`, `TC-A-resnext`, `TC-Y-resnext`) and a temporal classification model trained on data from all three children (`TC-SAY-resnext`). These models were all trained for 16 epochs (with batch size 256) with the following data augmentation pipeline:\n\n```python\nimport torchvision.transforms as tr\n\ntr.Compose([\n\ttr.RandomResizedCrop(224, scale=(0.2, 1.)),\n        tr.RandomApply([tr.ColorJitter(0.9, 0.9, 0.9, 0.5)], p=0.9),\n        tr.RandomGrayscale(p=0.2),\n        tr.RandomApply([GaussianBlur([.1, 2.])], p=0.5),\n        tr.RandomHorizontalFlip(),\n        tr.ToTensor(),\n        tr.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n])\n```\n\nThis data augmentation pipeline is similar to that used in [the SimCLR paper](https://arxiv.org/abs/2002.05709) with slightly larger random crops and slightly stronger color augmentation. Here are some evaluation results for these `resnext50_32x4d` models (to download the models, click on the links over the model names):\n\n| Model | Toybox (*iid*) | Toybox (*exemplar*) | ImageNet (*linear*) | ImageNet (*1% ft + linear*) | \n| ----- |:--------------:|:-------------------:|:-------------------:|:---------------------------:|\n| [`TC-SAY-resnext`](https://drive.google.com/file/d/1I-HvIeuupsE88yS6eff_nE6pHpEmTVPG/view?usp=sharing)  | **90.0** | **57.5** | **36.0** | **45.6** |\n| [`TC-S-resnext`](https://drive.google.com/file/d/14tZeOtK1Jd64ioxPwzwf2jblriN7Jgue/view?usp=sharing)    | 88.5 | 54.9 | -- | -- |\n| [`TC-A-resnext`](https://drive.google.com/file/d/1aQuWfb4O0xL0PALRJpYIUHk0tsyrujDF/view?usp=sharing)    | 86.8 | 50.4 | -- | -- |\n| [`TC-Y-resnext`](https://drive.google.com/file/d/1sB12pdnVEZsgVKiVdZyS0l4x24_T5zCj/view?usp=sharing)    | 87.0 | 53.0 | -- | -- |\n\nHere, **ImageNet (*linear*)** refers to the top-1 validation accuracy on ImageNet with only a linear classifier trained on top of the frozen features, and **ImageNet (*1% ft + linear*)** is similar but with the entire model first fine-tuned on 1% of the ImageNet training data (~12800 images). Note that these are results from a single run, so you may observe slightly different numbers.\n\nThese models come with the temporal classification heads attached. To load these models, please do something along the lines of:\n\n```python\nimport torch\nimport torchvision.models as models\n\nmodel = models.resnext50_32x4d(pretrained=False)\nmodel.fc = torch.nn.Linear(in_features=2048, out_features=n_out, bias=True)\nmodel = torch.nn.DataParallel(model).cuda()\n\ncheckpoint = torch.load('TC-SAY-resnext.tar')\nmodel.load_state_dict(checkpoint['model_state_dict'])\n```\n\nwhere `n_out` should be 6269 for `TC-SAY-resnext`, 2765 for `TC-S-resnext`, 1786 for `TC-A-resnext`, and 1718 for `TC-Y-resnext`. The differences here are due to the different lengths of the datasets. \n\nIn addition, please find below the best performing ImageNet models reported above: a model with a linear ImageNet classifier trained on top of the frozen features of `TC-SAY-resnext` (`TC-SAY-resnext-IN-linear`) and a model that was first fine-tuned with 1% of the ImageNet training data (`TC-SAY-resnext-IN-1pt-linear`):\n\n* [`TC-SAY-resnext-IN-linear`](https://drive.google.com/file/d/1Qo0_1RwgOsr-JM3lP4ILWRY0WflnS7On/view?usp=sharing)\n* [`TC-SAY-resnext-IN-1pt-linear`](https://drive.google.com/file/d/1lvCG3L1_-gdqWDMD41yTIbuNBpzpUOQq/view?usp=sharing)\n\nYou can load these models in the same way as described above. Since these are ImageNet models, `n_out` should be set to 1000.\n\n### MobileNet \nThe following are the pre-trained `mobilenet_v2` type models reported in the paper:\n\n* [TC-S-mobilenet](https://drive.google.com/file/d/1DVJjpaGhoBPNmlO7jXpwEX3lSCk2ZUCa/view?usp=sharing) (69.4 MB)\n* [TC-A-mobilenet](https://drive.google.com/file/d/1uQvJBbuy6P0uCW0HYs1wNgawRU8sGLhC/view?usp=sharing) (54.4 MB)\n* [TC-Y-mobilenet](https://drive.google.com/file/d/1TTndiiiqSiCMdZjwYZPKQySZot4ipCrG/view?usp=sharing) (53.3 MB)\n* [TC-SAY-mobilenet](https://drive.google.com/file/d/1zeidpBaXqqWCeeYj-fMI7V7x9EiAGH6Q/view?usp=sharing) (123.3 MB)\n\n## Acknowledgments\n\nWe are very grateful to the volunteers who contributed recordings to the SAYCam dataset. We thank Jessica Sullivan for her generous assistance with the dataset. We also thank the team behind the Toybox dataset, as well as the developers of PyTorch and torchvision for making this work possible. This project was partly funded by the NSF Award 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science.\n"
  },
  {
    "path": "feature_animation.py",
    "content": "'''Animating features on short clips'''\nimport os\nimport argparse\nimport numpy as np\nimport torch\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom torchvision.utils import make_grid\nimport matplotlib as mp\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\nimport matplotlib.cm as cm\n\n# TODO: combine the map extraction functions into a single function \n# TODO: combine model loading functions into a single function \n\ndef extract_map_layer_7x7_res(res_model):\n    layer_list = list(res_model.module.children())[:-2]\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_7x7(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_14x14(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_layer_list = layer_list[:-5]\n    new_layer_list.append(layer_list[-5].conv[0])\n    new_model = torch.nn.Sequential(*new_layer_list)\n    return new_model\n\ndef load_model_res(args):\n    model = models.resnext50_32x4d(pretrained=False)\n    model.fc = torch.nn.Linear(in_features=2048, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_model(args):\n    model = models.mobilenet_v2(pretrained=True)\n    model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_data(data_dir, args):\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    train_dataset = datasets.ImageFolder(\n        data_dir,\n        transforms.Compose([transforms.Resize(224), transforms.ToTensor(), normalize])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    return train_loader\n\ndef predict(data_loader, model, batch_size, feature_idx):\n\n    # switch to evaluate mode\n    model.eval()\n\n    preds_list = []\n    imgs_list = []\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(data_loader):\n            images = images.cuda()\n\n            # compute predictions\n            preds = model(images)\n            preds = preds[:, feature_idx, :, :]\n            preds_list.append(preds)\n            imgs_list.append(images)\n\n    preds = torch.cat(preds_list, 0)\n    images = torch.cat(imgs_list, 0)\n\n    print('Images shape:', images.size())\n    print('Preds shape:', preds.size())\n\n    # Copy activation map to all channels and upsample to image size\n    x = torch.zeros(preds.size()[0], 3, 7, 7)\n    x[:, 0, :, :] = preds\n    x[:, 1, :, :] = preds\n    x[:, 2, :, :] = preds\n\n    m = torch.nn.Upsample(scale_factor=32, mode='bicubic')\n\n    upsampled_maps = m(x).cuda()\n    # upsampled_maps = torch.sigmoid(10. * upsampled_maps / torch.std(upsampled_maps))\n\n    upsampled_maps = upsampled_maps.cpu().numpy()\n    images = images.cpu().numpy()\n\n    return  upsampled_maps, images\n\ndef show_img(ax, img, save_name):\n    '''Save maps'''\n    npimg = img.cpu().numpy()\n\n    print(npimg.shape)\n\n    ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n\n    ax.spines[\"bottom\"].set_visible(False)\n    ax.spines[\"left\"].set_visible(False)\n    ax.spines[\"right\"].set_visible(False)\n    ax.spines[\"top\"].set_visible(False)\n\n    mp.rcParams['axes.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 1.15\n    mp.rcParams['font.sans-serif'] = ['FreeSans']\n    mp.rcParams['mathtext.fontset'] = 'cm'\n\n    plt.savefig(save_name, bbox_inches='tight')\n\n\nif __name__ == '__main__':\n\n    parser = argparse.ArgumentParser(description='Plot spatial attention maps')\n    parser.add_argument('data', metavar='DIR', help='path to dataset')\n    parser.add_argument('--workers', default=32, type=int, help='number of data loading workers (default: 4)')\n    parser.add_argument('--batch-size', default=900, type=int, help='mini-batch size, this is the total '\n                                                                    'batch size of all GPUs on the current node when '\n                                                                    'using Data Parallel or Distributed Data Parallel')\n    parser.add_argument('--model-path', default='', type=str, help='path to model checkpoint (default: '\n                                                                   'ImageNet-pretrained)')\n    parser.add_argument('--n_out', default=2765, type=int, help='output dim of pre-trained model')\n    parser.add_argument('--feature-idx', default=1, type=int, help='feature index for which the maps will be computed')\n\n    args = parser.parse_args()\n\n    model = load_model(args)\n    map_layer = extract_map_layer_7x7(model)\n\n    data_loader = load_data(args.data, args)\n    preds, images = predict(data_loader, map_layer, args.batch_size, args.feature_idx)\n\n    preds = preds - preds.min()\n    preds = preds / preds.max()\n    preds = np.uint8(255 * preds)\n\n    images = images - images.min()\n    images = images / images.max()\n    # images = np.uint8(255 * images)\n\n    fig, ax = plt.subplots()\n    ax.set_axis_off()\n    ax.set_title('Feature: ' + str(args.feature_idx))\n\n    jet = cm.get_cmap(\"jet\")\n    jet_colors = jet(np.arange(256))[:, :3]\n    preds = jet_colors[preds[:, 0, :, :]]\n\n    masked_imgs = 1.0 * preds + np.transpose(images, (0, 2, 3, 1))\n    masked_imgs = np.uint8(255 * masked_imgs / masked_imgs.max())\n\n    imgs = []\n    for i in range(900):\n        im = ax.imshow(masked_imgs[i])\n\n        if i == 0:\n            im = ax.imshow(masked_imgs[i])\n\n        imgs.append([im])\n\n    ani = animation.ArtistAnimation(fig, imgs, interval=200, blit=True, repeat_delay=1000)\n\n    # To save the animation, use e.g.\n    ani.save('intphys_feature_animation_' + str(args.feature_idx) + '.mp4')"
  },
  {
    "path": "feature_animation_class.py",
    "content": "'''Animating features on short clips'''\nimport os\nimport argparse\nimport numpy as np\nimport torch\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom torchvision.utils import make_grid\nimport matplotlib as mp\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\nimport matplotlib.cm as cm\n\n# TODO: combine the map extraction functions into a single function \n# TODO: combine model loading functions into a single function \n\ndef extract_map_layer_7x7_res(res_model):\n    layer_list = list(res_model.module.children())[:-2]\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_7x7(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_14x14(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_layer_list = layer_list[:-5]\n    new_layer_list.append(layer_list[-5].conv[0])\n    new_model = torch.nn.Sequential(*new_layer_list)\n    return new_model\n\ndef load_model_res(args):\n    model = models.resnext50_32x4d(pretrained=False)\n    model.fc = torch.nn.Linear(in_features=2048, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_model(args):\n    model = models.mobilenet_v2(pretrained=True)\n    model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_data(data_dir, args):\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    train_dataset = datasets.ImageFolder(\n        data_dir,\n        transforms.Compose([transforms.ToTensor(), normalize])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    return train_loader\n\ndef predict(data_loader, model, batch_size, weights):\n\n    # switch to evaluate mode\n    model.eval()\n\n    preds_list = []\n    imgs_list = []\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(data_loader):\n            images = images.cuda()\n\n            # compute predictions\n            preds = model(images)\n            preds_list.append(preds)\n            imgs_list.append(images)\n\n    preds = torch.cat(preds_list, 0)\n    images = torch.cat(imgs_list, 0)\n\n    print('Images shape:', images.size())\n    print('Preds shape:', preds.size())\n\n    linear_combination_map = torch.einsum('ijkl,j->ikl', preds, weights)\n\n    # Copy activation map to all channels and upsample to image size\n    x = torch.zeros(preds.size()[0], 3, 7, 7)\n    x[:, 0, :, :] = linear_combination_map\n    x[:, 1, :, :] = linear_combination_map\n    x[:, 2, :, :] = linear_combination_map\n\n    m = torch.nn.Upsample(scale_factor=32, mode='bicubic')\n\n    upsampled_maps = m(x).cuda()\n    # upsampled_maps = torch.sigmoid(10. * upsampled_maps / torch.std(upsampled_maps))\n\n    upsampled_maps = upsampled_maps.cpu().numpy()\n    images = images.cpu().numpy()\n\n    return  upsampled_maps, images\n\ndef show_img(ax, img, save_name):\n    '''Save maps'''\n    npimg = img.cpu().numpy()\n\n    print(npimg.shape)\n\n    ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n\n    ax.spines[\"bottom\"].set_visible(False)\n    ax.spines[\"left\"].set_visible(False)\n    ax.spines[\"right\"].set_visible(False)\n    ax.spines[\"top\"].set_visible(False)\n\n    mp.rcParams['axes.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 1.15\n    mp.rcParams['font.sans-serif'] = ['FreeSans']\n    mp.rcParams['mathtext.fontset'] = 'cm'\n\n    plt.savefig(save_name, bbox_inches='tight')\n\n\nif __name__ == '__main__':\n\n    parser = argparse.ArgumentParser(description='Plot spatial attention maps')\n    parser.add_argument('data', metavar='DIR', help='path to dataset')\n    parser.add_argument('--workers', default=32, type=int, help='number of data loading workers (default: 4)')\n    parser.add_argument('--batch-size', default=500, type=int, help='mini-batch size, this is the total '\n                                                                    'batch size of all GPUs on the current node when '\n                                                                    'using Data Parallel or Distributed Data Parallel')\n    parser.add_argument('--model-path', default='', type=str, help='path to model checkpoint (default: '\n                                                                   'ImageNet-pretrained)')\n    parser.add_argument('--n_out', default=26, type=int, help='output dim of pre-trained model')\n    parser.add_argument('--class-idx', default=1, type=int, help='class index for which the maps will be computed')\n\n    args = parser.parse_args()\n\n    model = load_model(args)\n    map_layer = extract_map_layer_7x7(model)\n\n    weights = model.module.classifier.weight.data[args.class_idx, :].cuda()\n\n    data_loader = load_data(args.data, args)\n    preds, images = predict(data_loader, map_layer, args.batch_size, weights)\n\n    preds = preds - preds.min()\n    preds = preds / preds.max()\n    preds = np.uint8(255 * preds)\n\n    images = images - images.min()\n    images = images / images.max()\n    # images = np.uint8(255 * images)\n\n    fig, ax = plt.subplots()\n    ax.set_axis_off()\n    ax.set_title('Class: ' + str(args.class_idx))\n\n    jet = cm.get_cmap(\"jet\")\n    jet_colors = jet(np.arange(256))[:, :3]\n    preds = jet_colors[preds[:, 0, :, :]]\n\n    masked_imgs = 1.0 * preds + np.transpose(images, (0, 2, 3, 1))\n    masked_imgs = np.uint8(255 * masked_imgs / masked_imgs.max())\n\n    imgs = []\n    for i in range(200):\n        im = ax.imshow(masked_imgs[i])\n\n        if i == 0:\n            im = ax.imshow(masked_imgs[i])\n\n        imgs.append([im])\n\n    ani = animation.ArtistAnimation(fig, imgs, interval=200, blit=True, repeat_delay=1000)\n\n    # To save the animation, use e.g.\n    ani.save('computers_feature_animation_' + str(args.class_idx) + '.mp4')"
  },
  {
    "path": "highly_activating_imgs.py",
    "content": "'''Plots highly activating images'''\nimport os\nimport argparse\nimport numpy as np\nimport torch\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom torchvision.utils import make_grid\nimport matplotlib as mp\nimport matplotlib.pyplot as plt\n\n\ndef extract_map_layer_7x7(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_14x14(mobilenetV2_model, layer):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_layer_list = layer_list[:-layer]\n    new_layer_list.append(layer_list[-layer].conv[0])\n    new_model = torch.nn.Sequential(*new_layer_list)\n    return new_model\n\ndef load_model(args):\n    model = models.mobilenet_v2(pretrained=True)\n    model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_data(data_dir, args):\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    train_dataset = datasets.ImageFolder(\n        data_dir,\n        transforms.Compose([transforms.ToTensor(), normalize])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=True,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    return train_loader\n\ndef predict(data_loader, model, neuron_idx):\n\n    # switch to evaluate mode\n    model.eval()\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(data_loader):\n            images = images.cuda()\n\n            # compute predictions\n            pred = model(images)\n            pred_mean = torch.mean(pred, dim=(2, 3))\n            pred_mean = pred_mean[:, neuron_idx]\n\n            if i == 0:\n                break\n\n    _, indices = torch.sort(pred_mean, descending=True)\n    images = images[indices, :, :, :]\n\n    return images\n\ndef show_img(ax, img, save_name):\n    '''Save maps'''\n    npimg = img.cpu().numpy()\n\n    print(npimg.shape)\n\n    ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n\n    ax.spines[\"bottom\"].set_visible(False)\n    ax.spines[\"left\"].set_visible(False)\n    ax.spines[\"right\"].set_visible(False)\n    ax.spines[\"top\"].set_visible(False)\n\n    mp.rcParams['axes.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 1.15\n    mp.rcParams['font.sans-serif'] = ['FreeSans']\n    mp.rcParams['mathtext.fontset'] = 'cm'\n\n    plt.savefig(save_name, bbox_inches='tight')\n\n\nif __name__ == '__main__':\n\n    parser = argparse.ArgumentParser(description='Plot highly activating images for a given feature')\n    parser.add_argument('data', metavar='DIR', help='path to dataset')\n    parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 4)')\n    parser.add_argument('--batch-size', default=1024, type=int, help='mini-batch size, this is the total '\n                                                                    'batch size of all GPUs on the current node when '\n                                                                    'using Data Parallel or Distributed Data Parallel')\n    parser.add_argument('--model-path', default='', type=str, help='path to latest checkpoint (default: none)')\n    parser.add_argument('--n_out', default=1000, type=int, help='output dim')\n    parser.add_argument('--neuron_idx', default=276, type=int, help='neuron index')\n\n    args = parser.parse_args()\n\n    model = load_model(args)\n    map_layer = extract_map_layer_7x7(model)\n\n    data_loader = load_data(args.data, args)\n    imgs = predict(data_loader, map_layer, neuron_idx=args.neuron_idx)\n\n    print('Imgs shape', imgs.shape)\n\n    print('Plotting the top 10 images')\n    fig_img = plt.figure(figsize=(16, 16), dpi=300)\n    ax_img = fig_img.add_subplot('111')\n    grid_img = make_grid(imgs[:10, :, :, :], nrow=10, padding=2, normalize=True, scale_each=False)\n    show_img(ax_img, grid_img, 'highly_activating_imgs_neuron_' + str(args.neuron_idx) + '.pdf')"
  },
  {
    "path": "hog_baseline.py",
    "content": "'''HoG baseline'''\nimport os\nimport argparse\nimport numpy as np\nfrom skimage.feature import hog\nfrom skimage.io import imread\nfrom sklearn.linear_model import SGDClassifier\nfrom sklearn.model_selection import train_test_split\n\nparser = argparse.ArgumentParser(description='Linear decoding with HoG model')\nparser.add_argument('data', metavar='DIR', help='path to dataset')\nparser.add_argument('--subsample', default=False, action='store_true', help='subsample data?')\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n\n    c_list = os.listdir(args.data)\n    c_list.sort()\n    print('Class list:', c_list)\n\n    imgs = []\n    labels = []\n    label_counter = 0\n    file_counter = 0\n\n    for c in c_list:\n            curr_dir = os.path.join(args.data, c)\n            f_list = os.listdir(curr_dir)\n            f_list.sort()\n\n            print('Reading class:', c)\n\n            for f in f_list:\n                    f_path = os.path.join(curr_dir, f)\n                    img = imread(f_path)\n                    feats = hog(img, orientations=9, pixels_per_cell=(16, 16), cells_per_block=(3, 3), block_norm='L2',\n                                visualize=False, transform_sqrt=False, feature_vector=True, multichannel=True)\n                    if args.subsample:\n                            if file_counter % 10 == 0:\n                                    imgs.append(feats)\n                                    labels.append(label_counter)\n                    else:\n                            imgs.append(feats)\n                            labels.append(label_counter)\n\n                    file_counter += 1\n\n            label_counter += 1\n\n    imgs = np.vstack(imgs)\n    labels = np.array(labels)\n\n    print('Imgs shape:', imgs.shape)\n    print('Labels shape:', labels.shape)\n\n    print('Splitting dataset')\n    X_train, X_test, y_train, y_test = train_test_split(imgs, labels, test_size=0.5)\n\n    print('Fitting training data')\n    clf = SGDClassifier(loss=\"hinge\", penalty=\"l2\", alpha=0.0001, max_iter=250)\n    clf.fit(X_train, y_train)\n\n    print('Computing predictions')\n    pred_test = clf.predict(X_test)\n    test_acc = np.mean(y_test==pred_test)\n\n    pred_train = clf.predict(X_train)\n    train_acc = np.mean(y_train==pred_train)\n\n    print('Test accuracy', test_acc)\n    print('Train accuracy', train_acc)"
  },
  {
    "path": "imagenet_finetuning.py",
    "content": "import argparse\nimport os\nimport time\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.optim\nimport torch.multiprocessing as mp\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\n\n\nparser = argparse.ArgumentParser(description='ImageNet fine-tuning or linear classification')\nparser.add_argument('-j', '--workers', default=32, type=int, metavar='N', help='number of data loading workers (default:'\n                                                                              ' 4)')\nparser.add_argument('--epochs', default=25, type=int, metavar='N', help='number of total epochs to run')\nparser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')\nparser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N',\n                    help='mini-batch size (default: 256), this is the total batch size of all GPUs on the current node '\n                         'when using Data Parallel or Distributed Data Parallel')\nparser.add_argument('--lr', '--learning-rate', default=0.0005, type=float, metavar='LR', help='initial learning rate',\n                    dest='lr')\nparser.add_argument('--wd', '--weight-decay', default=0.0, type=float, metavar='W', help='weight decay (default: 0)',\n                    dest='weight_decay')\nparser.add_argument('-p', '--print-freq', default=5000, type=int, metavar='N', help='print frequency (default: 100)')\nparser.add_argument('--schedule', default=[23, 24], nargs='*', type=int,\n                    help='learning rate schedule (when to drop lr by a ratio)')\nparser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')\nparser.add_argument('--world-size', default=-1, type=int, help='number of nodes for distributed training')\nparser.add_argument('--rank', default=-1, type=int, help='node rank for distributed training')\nparser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str, help='url used to set up distributed '\n                                                                                     'training')\nparser.add_argument('--dist-backend', default='nccl', type=str, help='distributed backend')\nparser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')\nparser.add_argument('--multiprocessing-distributed', action='store_true',\n                    help='Use multi-processing distributed training to launch '\n                         'N processes per node, which has N GPUs. This is the '\n                         'fastest way to use PyTorch for either single node or '\n                         'multi node data parallel training')\nparser.add_argument('--n_out', default=20, type=int, help='output dim')\nparser.add_argument('--freeze-trunk', default=False, action='store_true', help='freeze trunk?')\nparser.add_argument('--frac-retained', default=1.0, type=float, help='fraction of tr data retained')\n\n\ndef set_parameter_requires_grad(model, feature_extracting=True):\n    '''Helper function for setting body to non-trainable'''\n    if feature_extracting:\n        for param in model.parameters():\n            param.requires_grad = False\n\n        for param in model.module.fc.parameters():\n            print(param.shape)\n            param.requires_grad = True\n\ndef main():\n    args = parser.parse_args()\n\n    if args.gpu is not None:\n        warnings.warn('You have chosen a specific GPU. This will completely disable data parallelism.')\n\n    if args.dist_url == \"env://\" and args.world_size == -1:\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    args.distributed = args.world_size > 1 or args.multiprocessing_distributed\n\n    ngpus_per_node = torch.cuda.device_count()\n    if args.multiprocessing_distributed:\n        # Since we have ngpus_per_node processes per node, the total world_size needs to be adjusted accordingly\n        args.world_size = ngpus_per_node * args.world_size\n        # Use torch.multiprocessing.spawn to launch distributed processes: the main_worker process function\n        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))\n    else:\n        # Simply call main_worker function\n        main_worker(args.gpu, ngpus_per_node, args)\n\n\ndef main_worker(gpu, ngpus_per_node, args):\n    args.gpu = gpu\n\n    if args.gpu is not None:\n        print(\"Use GPU: {} for training\".format(args.gpu))\n\n    if args.distributed:\n        if args.dist_url == \"env://\" and args.rank == -1:\n            args.rank = int(os.environ[\"RANK\"])\n        if args.multiprocessing_distributed:\n            # For multiprocessing distributed training, rank needs to be the\n            # global rank among all the processes\n            args.rank = args.rank * ngpus_per_node + gpu\n        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,\n                                world_size=args.world_size, rank=args.rank)\n\n    model = models.resnext50_32x4d(pretrained=False)\n    model.fc = torch.nn.Linear(in_features=2048, out_features=args.n_out, bias=True)\n\n    # DataParallel will divide and allocate batch_size to all available GPUs\n    model = torch.nn.DataParallel(model).cuda()\n\n    # if resume from a pretrained model\n    if args.resume:\n        if os.path.isfile(args.resume):\n            print(\"=> loading model '{}'\".format(args.resume))\n            checkpoint = torch.load(args.resume)\n            model.load_state_dict(checkpoint['model_state_dict'])\n            if args.freeze_trunk:\n                print('Freezing trunk.')\n                set_parameter_requires_grad(model)  # freeze the trunk\n            model.module.fc = torch.nn.Linear(in_features=2048, out_features=1000, bias=True).cuda()\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.resume))\n    else:\n        if args.freeze_trunk:\n            print('Freezing trunk.')\n            set_parameter_requires_grad(model)  # freeze the trunk\n            model.module.fc = torch.nn.Linear(in_features=2048, out_features=1000, bias=True).cuda()\n\n    print(model)\n\n    # define loss function (criterion) and optimizer\n    criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n    optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)\n    cudnn.benchmark = True\n\n    # Save file name\n    if args.resume:\n        sv_name = args.resume\n        savefile_name = 'ft_IN_' + sv_name  # str(args.freeze_trunk) + 'fz_IN_' + sv_name[26:]\n    else:\n        savefile_name = str(args.freeze_trunk) + 'fz_IN_MobileNetV2_scratch.tar'\n\n    # Data loaders\n    basedir = '/misc/vlgscratch4/LakeGroup/emin/robust_vision/imagenet/'\n\n    traindir = os.path.join(basedir, 'train')\n    valdir = os.path.join(basedir, 'val')\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                     std=[0.229, 0.224, 0.225])\n\n    train_dataset = datasets.ImageFolder(\n        traindir,\n        transforms.Compose([\n            transforms.RandomResizedCrop(224),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            normalize\n        ])\n    )\n\n    val_dataset = datasets.ImageFolder(valdir, transforms.Compose([\n            transforms.Resize(256),\n            transforms.CenterCrop(224),\n            transforms.ToTensor(),\n            normalize\n        ]))\n\n    if args.frac_retained < 1.0:\n        print('Fraction of train data retained:', args.frac_retained)\n\n        import numpy as np\n        num_train = len(train_dataset)\n        indices = list(range(num_train))\n        np.random.shuffle(indices)\n        train_idx = indices[:int(args.frac_retained * num_train)]\n        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx)\n\n        train_loader = torch.utils.data.DataLoader(\n            train_dataset, batch_size=args.batch_size, shuffle=False,\n            num_workers=args.workers, pin_memory=True, sampler=train_sampler)\n    else:\n        print('Using all of train data')\n\n        train_loader = torch.utils.data.DataLoader(\n            train_dataset, batch_size=args.batch_size, shuffle=True,\n            num_workers=args.workers, pin_memory=True, sampler=None)\n\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True)\n\n    acc1_list = []\n    val_acc1_list = []\n\n    for epoch in range(args.start_epoch, args.epochs):\n\n        adjust_learning_rate(optimizer, epoch, args)\n\n        # train for one epoch\n        acc1 = train(train_loader, model, criterion, optimizer, epoch, args)\n        acc1_list.append(acc1)\n\n        # ... then validate\n        val_acc1 = validate(val_loader, model, args)\n        val_acc1_list.append(val_acc1)\n\n    torch.save({'acc1_list': acc1_list,\n                'val_acc1_list': val_acc1_list,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict()}, savefile_name)\n\n\ndef train(train_loader, model, criterion, optimizer, epoch, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    data_time = AverageMeter('Data', ':6.3f')\n    losses = AverageMeter('Loss', ':.4e')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses, top1, top5],\n        prefix=\"Epoch: [{}]\".format(epoch))\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    for i, (images, target) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n\n        if args.gpu is not None:\n            images = images.cuda(args.gpu, non_blocking=True)\n        target = target.cuda(args.gpu, non_blocking=True)\n\n        # compute output\n        output = model(images)\n        loss = criterion(output, target)\n\n        # measure accuracy and record loss\n        acc1, acc5 = accuracy(output, target, topk=(1, 5))\n        losses.update(loss.item(), images.size(0))\n        top1.update(acc1[0], images.size(0))\n        top5.update(acc5[0], images.size(0))\n\n        # compute gradient and do SGD step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % args.print_freq == 0:\n            progress.display(i)\n\n    return top1.avg.cpu().numpy()\n\n\ndef validate(val_loader, model, args):\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n\n    # switch to eval mode\n    model.eval()\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(val_loader):\n\n            if args.gpu is not None:\n                images = images.cuda(args.gpu, non_blocking=True)\n            target = target.cuda(args.gpu, non_blocking=True)\n\n            # compute output\n            output = model(images)\n\n            # measure accuracy and record loss\n            acc1, acc5 = accuracy(output, target, topk=(1, 5))\n            top1.update(acc1[0], images.size(0))\n            top5.update(acc5[0], images.size(0))\n\n    print('End of epoch validation: * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))\n\n    return top1.avg.cpu().numpy()\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'\n\n\ndef adjust_learning_rate(optimizer, epoch, args):\n    \"\"\"Decay the learning rate based on schedule\"\"\"\n    lr = args.lr\n    for milestone in args.schedule:\n        lr *= 0.2 if epoch >= milestone else 1.\n    for param_group in optimizer.param_groups:\n        param_group['lr'] = lr\n\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n    with torch.no_grad():\n        maxk = max(topk)\n        batch_size = target.size(0)\n\n        _, pred = output.topk(maxk, 1, True, True)\n        pred = pred.t()\n        correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n        res = []\n        for k in topk:\n            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n            res.append(correct_k.mul_(100.0 / batch_size))\n        return res\n\n\nif __name__ == '__main__':\n    main()"
  },
  {
    "path": "linear_combination_maps.py",
    "content": "'''Plots spatial attention maps'''\nimport os\nimport argparse\nimport numpy as np\nimport torch\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom torchvision.utils import make_grid\nimport matplotlib as mp\nimport matplotlib.pyplot as plt\n\ndef extract_map_layer_7x7(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_14x14(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_layer_list = layer_list[:-5]\n    new_layer_list.append(layer_list[-5].conv[0])\n    new_model = torch.nn.Sequential(*new_layer_list)\n    return new_model\n\ndef load_model(args):\n    model = models.mobilenet_v2(pretrained=True)\n    model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_data(data_dir, args):\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    train_dataset = datasets.ImageFolder(\n        data_dir,\n        transforms.Compose([transforms.ToTensor(), normalize])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=True,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    return train_loader\n\ndef predict(data_loader, model, weights, batch_size):\n\n    # switch to evaluate mode\n    model.eval()\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(data_loader):\n            images = images.cuda()\n\n            print(images.size())\n\n            # compute predictions\n            pred = model(images)\n\n            if i == 0:\n                break\n\n    linear_combination_map = torch.einsum('ijkl,j->ikl', pred, weights)\n\n    x = torch.zeros(batch_size, 3, 7, 7)\n    x[:, 0, :, :] = linear_combination_map\n    x[:, 1, :, :] = linear_combination_map\n    x[:, 2, :, :] = linear_combination_map\n\n    m = torch.nn.Upsample(scale_factor=32, mode='bicubic')\n    mm = m(x).cuda()\n    mm = torch.sigmoid(10. * mm / torch.std(mm))\n\n    return mm * images\n\ndef show_img(ax, img, save_name):\n    '''Save maps'''\n    npimg = img.cpu().numpy()\n\n    print(npimg.shape)\n\n    ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n\n    ax.spines[\"bottom\"].set_visible(False)\n    ax.spines[\"left\"].set_visible(False)\n    ax.spines[\"right\"].set_visible(False)\n    ax.spines[\"top\"].set_visible(False)\n\n    mp.rcParams['axes.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 0.75\n    mp.rcParams['patch.linewidth'] = 1.15\n    mp.rcParams['font.sans-serif'] = ['FreeSans']\n    mp.rcParams['mathtext.fontset'] = 'cm'\n\n    plt.savefig(save_name, bbox_inches='tight')\n\n\nif __name__ == '__main__':\n\n    parser = argparse.ArgumentParser(description='Plot spatial attention maps')\n    parser.add_argument('data', metavar='DIR', help='path to dataset')\n    parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 4)')\n    parser.add_argument('--batch-size', default=36, type=int, help='mini-batch size, this is the total '\n                                                                    'batch size of all GPUs on the current node when '\n                                                                    'using Data Parallel or Distributed Data Parallel')\n    parser.add_argument('--model-path', default='', type=str, help='path to model checkpoint (default: '\n                                                                   'ImageNet-pretrained)')\n    parser.add_argument('--n_out', default=1000, type=int, help='output dim of pre-trained model')\n    parser.add_argument('--class-idx', default=6, type=int, help='class index for which the maps will be computed')\n\n    args = parser.parse_args()\n\n    model = load_model(args)\n    map_layer = extract_map_layer_7x7(model)\n\n    weights = model.module.classifier.weight.data[args.class_idx, :].cuda()\n\n    data_loader = load_data(args.data, args)\n    preds = predict(data_loader, map_layer, weights, args.batch_size)\n\n    print('Preds shape:', preds.shape)\n\n    fig_pred = plt.figure(figsize=(16, 16), dpi=300)\n    ax_pred = fig_pred.add_subplot('111')\n    grid_pred = make_grid(preds, nrow=12, padding=1, normalize=True, scale_each=False)\n    show_img(ax_pred, grid_pred, 'linear_combination_maps_class_' + str(args. class_idx) + '.pdf')"
  },
  {
    "path": "linear_decoding.py",
    "content": "import argparse\nimport os\nimport random\nimport shutil\nimport time\nimport warnings\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.optim\nimport torch.multiprocessing as mp\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\n\n\nparser = argparse.ArgumentParser(description='Linear decoding with headcam data')\nparser.add_argument('data', metavar='DIR', help='path to dataset')\nparser.add_argument('-j', '--workers', default=32, type=int, metavar='N', help='number of data loading workers (default: 32)')\nparser.add_argument('--epochs', default=100, type=int, metavar='N', help='number of total epochs to run')\nparser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')\nparser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N',\n                    help='mini-batch size (default: 1024), this is the total batch size of all GPUs on the current node '\n                         'when using Data Parallel or Distributed Data Parallel')\nparser.add_argument('--lr', '--learning-rate', default=0.0005, type=float, metavar='LR', help='initial learning rate', dest='lr')\nparser.add_argument('--wd', '--weight-decay', default=0.0, type=float, metavar='W', help='weight decay (default: 0)', dest='weight_decay')\nparser.add_argument('-p', '--print-freq', default=100, type=int, metavar='N', help='print frequency (default: 100)')\nparser.add_argument('--world-size', default=-1, type=int, help='number of nodes for distributed training')\nparser.add_argument('--rank', default=-1, type=int, help='node rank for distributed training')\nparser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str, help='url used to set up distributed training')\nparser.add_argument('--dist-backend', default='nccl', type=str, help='distributed backend')\nparser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')\nparser.add_argument('--multiprocessing-distributed', action='store_true',\n                    help='Use multi-processing distributed training to launch '\n                         'N processes per node, which has N GPUs. This is the '\n                         'fastest way to use PyTorch for either single node or '\n                         'multi node data parallel training')\nparser.add_argument('--model-name', type=str, default='random',\n                    choices=['random', 'imagenet', 'TC-S', 'TC-A', 'TC-Y', 'TC-SAY', 'moco_img_0011', 'moco_temp_0011'],\n                    help='evaluated model')\nparser.add_argument('--num-outs', default=16127, type=int, help='number of outputs in pretrained model')\nparser.add_argument('--num-classes', default=26, type=int, help='number of classes in downstream classification task')\nparser.add_argument('--subsample', default=False, action='store_true', help='subsample data?')\n\n\ndef set_parameter_requires_grad(model, feature_extracting=True):\n    '''Helper function for setting body to non-trainable'''\n    if feature_extracting:\n        for param in model.parameters():\n            param.requires_grad = False\n\n\ndef load_split_train_test(datadir, args, train_frac=0.5):\n\n    import numpy as np\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                     std=[0.229, 0.224, 0.225])\n\n    train_data = datasets.ImageFolder(datadir, transform=transforms.Compose([transforms.ToTensor(), normalize]))\n    test_data = datasets.ImageFolder(datadir, transform=transforms.Compose([transforms.ToTensor(), normalize]))\n\n    num_train = len(train_data)\n\n    print('Total data size is', num_train)\n\n    indices = list(range(num_train))\n    split = int(np.floor(train_frac * num_train))\n    np.random.shuffle(indices)\n\n    if args.subsample:\n        num_data = int(0.1 * num_train)\n        train_idx, test_idx = indices[:(num_data // 2)], indices[(num_data // 2):num_data]\n    else:\n        train_idx, test_idx = indices[:split], indices[split:]\n\n    print('Training data size is', len(train_idx))\n    print('Test data size is', len(test_idx))\n\n    train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx)\n    test_sampler = torch.utils.data.sampler.SubsetRandomSampler(test_idx)\n\n    trainloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True, sampler=train_sampler)\n    testloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True, sampler=test_sampler)\n\n    return trainloader, testloader\n\n\ndef main():\n    args = parser.parse_args()\n\n    if args.gpu is not None:\n        warnings.warn('You have chosen a specific GPU. This will completely disable data parallelism.')\n\n    if args.dist_url == \"env://\" and args.world_size == -1:\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    args.distributed = args.world_size > 1 or args.multiprocessing_distributed\n\n    ngpus_per_node = torch.cuda.device_count()\n    if args.multiprocessing_distributed:\n        # Since we have ngpus_per_node processes per node, the total world_size needs to be adjusted accordingly\n        args.world_size = ngpus_per_node * args.world_size\n        # Use torch.multiprocessing.spawn to launch distributed processes: the main_worker process function\n        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))\n    else:\n        # Simply call main_worker function\n        main_worker(args.gpu, ngpus_per_node, args)\n\n\ndef main_worker(gpu, ngpus_per_node, args):\n    args.gpu = gpu\n\n    if args.gpu is not None:\n        print(\"Use GPU: {} for training\".format(args.gpu))\n\n    # model definition\n    num_classes = args.num_classes\n\n    if args.model_name == 'random':\n        model = models.mobilenet_v2(pretrained=False)\n        set_parameter_requires_grad(model)\n        model.classifier = torch.nn.Linear(in_features=1280, out_features=num_classes, bias=True)\n        model = torch.nn.DataParallel(model).cuda()\n    elif args.model_name == 'imagenet':\n        model = models.mobilenet_v2(pretrained=True)\n        set_parameter_requires_grad(model)\n        model.classifier = torch.nn.Linear(in_features=1280, out_features=num_classes, bias=True)\n        model = torch.nn.DataParallel(model).cuda()\n    elif args.model_name.startswith('moco'):\n        model = models.mobilenet_v2(pretrained=False)\n        model.classifier = torch.nn.Linear(in_features=1280, out_features=args.num_outs, bias=True)\n        checkpoint = torch.load('../self_supervised_models/' + args.model_name + '.pth.tar')\n\n        # rename moco pre-trained keys\n        state_dict = checkpoint['state_dict']\n        for k in list(state_dict.keys()):\n            # retain only encoder_q up to before the embedding layer\n            if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.classifier'):\n                # remove prefix\n                state_dict[k[len(\"module.encoder_q.\"):]] = state_dict[k]\n            # delete renamed or unused k\n            del state_dict[k]\n\n        msg = model.load_state_dict(state_dict, strict=False)\n        assert set(msg.missing_keys) == {\"classifier.weight\", \"classifier.bias\"}\n\n        print(\"=> loaded pre-trained model '{}'\".format(args.model_name))\n\n        set_parameter_requires_grad(model)  # freeze the trunk\n        model.classifier = torch.nn.Linear(in_features=1280, out_features=num_classes, bias=True)\n        model = torch.nn.DataParallel(model).cuda()\n    else:\n        model = models.resnext50_32x4d(pretrained=False)\n        model.fc = torch.nn.Linear(in_features=2048, out_features=args.num_outs, bias=True)\n        model = torch.nn.DataParallel(model).cuda()\n        checkpoint = torch.load(args.model_name + '.tar')\n        model.load_state_dict(checkpoint['model_state_dict'])\n        set_parameter_requires_grad(model)  # freeze the trunk\n        model.module.fc = torch.nn.Linear(in_features=2048, out_features=num_classes, bias=True).cuda()\n\n    # define loss function (criterion) and optimizer\n    criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n    optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)\n    cudnn.benchmark = True\n\n    # Data loading code\n    savefile_name = args.model_name + '_labeledS_5_iid.tar'\n\n    train_loader, test_loader = load_split_train_test(args.data, args)\n    acc1_list = []\n    val_acc1_list = []\n\n    for epoch in range(args.start_epoch, args.epochs):\n\n        # train for one epoch\n        acc1 = train(train_loader, model, criterion, optimizer, epoch, args)\n        acc1_list.append(acc1)\n\n    # validate at end of epoch\n    val_acc1, preds, target, images = validate(test_loader, model, args)\n    val_acc1_list.append(val_acc1)\n\n    torch.save({'acc1_list': acc1_list,\n                'val_acc1_list': val_acc1_list,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'preds': preds,\n                'target': target,\n                'images': images\n                }, savefile_name)\n\n\ndef train(train_loader, model, criterion, optimizer, epoch, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    data_time = AverageMeter('Data', ':6.3f')\n    losses = AverageMeter('Loss', ':.4e')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses, top1, top5],\n        prefix=\"Epoch: [{}]\".format(epoch))\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    for i, (images, target) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n\n        if args.gpu is not None:\n            images = images.cuda(args.gpu, non_blocking=True)\n        target = target.cuda(args.gpu, non_blocking=True)\n\n        # compute output\n        output = model(images)\n        loss = criterion(output, target)\n\n        # measure accuracy and record loss\n        acc1, acc5 = accuracy(output, target, topk=(1, 2))\n        losses.update(loss.item(), images.size(0))\n        top1.update(acc1[0], images.size(0))\n        top5.update(acc5[0], images.size(0))\n\n        # compute gradient and do SGD step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # for param in model.parameters():\n        #     print(param.requires_grad)\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % args.print_freq == 0:\n            progress.display(i)\n\n    return top1.avg.cpu().numpy()\n\n\ndef validate(val_loader, model, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n\n    # switch to evaluate mode\n    model.eval()\n\n    with torch.no_grad():\n        end = time.time()\n        for i, (images, target) in enumerate(val_loader):\n            if args.gpu is not None:\n                images = images.cuda(args.gpu, non_blocking=True)\n            target = target.cuda(args.gpu, non_blocking=True)\n\n            # compute output\n            output = model(images)\n\n            preds = np.argmax(output.cpu().numpy(), axis=1)\n\n            # measure accuracy and record loss\n            acc1 = accuracy(output, target, topk=(1, ))\n            top1.update(acc1[0].cpu().numpy()[0], images.size(0))\n\n            # measure elapsed time\n            batch_time.update(time.time() - end)\n            end = time.time()\n\n        print('* Acc@1 {top1.avg:.3f} '.format(top1=top1))\n\n    return top1.avg, preds, target.cpu().numpy(), images.cpu().numpy()\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'\n\n\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n    with torch.no_grad():\n        maxk = max(topk)\n        batch_size = target.size(0)\n\n        _, pred = output.topk(maxk, 1, True, True)\n        pred = pred.t()\n        correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n        res = []\n        for k in topk:\n            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n            res.append(correct_k.mul_(100.0 / batch_size))\n        return res\n\n\nif __name__ == '__main__':\n    main()"
  },
  {
    "path": "moco/__init__.py",
    "content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n"
  },
  {
    "path": "moco/builder.py",
    "content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport torch\nimport torch.nn as nn\n\n\nclass MoCo(nn.Module):\n    \"\"\"\n    Build a MoCo model with: a query encoder, a key encoder, and a queue\n    https://arxiv.org/abs/1911.05722\n    \"\"\"\n    def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False):\n        \"\"\"\n        dim: feature dimension (default: 128)\n        K: queue size; number of negative keys (default: 65536)\n        m: moco momentum of updating key encoder (default: 0.999)\n        T: softmax temperature (default: 0.07)\n        \"\"\"\n        super(MoCo, self).__init__()\n\n        self.K = K\n        self.m = m\n        self.T = T\n\n        # create the encoders\n        # num_classes is the output fc dimension\n        self.encoder_q = base_encoder(num_classes=dim)\n        self.encoder_k = base_encoder(num_classes=dim)\n\n        # self.encoder_q.classifier = self.encoder_q.classifier[-1]  # remove dropout (only for mobilenet_v2)\n        # self.encoder_k.classifier = self.encoder_k.classifier[-1]  # remove dropout (only for mobilenet_v2)\n\n        if mlp:  # hack: brute-force replacement\n            dim_mlp = self.encoder_q.fc.weight.shape[1]\n            self.encoder_q.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc)\n            self.encoder_k.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc)\n\n        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n            param_k.data.copy_(param_q.data)  # initialize\n            param_k.requires_grad = False  # not update by gradient\n\n        # create the queue\n        self.register_buffer(\"queue\", torch.randn(dim, K))\n        self.queue = nn.functional.normalize(self.queue, dim=0)\n\n        self.register_buffer(\"queue_ptr\", torch.zeros(1, dtype=torch.long))\n\n    @torch.no_grad()\n    def _momentum_update_key_encoder(self):\n        \"\"\"\n        Momentum update of the key encoder\n        \"\"\"\n        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)\n\n    @torch.no_grad()\n    def _dequeue_and_enqueue(self, keys):\n        # gather keys before updating queue\n        keys = concat_all_gather(keys)\n\n        batch_size = keys.shape[0]\n\n        ptr = int(self.queue_ptr)\n        assert self.K % batch_size == 0  # for simplicity\n\n        # replace the keys at ptr (dequeue and enqueue)\n        self.queue[:, ptr:ptr + batch_size] = keys.T\n        ptr = (ptr + batch_size) % self.K  # move pointer\n\n        self.queue_ptr[0] = ptr\n\n    @torch.no_grad()\n    def _batch_shuffle_ddp(self, x):\n        \"\"\"\n        Batch shuffle, for making use of BatchNorm.\n        *** Only support DistributedDataParallel (DDP) model. ***\n        \"\"\"\n        # gather from all gpus\n        batch_size_this = x.shape[0]\n        x_gather = concat_all_gather(x)\n        batch_size_all = x_gather.shape[0]\n\n        num_gpus = batch_size_all // batch_size_this\n\n        # random shuffle index\n        idx_shuffle = torch.randperm(batch_size_all).cuda()\n\n        # broadcast to all gpus\n        torch.distributed.broadcast(idx_shuffle, src=0)\n\n        # index for restoring\n        idx_unshuffle = torch.argsort(idx_shuffle)\n\n        # shuffled index for this gpu\n        gpu_idx = torch.distributed.get_rank()\n        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]\n\n        return x_gather[idx_this], idx_unshuffle\n\n    @torch.no_grad()\n    def _batch_unshuffle_ddp(self, x, idx_unshuffle):\n        \"\"\"\n        Undo batch shuffle.\n        *** Only support DistributedDataParallel (DDP) model. ***\n        \"\"\"\n        # gather from all gpus\n        batch_size_this = x.shape[0]\n        x_gather = concat_all_gather(x)\n        batch_size_all = x_gather.shape[0]\n\n        num_gpus = batch_size_all // batch_size_this\n\n        # restored index for this gpu\n        gpu_idx = torch.distributed.get_rank()\n        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]\n\n        return x_gather[idx_this]\n\n    def forward(self, im_q, im_k):\n        \"\"\"\n        Input:\n            im_q: a batch of query images\n            im_k: a batch of key images\n        Output:\n            logits, targets\n        \"\"\"\n\n        # compute query features\n        q = self.encoder_q(im_q)  # queries: NxC\n        q = nn.functional.normalize(q, dim=1)\n\n        # compute key features\n        with torch.no_grad():  # no gradient to keys\n            self._momentum_update_key_encoder()  # update the key encoder\n\n            # shuffle for making use of BN\n            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)\n\n            k = self.encoder_k(im_k)  # keys: NxC\n            k = nn.functional.normalize(k, dim=1)\n\n            # undo shuffle\n            k = self._batch_unshuffle_ddp(k, idx_unshuffle)\n\n        # compute logits\n        # Einstein sum is more intuitive\n        # positive logits: Nx1\n        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)\n        # negative logits: NxK\n        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])\n\n        # logits: Nx(1+K)\n        logits = torch.cat([l_pos, l_neg], dim=1)\n\n        # apply temperature\n        logits /= self.T\n\n        # labels: positive key indicators\n        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()\n\n        # dequeue and enqueue\n        self._dequeue_and_enqueue(k)\n\n        return logits, labels\n\n\n# utils\n@torch.no_grad()\ndef concat_all_gather(tensor):\n    \"\"\"\n    Performs all_gather operation on the provided tensors.\n    *** Warning ***: torch.distributed.all_gather has no gradient.\n    \"\"\"\n    tensors_gather = [torch.ones_like(tensor)\n        for _ in range(torch.distributed.get_world_size())]\n    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)\n\n    output = torch.cat(tensors_gather, dim=0)\n    return output\n"
  },
  {
    "path": "moco/loader.py",
    "content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nfrom PIL import ImageFilter\nimport random\n\n\nclass TwoCropsTransform:\n    \"\"\"Take two random crops of one image as the query and key.\"\"\"\n\n    def __init__(self, base_transform):\n        self.base_transform = base_transform\n\n    def __call__(self, x):\n        q = self.base_transform(x)\n        k = self.base_transform(x)\n        return [q, k]\n\n\nclass GaussianBlur(object):\n    \"\"\"Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709\"\"\"\n\n    def __init__(self, sigma=[.1, 2.]):\n        self.sigma = sigma\n\n    def __call__(self, x):\n        sigma = random.uniform(self.sigma[0], self.sigma[1])\n        x = x.filter(ImageFilter.GaussianBlur(radius=sigma))\n        return x\n"
  },
  {
    "path": "moco_img.py",
    "content": "#!/usr/bin/env python\n# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport argparse\nimport builtins\nimport math\nimport os\nimport random\nimport shutil\nimport time\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.optim\nimport torch.multiprocessing as mp\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\n\nimport moco.loader\nimport moco.builder\n\nmodel_names = sorted(name for name in models.__dict__\n    if name.islower() and not name.startswith(\"__\")\n    and callable(models.__dict__[name]))\n\nparser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\nparser.add_argument('data', metavar='DIR',\n                    help='path to dataset')\nparser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50',\n                    choices=model_names,\n                    help='model architecture: ' +\n                        ' | '.join(model_names) +\n                        ' (default: resnet50)')\nparser.add_argument('-j', '--workers', default=32, type=int, metavar='N',\n                    help='number of data loading workers (default: 32)')\nparser.add_argument('--epochs', default=12, type=int, metavar='N',\n                    help='number of total epochs to run')\nparser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n                    help='manual epoch number (useful on restarts)')\nparser.add_argument('-b', '--batch-size', default=256, type=int,\n                    metavar='N',\n                    help='mini-batch size (default: 256), this is the total '\n                         'batch size of all GPUs on the current node when '\n                         'using Data Parallel or Distributed Data Parallel')\nparser.add_argument('--lr', '--learning-rate', default=0.03, type=float,\n                    metavar='LR', help='initial learning rate', dest='lr')\nparser.add_argument('--schedule', default=[11, 20], nargs='*', type=int,\n                    help='learning rate schedule (when to drop lr by 10x)')\nparser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n                    help='momentum of SGD solver')\nparser.add_argument('--wd', '--weight-decay', default=0, type=float,\n                    metavar='W', help='weight decay (default: 0)',\n                    dest='weight_decay')\nparser.add_argument('-p', '--print-freq', default=1000, type=int,\n                    metavar='N', help='print frequency (default: 10)')\nparser.add_argument('--resume', default='', type=str, metavar='PATH',\n                    help='path to latest checkpoint (default: none)')\nparser.add_argument('--world-size', default=-1, type=int,\n                    help='number of nodes for distributed training')\nparser.add_argument('--rank', default=-1, type=int,\n                    help='node rank for distributed training')\nparser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,\n                    help='url used to set up distributed training')\nparser.add_argument('--dist-backend', default='nccl', type=str,\n                    help='distributed backend')\nparser.add_argument('--seed', default=None, type=int,\n                    help='seed for initializing training. ')\nparser.add_argument('--gpu', default=None, type=int,\n                    help='GPU id to use.')\nparser.add_argument('--multiprocessing-distributed', action='store_true',\n                    help='Use multi-processing distributed training to launch '\n                         'N processes per node, which has N GPUs. This is the '\n                         'fastest way to use PyTorch for either single node or '\n                         'multi node data parallel training')\n\n# moco specific configs:\nparser.add_argument('--moco-dim', default=128, type=int,\n                    help='feature dimension (default: 128)')\nparser.add_argument('--moco-k', default=65536, type=int,\n                    help='queue size; number of negative keys (default: 65536)')\nparser.add_argument('--moco-m', default=0.999, type=float,\n                    help='moco momentum of updating key encoder (default: 0.999)')\nparser.add_argument('--moco-t', default=0.07, type=float,\n                    help='softmax temperature (default: 0.07)')\n\n# options for moco v2\nparser.add_argument('--mlp', action='store_true',\n                    help='use mlp head')\nparser.add_argument('--aug-plus', action='store_true',\n                    help='use moco v2 data augmentation')\nparser.add_argument('--cos', action='store_true',\n                    help='use cosine lr schedule')\n\n\ndef main():\n    args = parser.parse_args()\n\n    if args.seed is not None:\n        random.seed(args.seed)\n        torch.manual_seed(args.seed)\n        cudnn.deterministic = True\n        warnings.warn('You have chosen to seed training. '\n                      'This will turn on the CUDNN deterministic setting, '\n                      'which can slow down your training considerably! '\n                      'You may see unexpected behavior when restarting '\n                      'from checkpoints.')\n\n    if args.gpu is not None:\n        warnings.warn('You have chosen a specific GPU. This will completely '\n                      'disable data parallelism.')\n\n    if args.dist_url == \"env://\" and args.world_size == -1:\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    args.distributed = args.world_size > 1 or args.multiprocessing_distributed\n\n    ngpus_per_node = torch.cuda.device_count()\n    if args.multiprocessing_distributed:\n        # Since we have ngpus_per_node processes per node, the total world_size\n        # needs to be adjusted accordingly\n        args.world_size = ngpus_per_node * args.world_size\n        # Use torch.multiprocessing.spawn to launch distributed processes: the\n        # main_worker process function\n        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))\n    else:\n        # Simply call main_worker function\n        main_worker(args.gpu, ngpus_per_node, args)\n\n\ndef main_worker(gpu, ngpus_per_node, args):\n\n    args.gpu = gpu\n\n    print(args)\n\n    # suppress printing if not master\n    if args.multiprocessing_distributed and args.gpu != 0:\n        def print_pass(*args):\n            pass\n        builtins.print = print_pass\n\n    if args.gpu is not None:\n        print(\"Use GPU: {} for training\".format(args.gpu))\n\n    if args.distributed:\n        if args.dist_url == \"env://\" and args.rank == -1:\n            args.rank = int(os.environ[\"RANK\"])\n        if args.multiprocessing_distributed:\n            # For multiprocessing distributed training, rank needs to be the\n            # global rank among all the processes\n            args.rank = args.rank * ngpus_per_node + gpu\n        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,\n                                world_size=args.world_size, rank=args.rank)\n    # create model\n    print(\"=> creating model '{}'\".format(args.arch))\n    model = moco.builder.MoCo(\n        models.__dict__[args.arch],\n        args.moco_dim, args.moco_k, args.moco_m, args.moco_t, args.mlp)\n    print(model)\n\n    if args.distributed:\n        # For multiprocessing distributed, DistributedDataParallel constructor\n        # should always set the single device scope, otherwise,\n        # DistributedDataParallel will use all available devices.\n        if args.gpu is not None:\n            torch.cuda.set_device(args.gpu)\n            model.cuda(args.gpu)\n            # When using a single GPU per process and per\n            # DistributedDataParallel, we need to divide the batch size\n            # ourselves based on the total number of GPUs we have\n            args.batch_size = int(args.batch_size / ngpus_per_node)\n            args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])\n        else:\n            model.cuda()\n            # DistributedDataParallel will divide and allocate batch_size to all\n            # available GPUs if device_ids are not set\n            model = torch.nn.parallel.DistributedDataParallel(model)\n    elif args.gpu is not None:\n        torch.cuda.set_device(args.gpu)\n        model = model.cuda(args.gpu)\n        # comment out the following line for debugging\n        raise NotImplementedError(\"Only DistributedDataParallel is supported.\")\n    else:\n        # AllGather implementation (batch shuffle, queue update, etc.) in\n        # this code only supports DistributedDataParallel.\n        raise NotImplementedError(\"Only DistributedDataParallel is supported.\")\n\n    # define loss function (criterion) and optimizer\n    criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n\n    optimizer = torch.optim.SGD(model.parameters(), args.lr,\n                                momentum=args.momentum,\n                                weight_decay=args.weight_decay)\n\n    # optionally resume from a checkpoint\n    if args.resume:\n        if os.path.isfile(args.resume):\n            print(\"=> loading checkpoint '{}'\".format(args.resume))\n            if args.gpu is None:\n                checkpoint = torch.load(args.resume)\n            else:\n                # Map model to be loaded to specified single gpu.\n                loc = 'cuda:{}'.format(args.gpu)\n                checkpoint = torch.load(args.resume, map_location=loc)\n            args.start_epoch = checkpoint['epoch']\n            model.load_state_dict(checkpoint['state_dict'])\n            optimizer.load_state_dict(checkpoint['optimizer'])\n            print(\"=> loaded checkpoint '{}' (epoch {})\"\n                  .format(args.resume, checkpoint['epoch']))\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.resume))\n\n    cudnn.benchmark = True\n\n    # Data loading code\n    traindir = args.data\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                     std=[0.229, 0.224, 0.225])\n    if args.aug_plus:\n        # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709\n        augmentation = [\n            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n            transforms.RandomApply([\n                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened\n            ], p=0.8),\n            transforms.RandomGrayscale(p=0.2),\n            transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.5),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            normalize\n        ]\n    else:\n        # MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978\n        augmentation = [\n            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n            transforms.RandomGrayscale(p=0.2),\n            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            normalize\n        ]\n\n    train_dataset = datasets.ImageFolder(\n        traindir,\n        moco.loader.TwoCropsTransform(transforms.Compose(augmentation)))\n\n    print('Dataset size:', len(train_dataset))\n\n    if args.distributed:\n        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)\n    else:\n        train_sampler = None\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),\n        num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True)\n\n    print('Starting training ...')\n    for epoch in range(args.start_epoch, args.epochs):\n        if args.distributed:\n            train_sampler.set_epoch(epoch)\n        adjust_learning_rate(optimizer, epoch, args)\n\n        print('Start of epoch ', epoch)\n\n        # train for one epoch\n        train(train_loader, model, criterion, optimizer, epoch, args)\n\n        if not args.multiprocessing_distributed or (args.multiprocessing_distributed\n                and args.rank % ngpus_per_node == 0):\n            save_checkpoint({\n                'epoch': epoch + 1,\n                'arch': args.arch,\n                'state_dict': model.state_dict(),\n                'optimizer' : optimizer.state_dict(),\n            }, is_best=False, filename='moco_img_checkpoint_{:04d}.pth.tar'.format(epoch))\n\n\ndef train(train_loader, model, criterion, optimizer, epoch, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    data_time = AverageMeter('Data', ':6.3f')\n    losses = AverageMeter('Loss', ':.4e')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses, top1, top5],\n        prefix=\"Epoch: [{}]\".format(epoch))\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    for i, (images, _) in enumerate(train_loader):\n\n        # measure data loading time\n        data_time.update(time.time() - end)\n\n        if args.gpu is not None:\n            images[0] = images[0].cuda(args.gpu, non_blocking=True)\n            images[1] = images[1].cuda(args.gpu, non_blocking=True)\n\n        # compute output\n        output, target = model(im_q=images[0], im_k=images[1])\n        loss = criterion(output, target)\n\n        # acc1/acc5 are (K+1)-way contrast classifier accuracy\n        # measure accuracy and record loss\n        acc1, acc5 = accuracy(output, target, topk=(1, 5))\n        losses.update(loss.item(), images[0].size(0))\n        top1.update(acc1[0], images[0].size(0))\n        top5.update(acc5[0], images[0].size(0))\n\n        # compute gradient and do SGD step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % args.print_freq == 0:\n            progress.display(i)\n\n\ndef save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):\n    torch.save(state, filename)\n    if is_best:\n        shutil.copyfile(filename, 'model_best.pth.tar')\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'\n\n\ndef adjust_learning_rate(optimizer, epoch, args):\n    \"\"\"Decay the learning rate based on schedule\"\"\"\n    lr = args.lr\n    if args.cos:  # cosine lr schedule\n        lr *= 0.5 * (1. + math.cos(math.pi * epoch / args.epochs))\n    else:  # stepwise lr schedule\n        for milestone in args.schedule:\n            lr *= 0.1 if epoch >= milestone else 1.\n    for param_group in optimizer.param_groups:\n        param_group['lr'] = lr\n\n\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n    with torch.no_grad():\n        maxk = max(topk)\n        batch_size = target.size(0)\n\n        _, pred = output.topk(maxk, 1, True, True)\n        pred = pred.t()\n        correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n        res = []\n        for k in topk:\n            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n            res.append(correct_k.mul_(100.0 / batch_size))\n        return res\n\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "moco_temp.py",
    "content": "#!/usr/bin/env python\n# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport argparse\nimport builtins\nimport math\nimport os\nimport random\nimport shutil\nimport time\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.optim\nimport torch.multiprocessing as mp\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\n\nimport moco.loader\nimport moco.builder\nfrom moco_utils import DistributedProxySampler, ContrastiveBatchSampler\n\nmodel_names = sorted(name for name in models.__dict__\n    if name.islower() and not name.startswith(\"__\")\n    and callable(models.__dict__[name]))\n\nparser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\nparser.add_argument('data', metavar='DIR',\n                    help='path to dataset')\nparser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50',\n                    choices=model_names,\n                    help='model architecture: ' +\n                        ' | '.join(model_names) +\n                        ' (default: resnet50)')\nparser.add_argument('-j', '--workers', default=32, type=int, metavar='N',\n                    help='number of data loading workers (default: 32)')\nparser.add_argument('--epochs', default=12, type=int, metavar='N',\n                    help='number of total epochs to run')\nparser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n                    help='manual epoch number (useful on restarts)')\nparser.add_argument('-b', '--batch-size', default=256, type=int,\n                    metavar='N',\n                    help='mini-batch size (default: 256), this is the total '\n                         'batch size of all GPUs on the current node when '\n                         'using Data Parallel or Distributed Data Parallel')\nparser.add_argument('--lr', '--learning-rate', default=0.03, type=float,\n                    metavar='LR', help='initial learning rate', dest='lr')\nparser.add_argument('--schedule', default=[11, 20], nargs='*', type=int,\n                    help='learning rate schedule (when to drop lr by 10x)')\nparser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n                    help='momentum of SGD solver')\nparser.add_argument('--wd', '--weight-decay', default=0, type=float,\n                    metavar='W', help='weight decay (default: 0)',\n                    dest='weight_decay')\nparser.add_argument('-p', '--print-freq', default=1000, type=int,\n                    metavar='N', help='print frequency (default: 10)')\nparser.add_argument('--resume', default='', type=str, metavar='PATH',\n                    help='path to latest checkpoint (default: none)')\nparser.add_argument('--world-size', default=-1, type=int,\n                    help='number of nodes for distributed training')\nparser.add_argument('--rank', default=-1, type=int,\n                    help='node rank for distributed training')\nparser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,\n                    help='url used to set up distributed training')\nparser.add_argument('--dist-backend', default='nccl', type=str,\n                    help='distributed backend')\nparser.add_argument('--seed', default=None, type=int,\n                    help='seed for initializing training. ')\nparser.add_argument('--gpu', default=None, type=int,\n                    help='GPU id to use.')\nparser.add_argument('--multiprocessing-distributed', action='store_true',\n                    help='Use multi-processing distributed training to launch '\n                         'N processes per node, which has N GPUs. This is the '\n                         'fastest way to use PyTorch for either single node or '\n                         'multi node data parallel training')\n\n# moco specific configs:\nparser.add_argument('--moco-dim', default=128, type=int,\n                    help='feature dimension (default: 128)')\nparser.add_argument('--moco-k', default=65536, type=int,\n                    help='queue size; number of negative keys (default: 65536)')\nparser.add_argument('--moco-m', default=0.999, type=float,\n                    help='moco momentum of updating key encoder (default: 0.999)')\nparser.add_argument('--moco-t', default=0.07, type=float,\n                    help='softmax temperature (default: 0.07)')\n\n# options for moco v2\nparser.add_argument('--mlp', action='store_true',\n                    help='use mlp head')\nparser.add_argument('--aug-plus', action='store_true',\n                    help='use moco v2 data augmentation')\nparser.add_argument('--cos', action='store_true',\n                    help='use cosine lr schedule')\n\n\ndef main():\n    args = parser.parse_args()\n\n    if args.seed is not None:\n        random.seed(args.seed)\n        torch.manual_seed(args.seed)\n        cudnn.deterministic = True\n        warnings.warn('You have chosen to seed training. '\n                      'This will turn on the CUDNN deterministic setting, '\n                      'which can slow down your training considerably! '\n                      'You may see unexpected behavior when restarting '\n                      'from checkpoints.')\n\n    if args.gpu is not None:\n        warnings.warn('You have chosen a specific GPU. This will completely '\n                      'disable data parallelism.')\n\n    if args.dist_url == \"env://\" and args.world_size == -1:\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    args.distributed = args.world_size > 1 or args.multiprocessing_distributed\n\n    ngpus_per_node = torch.cuda.device_count()\n    if args.multiprocessing_distributed:\n        # Since we have ngpus_per_node processes per node, the total world_size\n        # needs to be adjusted accordingly\n        args.world_size = ngpus_per_node * args.world_size\n        # Use torch.multiprocessing.spawn to launch distributed processes: the\n        # main_worker process function\n        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))\n    else:\n        # Simply call main_worker function\n        main_worker(args.gpu, ngpus_per_node, args)\n\n\ndef main_worker(gpu, ngpus_per_node, args):\n\n    args.gpu = gpu\n\n    # suppress printing if not master\n    if args.multiprocessing_distributed and args.gpu != 0:\n        def print_pass(*args):\n            pass\n        builtins.print = print_pass\n\n    if args.gpu is not None:\n        print(\"Use GPU: {} for training\".format(args.gpu))\n\n    if args.distributed:\n        if args.dist_url == \"env://\" and args.rank == -1:\n            args.rank = int(os.environ[\"RANK\"])\n        if args.multiprocessing_distributed:\n            # For multiprocessing distributed training, rank needs to be the\n            # global rank among all the processes\n            args.rank = args.rank * ngpus_per_node + gpu\n        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,\n                                world_size=args.world_size, rank=args.rank)\n    # create model\n    print(\"=> creating model '{}'\".format(args.arch))\n    model = moco.builder.MoCo(\n        models.__dict__[args.arch],\n        args.moco_dim, args.moco_k, args.moco_m, args.moco_t, args.mlp)\n    print(model)\n\n    if args.distributed:\n        # For multiprocessing distributed, DistributedDataParallel constructor\n        # should always set the single device scope, otherwise,\n        # DistributedDataParallel will use all available devices.\n        if args.gpu is not None:\n            torch.cuda.set_device(args.gpu)\n            model.cuda(args.gpu)\n            # When using a single GPU per process and per\n            # DistributedDataParallel, we need to divide the batch size\n            # ourselves based on the total number of GPUs we have\n            args.batch_size = int(args.batch_size / ngpus_per_node)\n            args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])\n        else:\n            model.cuda()\n            # DistributedDataParallel will divide and allocate batch_size to all\n            # available GPUs if device_ids are not set\n            model = torch.nn.parallel.DistributedDataParallel(model)\n    elif args.gpu is not None:\n        torch.cuda.set_device(args.gpu)\n        model = model.cuda(args.gpu)\n        # comment out the following line for debugging\n        raise NotImplementedError(\"Only DistributedDataParallel is supported.\")\n    else:\n        # AllGather implementation (batch shuffle, queue update, etc.) in\n        # this code only supports DistributedDataParallel.\n        raise NotImplementedError(\"Only DistributedDataParallel is supported.\")\n\n    # define loss function (criterion) and optimizer\n    criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n\n    optimizer = torch.optim.SGD(model.parameters(), args.lr,\n                                momentum=args.momentum,\n                                weight_decay=args.weight_decay)\n\n    # optionally resume from a checkpoint\n    if args.resume:\n        if os.path.isfile(args.resume):\n            print(\"=> loading checkpoint '{}'\".format(args.resume))\n            if args.gpu is None:\n                checkpoint = torch.load(args.resume)\n            else:\n                # Map model to be loaded to specified single gpu.\n                loc = 'cuda:{}'.format(args.gpu)\n                checkpoint = torch.load(args.resume, map_location=loc)\n            args.start_epoch = checkpoint['epoch']\n            model.load_state_dict(checkpoint['state_dict'])\n            optimizer.load_state_dict(checkpoint['optimizer'])\n            print(\"=> loaded checkpoint '{}' (epoch {})\"\n                  .format(args.resume, checkpoint['epoch']))\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.resume))\n\n    cudnn.benchmark = True\n\n    # Data loading code\n    traindir = args.data\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                     std=[0.229, 0.224, 0.225])\n    if args.aug_plus:\n        # MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709\n        augmentation = [\n            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n            transforms.RandomApply([\n                transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened\n            ], p=0.8),\n            transforms.RandomGrayscale(p=0.2),\n            transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.5),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            normalize\n        ]\n    else:\n        # MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978\n        augmentation = [\n            transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n            transforms.RandomGrayscale(p=0.2),\n            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),\n            transforms.RandomHorizontalFlip(),\n            transforms.ToTensor(),\n            normalize\n        ]\n\n    train_dataset = datasets.ImageFolder(\n        traindir, transforms.Compose(augmentation))\n\n    print('Dataset size:', len(train_dataset))\n\n    if args.distributed:\n        train_sampler = DistributedProxySampler(ContrastiveBatchSampler(\n            train_dataset, args.batch_size, 1, False))\n    else:\n        train_sampler = None\n\n    train_loader = torch.utils.data.DataLoader(train_dataset, shuffle=(train_sampler is None),\n                                               num_workers=args.workers, pin_memory=True,\n                                               batch_sampler=train_sampler)\n\n    print('Starting training ...')\n    for epoch in range(args.start_epoch, args.epochs):\n        if args.distributed:\n            train_sampler.set_epoch(epoch)\n        adjust_learning_rate(optimizer, epoch, args)\n\n        print('Start of epoch ', epoch)\n\n        # train for one epoch\n        train(train_loader, model, criterion, optimizer, epoch, args)\n\n        if not args.multiprocessing_distributed or (args.multiprocessing_distributed\n                and args.rank % ngpus_per_node == 0):\n            save_checkpoint({\n                'epoch': epoch + 1,\n                'arch': args.arch,\n                'state_dict': model.state_dict(),\n                'optimizer' : optimizer.state_dict(),\n            }, is_best=False, filename='moco_temp_checkpoint_{:04d}.pth.tar'.format(epoch))\n\n\ndef train(train_loader, model, criterion, optimizer, epoch, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    data_time = AverageMeter('Data', ':6.3f')\n    losses = AverageMeter('Loss', ':.4e')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses, top1, top5],\n        prefix=\"Epoch: [{}]\".format(epoch))\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    for i, (images, _) in enumerate(train_loader):\n\n        # measure data loading time\n        data_time.update(time.time() - end)\n\n        if args.gpu is not None:\n            images_0 = images[:images.size(0)//2].cuda(args.gpu, non_blocking=True)\n            images_1 = images[images.size(0)//2:].cuda(args.gpu, non_blocking=True)\n\n        # compute output\n        output, target = model(im_q=images_0, im_k=images_1)\n        loss = criterion(output, target)\n\n        # acc1/acc5 are (K+1)-way contrast classifier accuracy\n        # measure accuracy and record loss\n        acc1, acc5 = accuracy(output, target, topk=(1, 5))\n        losses.update(loss.item(), images_0.size(0))\n        top1.update(acc1[0], images_0.size(0))\n        top5.update(acc5[0], images_0.size(0))\n\n        # compute gradient and do SGD step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % args.print_freq == 0:\n            progress.display(i)\n\n\ndef save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):\n    torch.save(state, filename)\n    if is_best:\n        shutil.copyfile(filename, 'model_best.pth.tar')\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'\n\n\ndef adjust_learning_rate(optimizer, epoch, args):\n    \"\"\"Decay the learning rate based on schedule\"\"\"\n    lr = args.lr\n    if args.cos:  # cosine lr schedule\n        lr *= 0.5 * (1. + math.cos(math.pi * epoch / args.epochs))\n    else:  # stepwise lr schedule\n        for milestone in args.schedule:\n            lr *= 0.1 if epoch >= milestone else 1.\n    for param_group in optimizer.param_groups:\n        param_group['lr'] = lr\n\n\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n    with torch.no_grad():\n        maxk = max(topk)\n        batch_size = target.size(0)\n\n        _, pred = output.topk(maxk, 1, True, True)\n        pred = pred.t()\n        correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n        res = []\n        for k in topk:\n            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n            res.append(correct_k.mul_(100.0 / batch_size))\n        return res\n\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "moco_utils.py",
    "content": "# Defines some util functions\nimport torch\nfrom torch.utils.data import Sampler\nfrom torch.utils.data.distributed import DistributedSampler\n\n\nclass DistributedProxySampler(DistributedSampler):\n    \"\"\"Sampler that restricts data loading to a subset of input sampler indices.\n\n    It is especially useful in conjunction with\n    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each\n    process can pass a DistributedSampler instance as a DataLoader sampler,\n    and load a subset of the original dataset that is exclusive to it.\n\n    .. note::\n        Input sampler is assumed to be of constant size.\n\n    Arguments:\n        sampler: Input data sampler.\n        num_replicas (optional): Number of processes participating in\n            distributed training.\n        rank (optional): Rank of the current process within num_replicas.\n    \"\"\"\n\n    def __init__(self, sampler, num_replicas=None, rank=None):\n        super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)\n        self.sampler = sampler\n\n    def __iter__(self):\n        # deterministically shuffle based on epoch\n        torch.manual_seed(self.epoch)\n        indices = list(self.sampler)\n\n        # add extra samples to make it evenly divisible\n        indices += indices[:(self.total_size - len(indices))]\n        if len(indices) != self.total_size:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.total_size))\n\n        # subsample\n        indices = indices[self.rank:self.total_size:self.num_replicas]\n        if len(indices) != self.num_samples:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.num_samples))\n\n        return iter(indices)\n\n    def set_epoch(self, epoch):\n        self.epoch = epoch\n\nclass ContrastiveBatchSampler(Sampler):\n    def __init__(self, data_source, batch_size, pos_window, drop_last):\n        self.data_source = data_source\n        self.batch_size = batch_size\n        self.pos_window = pos_window\n        self.drop_last = drop_last\n        self.n = len(self.data_source)\n\n    def __iter__(self):\n        for i in range(self.n // self.batch_size):\n            x = torch.randint(low=0, high=self.n-1, size=(self.batch_size//2,),\n                              dtype=torch.int64)\n            y = x + torch.randint(low=-self.pos_window, high=self.pos_window, size=(self.batch_size//2,),\n                                  dtype=torch.int64)\n            y = torch.clamp(y, 0, self.n-1)\n            z = x.tolist() + y.tolist()\n\n            yield z\n\n    def __len__(self):\n        if self.drop_last:\n            return self.n // self.batch_size\n        else:\n            return (self.n + self.batch_size - 1) // self.batch_size"
  },
  {
    "path": "read_saycam.py",
    "content": "import os\nimport sys\nimport argparse\nimport cv2\nimport numpy as np\n\n\nparser = argparse.ArgumentParser(description='Read SAYCam videos')\nparser.add_argument('data', metavar='DIR', help='path to SAYCam videos')\nparser.add_argument('--save-dir', default='', type=str, help='save directory')\nparser.add_argument('--fps', default=5, type=int, help='sampling rate (frames per second)')\nparser.add_argument('--seg-len', default=288, type=int, help='segment length (seconds)')\n\n\nif __name__ == '__main__':\n\n    args = parser.parse_args()\n\n    file_list = os.listdir(args.data)\n    file_list.sort()\n\n    class_counter = 0\n    img_counter = 0\n    file_counter = 0\n\n    final_size = 224\n    resized_minor_length = 256\n    edge_filter = False\n    n_imgs_per_class = args.seg_len * args.fps\n\n    curr_dir_name = os.path.join(args.save_dir, 'class_{:04d}'.format(class_counter))\n    os.mkdir(curr_dir_name)\n\n    for file_indx in file_list:\n        file_name = os.path.join(args.data, file_indx)\n\n        cap = cv2.VideoCapture(file_name)\n        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n        frame_rate = int(cap.get(cv2.CAP_PROP_FPS))\n\n        # take every sample_rate frames (30: 1fps, 15: 2fps, 10: 3fps, 6: 5fps, 5: 6fps, 3: 10fps, 2: 15fps, 1: 30fps)\n        sample_rate = frame_rate // args.fps + 1\n\n        print('Total frame count: ', frame_count)\n        print('Native frame rate: ', frame_rate)\n\n        fc = 0\n        ret = True\n\n        # Resize\n        new_height = frame_height * resized_minor_length // min(frame_height, frame_width)\n        new_width = frame_width * resized_minor_length // min(frame_height, frame_width)\n\n        while (fc < frame_count):\n\n            ret, frame = cap.read()\n\n            if fc % sample_rate == 0 and ret:\n\n                # Resize\n                resized_frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_CUBIC)\n\n                # Crop\n                height, width, _ = resized_frame.shape\n                startx = width // 2 - (final_size // 2)\n                starty = height // 2 - (final_size // 2) - 16\n                cropped_frame = resized_frame[starty:starty + final_size, startx:startx + final_size]\n                assert cropped_frame.shape[0] == final_size and cropped_frame.shape[1] == final_size, \\\n                    (cropped_frame.shape, height, width)\n\n                if edge_filter:\n                    cropped_frame = cv2.Laplacian(cropped_frame, cv2.CV_64F, ksize=5)\n                    img_min = cropped_frame.min()\n                    img_max = cropped_frame.max()\n                    cropped_frame = np.uint8(255 * (cropped_frame - img_min) / (img_max - img_min))\n\n                cv2.imwrite(os.path.join(curr_dir_name, 'img_{:04d}.jpeg'.format(img_counter)), cropped_frame[::-1, ::-1, :])\n                img_counter += 1\n\n                if img_counter == n_imgs_per_class:\n                    img_counter = 0\n                    class_counter += 1\n                    curr_dir_name = os.path.join(args.save_dir, 'class_{:04d}'.format(class_counter))\n                    os.mkdir(curr_dir_name)\n\n            fc += 1\n\n        cap.release()\n\n        file_counter += 1\n        print('Completed video {:4d} of {:4d}'.format(file_counter, len(file_list)))"
  },
  {
    "path": "scripts/feature_animation.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=8\n#SBATCH --gres=gpu:1080ti:2\n#SBATCH --mem=150GB\n#SBATCH --time=1:00:00\n#SBATCH --array=0\n#SBATCH --job-name=feature_animation\n#SBATCH --output=feature_animation_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/feature_animation.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/feature_animation_imgs_intphys/' --model-path '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/self_supervised_models/TC-SAY.tar' --batch-size 900 --n_out 6269 --feature-idx 600\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/feature_animation_class.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --exclude=hpc1,hpc2,hpc3,hpc4,hpc5,hpc6,hpc7,hpc8,hpc9,vine3,vine4,vine6,vine11,vine12,lion17,rose7,rose8,rose9\n#SBATCH --ntasks=1\n#SBATCH --gres=gpu:4\n#SBATCH --mem=100GB\n#SBATCH --time=1:00:00\n#SBATCH --array=0\n#SBATCH --job-name=feature_animation\n#SBATCH --output=feature_animation_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/feature_animation_class.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/feature_animation_computers/' --model-path '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/self_supervised_models/TC-S_labeledS_5_iid.tar' --batch-size 200 --n_out 26 --class-idx 6\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/highly_activating_imgs.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=1\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=1\n#SBATCH --mem=16GB\n#SBATCH --time=1:00:00\n#SBATCH --array=0\n#SBATCH --job-name=activating_imgs\n#SBATCH --output=activating_imgs_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/highly_activating_imgs.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --n_out 26 --model-path 'mobilenetV2_S_5fps_2000cls_coloraug_labeled.tar' \n\necho \"Done\"\n"
  },
  {
    "path": "scripts/hog_baseline.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=1\n#SBATCH --cpus-per-task=1\n#SBATCH --mem=100GB\n#SBATCH --time=6:00:00\n#SBATCH --array=0\n#SBATCH --job-name=hog\n#SBATCH --output=hog_%A_%a.out\n\nmodule purge\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/hog_baseline.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/toybox_1fps/'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/imagenet_finetuning.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=16\n#SBATCH --gres=gpu:titanrtx:4\n#SBATCH --mem=150GB\n#SBATCH --time=48:00:00\n#SBATCH --array=0\n#SBATCH --job-name=finetune_imgnet\n#SBATCH --output=finetune_imgnet_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --freeze-trunk --n_out 6269 --resume 'resnext50_32x4d_augmentstrong_batch256_True_SAY_5_288_epoch_15.tar'\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --freeze-trunk --n_out 2765 --resume 'resnext50_32x4d_augmentstrong_batch256_True_S_5_288_epoch_10.tar'\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --freeze-trunk --n_out 1786 --resume 'resnext50_32x4d_augmentstrong_batch256_True_A_5_288_epoch_10.tar'\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --freeze-trunk --n_out 1718 --resume 'resnext50_32x4d_augmentstrong_batch256_True_Y_5_288_epoch_10.tar'\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --frac-retained 0.01 --n_out 6269 --resume 'resnext50_32x4d_augmentstrong_batch256_True_SAY_5_288_epoch_15.tar'\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/imagenet_finetuning.py --freeze-trunk --n_out 1000 --resume 'ft_IN_resnext50_32x4d_augmentstrong_batch256_True_SAY_5_288_epoch_15.tar'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/linear_combination_maps.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=1\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=1\n#SBATCH --mem=16GB\n#SBATCH --time=1:00:00\n#SBATCH --array=0\n#SBATCH --job-name=linear_maps\n#SBATCH --output=linear_maps_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_combination_maps.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --n_out 26 --model-path 'mobilenetV2_S_5fps_2000cls_coloraug_labeled.tar' \n\necho \"Done\"\n"
  },
  {
    "path": "scripts/linear_decoding.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=8\n#SBATCH --gres=gpu:titanrtx:2\n#SBATCH --mem=150GB\n#SBATCH --time=12:00:00\n#SBATCH --array=0\n#SBATCH --job-name=linear_decoding\n#SBATCH --output=linear_decoding_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby_vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --model-name 'random' --num-classes 26 --subsample\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --model-name 'moco_img_0005' --num-classes 26 --subsample\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --model-name 'moco_temp_0005' --num-classes 26 --subsample \n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_5/' --model-name 'TC-S' --num-outs 2765\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --model-name 'mobilenetV2_A_5fps_2000cls_coloraug' --num-outs 1786\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/' --model-name 'mobilenetV2_Y_5fps_2000cls_coloraug' --num-outs 1718\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_5/' --model-name 'TC-SAY' --num-outs 6269\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/linear_decoding.py '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_5/' --model-name 'TC-S' --num-outs 2765\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/moco_img.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=16\n#SBATCH --gres=gpu:titanrtx:4\n#SBATCH --mem=150GB\n#SBATCH --time=48:00:00\n#SBATCH --array=0\n#SBATCH --job-name=moco_img\n#SBATCH --output=moco_img_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/moco_img.py \\\n  -a resnext50_32x4d \\\n  --lr 0.015 \\\n  --batch-size 256 \\\n  --mlp \\\n  --moco-t 0.2 \\\n  --aug-plus --cos \\\n  --dist-url 'tcp://localhost:10001' \\\n  --multiprocessing-distributed \\\n  --world-size 1 --rank 0 \\\n  --start-epoch 0 \\\n  --resume '' \\\n  '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_data_5fps_2000cls_pytorch/'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/moco_temp.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=16\n#SBATCH --gres=gpu:v100:4\n#SBATCH --mem=150GB\n#SBATCH --time=48:00:00\n#SBATCH --array=0\n#SBATCH --job-name=moco_temp\n#SBATCH --output=moco_temp_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/moco_temp.py \\\n  -a resnext50_32x4d \\\n  --lr 0.015 \\\n  --batch-size 256 \\\n  --mlp \\\n  --moco-t 0.2 \\\n  --aug-plus --cos \\\n  --dist-url 'tcp://localhost:10001' \\\n  --multiprocessing-distributed \\\n  --world-size 1 --rank 0 \\\n  --start-epoch 0 \\\n  --resume '' \\\n  '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_data_5fps_2000cls_pytorch/'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/read_saycam.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=1\n#SBATCH --cpus-per-task=1\n#SBATCH --mem=16GB\n#SBATCH --time=48:00:00\n#SBATCH --array=0\n#SBATCH --job-name=read_saycam\n#SBATCH --output=read_saycam_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/read_saycam.py \\\n  --save-dir '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_15fps_288s' \\\n  --fps 15 \\\n  --seg-len 288 \\\n  '/misc/vlgscratch4/LakeGroup/emin/headcam/data_2/S'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/selectivities.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=1\n#SBATCH --gres=gpu:2\n#SBATCH --mem=64GB\n#SBATCH --time=1:00:00\n#SBATCH --array=0\n#SBATCH --job-name=selectivity\n#SBATCH --output=selectivity_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/selectivities.py \\\n--n_out 1000 \\\n--model-path '' \\\n--layer 18 \\\n'/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_clean_labeled_data_1fps_4/'\n\necho \"Done\"\n"
  },
  {
    "path": "scripts/temporal_classification.sh",
    "content": "#!/bin/bash\n\n#SBATCH --nodes=1\n#SBATCH --ntasks=16\n#SBATCH --gres=gpu:v100:4\n#SBATCH --mem=150GB\n#SBATCH --time=48:00:00\n#SBATCH --array=0\n#SBATCH --job-name=tempclas\n#SBATCH --output=tempclas_%A_%a.out\n\nmodule purge\nmodule load cuda-10.1\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/temporal_classification.py --model 'resnext50_32x4d' --n_out 6269 --resume 'resnext50_32x4d_augmentstrong_batch256_True_SAY_5_288_epoch_15.tar' --start-epoch 16 '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/SAY_data_5fps_2000cls_pytorch/'\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/temporal_classification.py --model 'resnext50_32x4d' --n_out 2765 '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/S_data_5fps_2000cls_pytorch/'\n\n#python -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/temporal_classification.py --model 'resnext50_32x4d' --n_out 1786 '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/A_data_5fps_2000cls_pytorch/'\n\npython -u /misc/vlgscratch4/LakeGroup/emin/baby-vision/temporal_classification.py --model 'resnext50_32x4d' --n_out 1718 '/misc/vlgscratch4/LakeGroup/emin/headcam/preprocessing/Y_data_5fps_2000cls_pytorch/'\n\necho \"Done\"\n"
  },
  {
    "path": "selectivities.py",
    "content": "'''Measure single feature class selectivities'''\nimport os\nimport argparse\nimport numpy as np\nimport torch\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom torchvision.utils import make_grid\n\n\ndef extract_map_layer_7x7(mobilenetV2_model):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_model = torch.nn.Sequential(*layer_list)\n    return new_model\n\ndef extract_map_layer_14x14(mobilenetV2_model, layer):\n    layer_list = list(mobilenetV2_model.module.features.children())\n    new_layer_list = layer_list[:-layer]\n    new_layer_list.append(layer_list[-layer].conv[0])\n    new_model = torch.nn.Sequential(*new_layer_list)\n    return new_model\n\ndef load_model(args):\n    model = models.mobilenet_v2(pretrained=True)\n    model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n    model = torch.nn.DataParallel(model).cuda()\n\n    if args.model_path:\n        if os.path.isfile(args.model_path):\n            checkpoint = torch.load(args.model_path)\n            model.load_state_dict(checkpoint['model_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.model_path))\n\n    return model\n\ndef load_data(data_dir, args):\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    train_dataset = datasets.ImageFolder(\n        data_dir,\n        transforms.Compose([transforms.ToTensor(), normalize])\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=False,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    return train_loader\n\ndef predict(data_loader, model):\n\n    targets = []\n    preds = []\n\n    # switch to evaluate mode\n    model.eval()\n\n    with torch.no_grad():\n        for i, (images, target) in enumerate(data_loader):\n            images = images.cuda()\n\n            # compute predictions\n            pred = model(images)\n            pred = torch.mean(pred, dim=(2, 3))\n\n            targets.append(target.cpu().numpy())\n            preds.append(pred.cpu().numpy())\n\n            print('Iter:', i)\n\n    targets = np.concatenate(targets, axis=0)\n    preds = np.concatenate(preds, axis=0)\n\n    print('Targets size:', targets.shape)\n    print('Preds size:', preds.shape)\n\n    return targets, preds\n\n\nif __name__ == '__main__':\n\n    parser = argparse.ArgumentParser(description='Measure single feature class selectivities')\n    parser.add_argument('data', metavar='DIR', help='path to dataset')\n    parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 4)')\n    parser.add_argument('--batch-size', default=580, type=int, help='mini-batch size, this is the total '\n                                                                    'batch size of all GPUs on the current node when '\n                                                                    'using Data Parallel or Distributed Data Parallel')\n    parser.add_argument('--model-path', default='', type=str, help='path to model checkpoint '\n                                                                   '(default: ImageNet-pretrained)')\n    parser.add_argument('--n_out', default=1000, type=int, help='output dim of pre-trained model')\n    parser.add_argument('--layer', default=1, type=int, choices=[1, 2, 6, 10, 14, 18], help='which layer?')\n\n    args = parser.parse_args()\n\n    model = load_model(args)\n    if args.layer == 1:\n        map_layer = extract_map_layer_7x7(model)\n    else:\n        map_layer = extract_map_layer_14x14(model, args.layer)\n\n    data_loader = load_data(args.data, args)\n    targets, preds = predict(data_loader, map_layer)\n\n    n_classes = 26\n    n_neurons = preds.shape[1]\n    class_matrix_mean = np.zeros((n_neurons, n_classes))\n    class_matrix_std = np.zeros((n_neurons, n_classes))\n\n    for i in range(n_neurons):\n        for j in range(n_classes):\n            aux_vec = preds[targets==j, i]\n            class_matrix_mean[i, j] = np.mean(aux_vec)\n            class_matrix_std[i, j] = np.std(aux_vec)\n\n    sorted_mean = np.sort(class_matrix_mean, axis=1)\n    selectivity = (sorted_mean[:, -1] - np.mean(sorted_mean[:, :-1], axis=1)) / \\\n                  (sorted_mean[:, -1] + np.mean(sorted_mean[:, :-1], axis=1))\n\n    print('Most selective 10 features:', np.argsort(selectivity)[-10:])\n    print('Highest 10 selectivities:', np.sort(selectivity)[-10:])\n    print('Selectivity shape:', selectivity.shape)\n\n    np.save('selectivity_' + str(args.layer) + '.npy', selectivity)"
  },
  {
    "path": "temporal_classification.py",
    "content": "import argparse\nimport os\nimport random\nimport shutil\nimport time\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.optim\nimport torch.multiprocessing as mp\nimport torch.utils.data\nimport torch.utils.data.distributed\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nfrom utils import GaussianBlur\n\n\nparser = argparse.ArgumentParser(description='Temporal classification with headcam data')\nparser.add_argument('data', metavar='DIR', help='path to dataset')\nparser.add_argument('--model', default='resnet50', choices=['resnet50', 'resnext101_32x8d', 'resnext50_32x4d',\n                                                            'mobilenet_v2'], help='model')\nparser.add_argument('-j', '--workers', default=32, type=int, metavar='N', help='number of data loading workers (default'\n                                                                               ':16)')\nparser.add_argument('--epochs', default=16, type=int, metavar='N', help='number of total epochs to run')\nparser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')\nparser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N',\n                    help='mini-batch size (default: 128), this is the total batch size of all GPUs on the current node '\n                         'when using Data Parallel or Distributed Data Parallel')\nparser.add_argument('--lr', '--learning-rate', default=0.0005, type=float, metavar='LR', help='initial learning rate',\n                    dest='lr')\nparser.add_argument('--wd', '--weight-decay', default=0.0, type=float, metavar='W', help='weight decay (default: 0)',\n                    dest='weight_decay')\nparser.add_argument('-p', '--print-freq', default=10000, type=int, metavar='N', help='print frequency (default: 250)')\nparser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')\nparser.add_argument('--world-size', default=-1, type=int, help='number of nodes for distributed training')\nparser.add_argument('--rank', default=-1, type=int, help='node rank for distributed training')\nparser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str, help='url used to set up distributed '\n                                                                                     'training')\nparser.add_argument('--dist-backend', default='nccl', type=str, help='distributed backend')\nparser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')\nparser.add_argument('--multiprocessing-distributed', action='store_true',\n                    help='Use multi-processing distributed training to launch '\n                         'N processes per node, which has N GPUs. This is the '\n                         'fastest way to use PyTorch for either single node or '\n                         'multi node data parallel training')\nparser.add_argument('--n_out', default=1000, type=int, help='output dim')\nparser.add_argument('--augmentation', default=True, action='store_false', help='whether to use data augmentation?')\n\n\ndef main():\n    args = parser.parse_args()\n\n    print(args)\n\n    if args.gpu is not None:\n        warnings.warn('You have chosen a specific GPU. This will completely disable data parallelism.')\n\n    if args.dist_url == \"env://\" and args.world_size == -1:\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    args.distributed = args.world_size > 1 or args.multiprocessing_distributed\n\n    ngpus_per_node = torch.cuda.device_count()\n    if args.multiprocessing_distributed:\n        # Since we have ngpus_per_node processes per node, the total world_size needs to be adjusted accordingly\n        args.world_size = ngpus_per_node * args.world_size\n        # Use torch.multiprocessing.spawn to launch distributed processes: the main_worker process function\n        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))\n    else:\n        # Simply call main_worker function\n        main_worker(args.gpu, ngpus_per_node, args)\n\n\ndef main_worker(gpu, ngpus_per_node, args):\n    args.gpu = gpu\n\n    if args.gpu is not None:\n        print(\"Use GPU: {} for training\".format(args.gpu))\n\n    print('Model:', args.model)\n    model = models.__dict__[args.model](pretrained=False)\n    if args.model.startswith('res'):\n        model.fc = torch.nn.Linear(in_features=2048, out_features=args.n_out, bias=True)\n    else:\n        model.classifier = torch.nn.Linear(in_features=1280, out_features=args.n_out, bias=True)\n\n    # DataParallel will divide and allocate batch_size to all available GPUs\n    model = torch.nn.DataParallel(model).cuda()\n\n    # define loss function (criterion) and optimizer\n    criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n    optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)\n    cudnn.benchmark = True\n\n    if args.resume:\n        if os.path.isfile(args.resume):\n            print(args.resume)\n            checkpoint = torch.load(args.resume)\n            model.load_state_dict(checkpoint['model_state_dict'])\n            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n        else:\n            print(\"=> no checkpoint found at '{}'\".format(args.resume))\n\n    savefile_name = args.model + '_augmentstrong_batch256_' + str(args.augmentation) + '_Y_5_288'\n\n    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n\n    if args.augmentation:\n        train_dataset = datasets.ImageFolder(\n            args.data,\n            transforms.Compose([\n                        transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n                        transforms.RandomApply([transforms.ColorJitter(0.9, 0.9, 0.9, 0.5)], p=0.9),\n                        transforms.RandomGrayscale(p=0.2),\n                        transforms.RandomApply([GaussianBlur([.1, 2.])], p=0.5),\n                        transforms.RandomHorizontalFlip(),\n                        transforms.ToTensor(),\n                        normalize\n                        ])\n        )\n    else:\n        train_dataset = datasets.ImageFolder(\n            args.data,\n            transforms.Compose([\n                transforms.ToTensor(),\n                normalize\n            ])\n        )\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, shuffle=True,\n        num_workers=args.workers, pin_memory=True, sampler=None\n    )\n\n    acc1_list = []\n\n    for epoch in range(args.start_epoch, args.epochs):\n\n        # train for one epoch\n        acc1 = train(train_loader, model, criterion, optimizer, epoch, args)\n        acc1_list.append(acc1)\n\n        torch.save({'acc1_list': acc1_list,\n                    'model_state_dict': model.state_dict(),\n                    'optimizer_state_dict': optimizer.state_dict()}, savefile_name + '_epoch_' + str(epoch) + '.tar')\n\n\ndef train(train_loader, model, criterion, optimizer, epoch, args):\n    batch_time = AverageMeter('Time', ':6.3f')\n    data_time = AverageMeter('Data', ':6.3f')\n    losses = AverageMeter('Loss', ':.4e')\n    top1 = AverageMeter('Acc@1', ':6.2f')\n    top5 = AverageMeter('Acc@5', ':6.2f')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses, top1, top5],\n        prefix=\"Epoch: [{}]\".format(epoch))\n\n    # switch to train mode\n    model.train()\n\n    end = time.time()\n    for i, (images, target) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n\n        if args.gpu is not None:\n            images = images.cuda(args.gpu, non_blocking=True)\n        target = target.cuda(args.gpu, non_blocking=True)\n\n        # compute output\n        output = model(images)\n        loss = criterion(output, target)\n\n        # measure accuracy and record loss\n        acc1, acc5 = accuracy(output, target, topk=(1, 5))\n        losses.update(loss.item(), images.size(0))\n        top1.update(acc1[0], images.size(0))\n        top5.update(acc5[0], images.size(0))\n\n        # compute gradient and do SGD step\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % args.print_freq == 0:\n            progress.display(i)\n\n    return top1.avg.cpu().numpy()\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'\n\n\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n    with torch.no_grad():\n        maxk = max(topk)\n        batch_size = target.size(0)\n\n        _, pred = output.topk(maxk, 1, True, True)\n        pred = pred.t()\n        correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n        res = []\n        for k in topk:\n            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n            res.append(correct_k.mul_(100.0 / batch_size))\n        return res\n\n\nif __name__ == '__main__':\n    main()"
  },
  {
    "path": "utils.py",
    "content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nfrom PIL import ImageFilter\nimport random\n\n\nclass GaussianBlur(object):\n    \"\"\"Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709\"\"\"\n\n    def __init__(self, sigma=[.1, 2.]):\n        self.sigma = sigma\n\n    def __call__(self, x):\n        sigma = random.uniform(self.sigma[0], self.sigma[1])\n        x = x.filter(ImageFilter.GaussianBlur(radius=sigma))\n        return x"
  }
]