Merge pull request #1 from rwightman/master

merge master

Author: mrT23

.gitignore

@@ -104,3 +104,5 @@ venv.bak/
 *.tar
 *.pth
 *.gz
+Untitled.ipynb
+Testing notebook.ipynb

README.md

@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+""" Checkpoint Averaging Script
+
+This script averages all model weights for checkpoints in specified path that match
+the specified filter wildcard. All checkpoints must be from the exact same model.
+
+For any hope of decent results, the checkpoints should be from the same or child
+(via resumes) training session. This can be viewed as similar to maintaining running
+EMA (exponential moving average) of the model weights or performing SWA (stochastic
+weight averaging), but post-training.
+
+Hacked together by Ross Wightman (https://github.com/rwightman)
+"""
+import torch
+import argparse
+import os
+import glob
+import hashlib
+from timm.models.helpers import load_state_dict
+
+parser = argparse.ArgumentParser(description='PyTorch Checkpoint Averager')
+parser.add_argument('--input', default='', type=str, metavar='PATH',
+                    help='path to base input folder containing checkpoints')
+parser.add_argument('--filter', default='*.pth.tar', type=str, metavar='WILDCARD',
+                    help='checkpoint filter (path wildcard)')
+parser.add_argument('--output', default='./averaged.pth', type=str, metavar='PATH',
+                    help='output filename')
+parser.add_argument('--no-use-ema', dest='no_use_ema', action='store_true',
+                    help='Force not using ema version of weights (if present)')
+parser.add_argument('--no-sort', dest='no_sort', action='store_true',
+                    help='Do not sort and select by checkpoint metric, also makes "n" argument irrelevant')
+parser.add_argument('-n', type=int, default=10, metavar='N',
+                    help='Number of checkpoints to average')
+
+
+def checkpoint_metric(checkpoint_path):
+    if not checkpoint_path or not os.path.isfile(checkpoint_path):
+        return {}
+    print("=> Extracting metric from checkpoint '{}'".format(checkpoint_path))
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    metric = None
+    if 'metric' in checkpoint:
+        metric = checkpoint['metric']
+    return metric
+
+
+def main():
+    args = parser.parse_args()
+    # by default use the EMA weights (if present)
+    args.use_ema = not args.no_use_ema
+    # by default sort by checkpoint metric (if present) and avg top n checkpoints
+    args.sort = not args.no_sort
+
+    if os.path.exists(args.output):
+        print("Error: Output filename ({}) already exists.".format(args.output))
+        exit(1)
+
+    pattern = args.input
+    if not args.input.endswith(os.path.sep) and not args.filter.startswith(os.path.sep):
+        pattern += os.path.sep
+    pattern += args.filter
+    checkpoints = glob.glob(pattern, recursive=True)
+
+    if args.sort:
+        checkpoint_metrics = []
+        for c in checkpoints:
+            metric = checkpoint_metric(c)
+            if metric is not None:
+                checkpoint_metrics.append((metric, c))
+        checkpoint_metrics = list(sorted(checkpoint_metrics))
+        checkpoint_metrics = checkpoint_metrics[-args.n:]
+        print("Selected checkpoints:")
+        [print(m, c) for m, c in checkpoint_metrics]
+        avg_checkpoints = [c for m, c in checkpoint_metrics]
+    else:
+        avg_checkpoints = checkpoints
+        print("Selected checkpoints:")
+        [print(c) for c in checkpoints]
+
+    avg_state_dict = {}
+    avg_counts = {}
+    for c in avg_checkpoints:
+        new_state_dict = load_state_dict(c, args.use_ema)
+        if not new_state_dict:
+            print("Error: Checkpoint ({}) doesn't exist".format(args.checkpoint))
+            continue
+
+        for k, v in new_state_dict.items():
+            if k not in avg_state_dict:
+                avg_state_dict[k] = v.clone().to(dtype=torch.float64)
+                avg_counts[k] = 1
+            else:
+                avg_state_dict[k] += v.to(dtype=torch.float64)
+                avg_counts[k] += 1
+
+    for k, v in avg_state_dict.items():
+        v.div_(avg_counts[k])
+
+    # float32 overflow seems unlikely based on weights seen to date, but who knows
+    float32_info = torch.finfo(torch.float32)
+    final_state_dict = {}
+    for k, v in avg_state_dict.items():
+        v = v.clamp(float32_info.min, float32_info.max)
+        final_state_dict[k] = v.to(dtype=torch.float32)
+
+    torch.save(final_state_dict, args.output)
+    with open(args.output, 'rb') as f:
+        sha_hash = hashlib.sha256(f.read()).hexdigest()
+    print("=> Saved state_dict to '{}, SHA256: {}'".format(args.output, sha_hash))
+
+
+if __name__ == '__main__':
+    main()

