Bring activations mostly in line with timm, update some docstrings, fixup ONNX export and caffe2 scripts.

Author: rwightman

caffe2_benchmark.py

@@ -1,16 +1,22 @@
 """ Caffe2 validation script
-This script runs Caffe2 benchmark on exported model.
+
+This script runs Caffe2 benchmark on exported ONNX model.
+It is a useful tool for reporting model FLOPS.
+
+Copyright 2020 Ross Wightman
 """
 import argparse
 from caffe2.python import core, workspace, model_helper
 from caffe2.proto import caffe2_pb2
 
 
 parser = argparse.ArgumentParser(description='Caffe2 Model Benchmark')
+parser.add_argument('--c2-prefix', default='', type=str, metavar='NAME',
+                    help='caffe2 model pb name prefix')
 parser.add_argument('--c2-init', default='', type=str, metavar='PATH',
-                    help='path to latest checkpoint (default: none)')
+                    help='caffe2 model init .pb')
 parser.add_argument('--c2-predict', default='', type=str, metavar='PATH',
-                    help='path to latest checkpoint (default: none)')
+                    help='caffe2 model predict .pb')
 parser.add_argument('-b', '--batch-size', default=1, type=int,
                     metavar='N', help='mini-batch size (default: 1)')
 parser.add_argument('--img-size', default=224, type=int,
@@ -20,20 +26,23 @@
 def main():
     args = parser.parse_args()
     args.gpu_id = 0
+    if args.c2_prefix:
+        args.c2_init = args.c2_prefix + '.init.pb'
+        args.c2_predict = args.c2_prefix + '.predict.pb'
 
     model = model_helper.ModelHelper(name="le_net", init_params=False)
 
     # Bring in the init net from init_net.pb
     init_net_proto = caffe2_pb2.NetDef()
     with open(args.c2_init, "rb") as f:
         init_net_proto.ParseFromString(f.read())
-    model.param_init_net = core.Net(init_net_proto)  # model.param_init_net.AppendNet(core.Net(init_net_proto)) #
+    model.param_init_net = core.Net(init_net_proto)
 
     # bring in the predict net from predict_net.pb
     predict_net_proto = caffe2_pb2.NetDef()
     with open(args.c2_predict, "rb") as f:
         predict_net_proto.ParseFromString(f.read())
-    model.net = core.Net(predict_net_proto)  # model.net.AppendNet(core.Net(predict_net_proto))
+    model.net = core.Net(predict_net_proto)
 
     # CUDA performance not impressive
     #device_opts = core.DeviceOption(caffe2_pb2.PROTO_CUDA, args.gpu_id)

caffe2_validate.py

@@ -1,7 +1,10 @@
 """ Caffe2 validation script
-This script is intended to verify exported models running in Caffe2
-It utilizes the same PyTorch dataloader/processing pipeline for comparison against
-the originals, I also have no desire to write that code in Caffe2.
+
+This script is created to verify exported ONNX models running in Caffe2
+It utilizes the same PyTorch dataloader/processing pipeline for a
+fair comparison against the originals.
+
+Copyright 2020 Ross Wightman
 """
 import argparse
 import numpy as np
@@ -14,12 +17,12 @@
 parser = argparse.ArgumentParser(description='Caffe2 ImageNet Validation')
 parser.add_argument('data', metavar='DIR',
                     help='path to dataset')
-parser.add_argument('--model', '-m', metavar='MODEL', default='spnasnet1_00',
-                    help='model architecture (default: dpn92)')
+parser.add_argument('--c2-prefix', default='', type=str, metavar='NAME',
+                    help='caffe2 model pb name prefix')
 parser.add_argument('--c2-init', default='', type=str, metavar='PATH',
-                    help='path to latest checkpoint (default: none)')
+                    help='caffe2 model init .pb')
 parser.add_argument('--c2-predict', default='', type=str, metavar='PATH',
-                    help='path to latest checkpoint (default: none)')
+                    help='caffe2 model predict .pb')
 parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
                     help='number of data loading workers (default: 2)')
 parser.add_argument('-b', '--batch-size', default=256, type=int,
@@ -43,22 +46,25 @@
 def main():
     args = parser.parse_args()
     args.gpu_id = 0
+    if args.c2_prefix:
+        args.c2_init = args.c2_prefix + '.init.pb'
+        args.c2_predict = args.c2_prefix + '.predict.pb'
 
     model = model_helper.ModelHelper(name="validation_net", init_params=False)
 
     # Bring in the init net from init_net.pb
     init_net_proto = caffe2_pb2.NetDef()
     with open(args.c2_init, "rb") as f:
         init_net_proto.ParseFromString(f.read())
-    model.param_init_net = core.Net(init_net_proto)  # model.param_init_net.AppendNet(core.Net(init_net_proto)) #
+    model.param_init_net = core.Net(init_net_proto)
 
     # bring in the predict net from predict_net.pb
     predict_net_proto = caffe2_pb2.NetDef()
     with open(args.c2_predict, "rb") as f:
         predict_net_proto.ParseFromString(f.read())
-    model.net = core.Net(predict_net_proto)  # model.net.AppendNet(core.Net(predict_net_proto))
+    model.net = core.Net(predict_net_proto)
 
-    data_config = resolve_data_config(args.model, args)
+    data_config = resolve_data_config(None, args)
     loader = create_loader(
         Dataset(args.data, load_bytes=args.tf_preprocessing),
         input_size=data_config['input_size'],

data/tf_preprocessing.py

@@ -1,3 +1,11 @@
+""" Tensorflow Preprocessing Adapter
+
+Allows use of Tensorflow preprocessing pipeline in PyTorch Transform
+
+Copyright of original Tensorflow code below.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
 # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,7 +20,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""ImageNet preprocessing for MnasNet."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function

data/transforms.py

@@ -17,7 +17,7 @@
 def resolve_data_config(model, args, default_cfg={}, verbose=True):
     new_config = {}
     default_cfg = default_cfg
-    if not default_cfg and hasattr(model, 'default_cfg'):
+    if not default_cfg and model is not None and hasattr(model, 'default_cfg'):
         default_cfg = model.default_cfg
 
     # Resolve input/image size
@@ -40,7 +40,7 @@ def resolve_data_config(model, args, default_cfg={}, verbose=True):
         new_config['interpolation'] = default_cfg['interpolation']
 
     # resolve dataset + model mean for normalization
-    new_config['mean'] = get_mean_by_model(args.model)
+    new_config['mean'] = IMAGENET_DEFAULT_MEAN
     if args.mean is not None:
         mean = tuple(args.mean)
         if len(mean) == 1:
@@ -52,7 +52,7 @@ def resolve_data_config(model, args, default_cfg={}, verbose=True):
         new_config['mean'] = default_cfg['mean']
 
     # resolve dataset + model std deviation for normalization
-    new_config['std'] = get_std_by_model(args.model)
+    new_config['std'] = IMAGENET_DEFAULT_STD
     if args.std is not None:
         std = tuple(args.std)
         if len(std) == 1:
@@ -78,26 +78,6 @@ def resolve_data_config(model, args, default_cfg={}, verbose=True):
     return new_config
 
 
-def get_mean_by_model(model_name):
-    model_name = model_name.lower()
-    if 'dpn' in model_name:
-        return IMAGENET_DPN_STD
-    elif 'ception' in model_name:
-        return IMAGENET_INCEPTION_MEAN
-    else:
-        return IMAGENET_DEFAULT_MEAN
-
-
-def get_std_by_model(model_name):
-    model_name = model_name.lower()
-    if 'dpn' in model_name:
-        return IMAGENET_DEFAULT_STD
-    elif 'ception' in model_name:
-        return IMAGENET_INCEPTION_STD
-    else:
-        return IMAGENET_DEFAULT_STD
-
-
 class ToNumpy:
 
     def __call__(self, pil_img):

geffnet/activations/__init__.py

@@ -1,5 +1,5 @@
 from geffnet import config
-from geffnet.activations.activations_autofn import *
+from geffnet.activations.activations_me import *
 from geffnet.activations.activations_jit import *
 from geffnet.activations.activations import *
 
@@ -15,16 +15,16 @@
     hard_swish=hard_swish,
 )
 
-_ACT_FN_AUTO = dict(
-    swish=swish_auto,
-    mish=mish_auto,
-)
-
 _ACT_FN_JIT = dict(
     swish=swish_jit,
     mish=mish_jit,
-    #hard_swish=hard_swish_jit,
-    #hard_sigmoid_jit=hard_sigmoid_jit,
+)
+
+_ACT_FN_ME = dict(
+    swish=swish_me,
+    mish=mish_me,
+    hard_swish=hard_swish_me,
+    hard_sigmoid_jit=hard_sigmoid_me,
 )
 
 _ACT_LAYER_DEFAULT = dict(
@@ -38,16 +38,16 @@
     hard_swish=HardSwish,
 )
 
-_ACT_LAYER_AUTO = dict(
-    swish=SwishAuto,
-    mish=MishAuto,
-)
-
 _ACT_LAYER_JIT = dict(
     swish=SwishJit,
     mish=MishJit,
-    #hard_swish=HardSwishJit,
-    #hard_sigmoid=HardSigmoidJit
+)
+
+_ACT_LAYER_ME = dict(
+    swish=SwishMe,
+    mish=MishMe,
+    hard_swish=HardSwishMe,
+    hard_sigmoid=HardSigmoidMe
 )
 
 _OVERRIDE_FN = dict()
@@ -92,14 +92,15 @@ def get_act_fn(name='relu'):
     """
     if name in _OVERRIDE_FN:
         return _OVERRIDE_FN[name]
-    if not config.is_exportable() and not config.is_scriptable():
-        # If not exporting or scripting the model, first look for a JIT optimized version
-        # of our activation, then a custom autograd.Function variant before defaulting to
-        # a Python or Torch builtin impl
-        if name in _ACT_FN_JIT:
-            return _ACT_FN_JIT[name]
-        if name in _ACT_FN_AUTO:
-            return _ACT_FN_AUTO[name]
+    no_me = config.is_exportable() or config.is_scriptable() or config.is_no_jit()
+    if not no_me and name in _ACT_FN_ME:
+        # If not exporting or scripting the model, first look for a memory optimized version
+        # activation with custom autograd, then fallback to jit scripted, then a Python or Torch builtin
+        return _ACT_FN_ME[name]
+    no_jit = config.is_exportable() or config.is_no_jit()
+    # NOTE: export tracing should work with jit scripted components, but I keep running into issues
+    if no_jit and name in _ACT_FN_JIT:  # jit scripted models should be okay for export/scripting
+        return _ACT_FN_JIT[name]
     return _ACT_FN_DEFAULT[name]
 
 
@@ -110,11 +111,12 @@ def get_act_layer(name='relu'):
     """
     if name in _OVERRIDE_LAYER:
         return _OVERRIDE_LAYER[name]
-    if not config.is_exportable() and not config.is_scriptable():
-        if name in _ACT_LAYER_JIT:
-            return _ACT_LAYER_JIT[name]
-        if name in _ACT_LAYER_AUTO:
-            return _ACT_LAYER_AUTO[name]
+    no_me = config.is_exportable() or config.is_scriptable() or config.is_no_jit()
+    if not no_me and name in _ACT_LAYER_ME:
+        return _ACT_LAYER_ME[name]
+    no_jit = config.is_exportable() or config.is_no_jit()
+    if not no_jit and name in _ACT_LAYER_JIT:  # jit scripted models should be okay for export/scripting
+        return _ACT_LAYER_JIT[name]
     return _ACT_LAYER_DEFAULT[name]
 
 

geffnet/activations/activations.py

@@ -1,9 +1,19 @@
+""" Activations
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+Copyright 2020 Ross Wightman
+"""
 from torch import nn as nn
 from torch.nn import functional as F
 
 
 def swish(x, inplace: bool = False):
-    """Swish - Described in: https://arxiv.org/abs/1710.05941
+    """Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3)
+    and also as Swish (https://arxiv.org/abs/1710.05941).
+
+    TODO Rename to SiLU with addition to PyTorch
     """
     return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())