avg_checkpoints.py

@@ -1,16 +1,30 @@
+#!/usr/bin/env python
+""" Checkpoint Cleaning Script
+
+Takes training checkpoints with GPU tensors, optimizer state, extra dict keys, etc.
+and outputs a CPU  tensor checkpoint with only the `state_dict` along with SHA256
+calculation for model zoo compatibility.
+
+Hacked together by Ross Wightman (https://github.com/rwightman)
+"""
 import torch
 import argparse
 import os
 import hashlib
+import shutil
 from collections import OrderedDict
 
-parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
+parser = argparse.ArgumentParser(description='PyTorch Checkpoint Cleaner')
 parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
                     help='path to latest checkpoint (default: none)')
-parser.add_argument('--output', default='./cleaned.pth', type=str, metavar='PATH',
+parser.add_argument('--output', default='', type=str, metavar='PATH',
                     help='output path')
 parser.add_argument('--use-ema', dest='use_ema', action='store_true',
                     help='use ema version of weights if present')
+parser.add_argument('--clean-aux-bn', dest='clean_aux_bn', action='store_true',
+                    help='remove auxiliary batch norm layers (from SplitBN training) from checkpoint')
+
+_TEMP_NAME = './_checkpoint.pth'
 
 
 def main():
@@ -31,19 +45,31 @@ def main():
             if state_dict_key in checkpoint:
                 state_dict = checkpoint[state_dict_key]
             else:
-                print("Error: No state_dict found in checkpoint {}.".format(args.checkpoint))
-                exit(1)
+                state_dict = checkpoint
         else:
-            state_dict = checkpoint
+            assert False
         for k, v in state_dict.items():
+            if args.clean_aux_bn and 'aux_bn' in k:
+                # If all aux_bn keys are removed, the SplitBN layers will end up as normal and
+                # load with the unmodified model using BatchNorm2d.
+                continue
             name = k[7:] if k.startswith('module') else k
             new_state_dict[name] = v
         print("=> Loaded state_dict from '{}'".format(args.checkpoint))
 
-        torch.save(new_state_dict, args.output)
-        with open(args.output, 'rb') as f:
+        torch.save(new_state_dict, _TEMP_NAME)
+        with open(_TEMP_NAME, 'rb') as f:
             sha_hash = hashlib.sha256(f.read()).hexdigest()
-        print("=> Saved state_dict to '{}, SHA256: {}'".format(args.output, sha_hash))
+
+        if args.output:
+            checkpoint_root, checkpoint_base = os.path.split(args.output)
+            checkpoint_base = os.path.splitext(checkpoint_base)[0]
+        else:
+            checkpoint_root = ''
+            checkpoint_base = os.path.splitext(args.checkpoint)[0]
+        final_filename = '-'.join([checkpoint_base, sha_hash[:8]]) + '.pth'
+        shutil.move(_TEMP_NAME, os.path.join(checkpoint_root, final_filename))
+        print("=> Saved state_dict to '{}, SHA256: {}'".format(final_filename, sha_hash))
     else:
         print("Error: Checkpoint ({}) doesn't exist".format(args.checkpoint))
 

clean_checkpoint.py

@@ -5,7 +5,7 @@
 import mxnet as mx
 import gluoncv
 import torch
-from models.model_factory import create_model
+from timm import create_model
 
 parser = argparse.ArgumentParser(description='Convert from MXNet')
 parser.add_argument('--model', default='all', type=str, metavar='MODEL',

convert/convert_from_mxnet.py

@@ -0,0 +1,4 @@
+dependencies = ['torch']
+from timm.models import registry
+
+globals().update(registry._model_entrypoints)

hubconf.py

@@ -1,10 +1,10 @@
-"""Sample PyTorch Inference script
-"""
+#!/usr/bin/env python
+"""PyTorch Inference Script
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+An example inference script that outputs top-k class ids for images in a folder into a csv.
 
+Hacked together by Ross Wightman (https://github.com/rwightman)
+"""
 import os
 import time
 import argparse
@@ -29,7 +29,7 @@
                     help='number of data loading workers (default: 2)')
 parser.add_argument('-b', '--batch-size', default=256, type=int,
                     metavar='N', help='mini-batch size (default: 256)')
-parser.add_argument('--img-size', default=224, type=int,
+parser.add_argument('--img-size', default=None, type=int,
                     metavar='N', help='Input image dimension')
 parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
                     help='Override mean pixel value of dataset')

inference.py

@@ -1,2 +1,3 @@
-torch~=1.0
-torchvision
+torch>=1.2.0
+torchvision>=0.4.0
+pyyaml

requirements.txt

@@ -0,0 +1,39 @@
+# Validation Results
+
+This folder contains validation results for the models in this collection having pretrained weights. Since the focus for this repository is currently ImageNet-1k classification, all of the results are based on datasets compatible with ImageNet-1k classes.
+
+## Datasets
+
+There are currently results for the ImageNet validation set and 3 additional test sets.
+
+### ImageNet Validation - [`results-imagenet.csv`](results-imagenet.csv)
+
+* Source: http://image-net.org/challenges/LSVRC/2012/index
+* Paper: "ImageNet Large Scale Visual Recognition Challenge" - https://arxiv.org/abs/1409.0575
+
+The standard 50,000 image ImageNet-1k validation set. Model selection during training utilizes this validation set, so it is not a true test set. Question: Does anyone have the official ImageNet-1k test set classification labels now that challenges are done?
+
+### ImageNetV2 Matched Frequency - [`results-imagenetv2-matched-frequency.csv`](results-imagenetv2-matched-frequency.csv)
+
+* Source: https://github.com/modestyachts/ImageNetV2
+* Paper: "Do ImageNet Classifiers Generalize to ImageNet?" - https://arxiv.org/abs/1902.10811
+
+An ImageNet test set of 10,000 images sampled from new images roughly 10 years after the original. Care was taken to replicate the original ImageNet curation/sampling process.
+
+### ImageNet-Sketch - [`results-sketch.csv`](results-sketch.csv)
+
+* Source: https://github.com/HaohanWang/ImageNet-Sketch
+* Paper: "Learning Robust Global Representations by Penalizing Local Predictive Power" - https://arxiv.org/abs/1905.13549
+
+50,000 non photographic (or photos of such) images (sketches, doodles, mostly monochromatic) covering all 1000 ImageNet classes.
+
+### ImageNet-Adversarial - [`results-imagenet-a.csv`](results-imagenet-a.csv)
+
+* Source: https://github.com/hendrycks/natural-adv-examples
+* Paper: "Natural Adversarial Examples" - https://arxiv.org/abs/1907.07174
+
+A collection of 7500 images covering 200 of the 1000 ImageNet classes. Images are naturally occuring adversarial examples that confuse typical ImageNet classifiers. This is a challenging dataset, your typical ResNet-50 will score 0% top-1.
+
+## TODO
+* Add rank difference, and top-1/top-5 difference from ImageNet-1k validation for the 3 additional test sets
+* Explore adding a reduced version of ImageNet-C (Corruptions) and ImageNet-P (Perturbations) from https://github.com/hendrycks/robustness. The originals are huge and image size specific.