Remove all prints, change most to logging calls, tweak alignment of batch logs, improve setup.py

Author: rwightman

README.md

@@ -30,7 +30,7 @@ I've included a few of my favourite models, but this is not an exhaustive collec
 * DPN (from [me](https://github.com/rwightman/pytorch-dpn-pretrained), weights hosted by Cadene)
     * DPN-68, DPN-68b, DPN-92, DPN-98, DPN-131, DPN-107
 * Generic EfficientNet (from my standalone [GenMobileNet](https://github.com/rwightman/genmobilenet-pytorch)) - A generic model that implements many of the mobile optimized architecture search derived models that utilize similar DepthwiseSeparable and InvertedResidual blocks
-    * EfficientNet (B0-B4) (https://arxiv.org/abs/1905.11946) -- validated, compat with TF weights
+    * EfficientNet (B0-B5) (https://arxiv.org/abs/1905.11946) -- validated, compat with TF weights
     * MNASNet B1, A1 (Squeeze-Excite), and Small (https://arxiv.org/abs/1807.11626)
     * MobileNet-V1 (https://arxiv.org/abs/1704.04861)
     * MobileNet-V2 (https://arxiv.org/abs/1801.04381)
@@ -187,9 +187,6 @@ To run inference from a checkpoint:
 
 ## TODO
 A number of additions planned in the future for various projects, incl
-* Find optimal training hyperparams and create/port pretraiend weights for the generic MobileNet variants
 * Do a model performance (speed + accuracy) benchmarking across all models (make runable as script)
-* More training experiments
-* Make folder/file layout compat with usage as a module
 * Add usage examples to comments, good hyper params for training
 * Comments, cleanup and the usual things that get pushed back

inference.py

@@ -8,12 +8,13 @@
 import os
 import time
 import argparse
+import logging
 import numpy as np
 import torch
 
 from timm.models import create_model, apply_test_time_pool
 from timm.data import Dataset, create_loader, resolve_data_config
-from timm.utils import AverageMeter
+from timm.utils import AverageMeter, setup_default_logging
 
 torch.backends.cudnn.benchmark = True
 
@@ -38,8 +39,8 @@
                     help='Image resize interpolation type (overrides model)')
 parser.add_argument('--num-classes', type=int, default=1000,
                     help='Number classes in dataset')
-parser.add_argument('--print-freq', '-p', default=10, type=int,
-                    metavar='N', help='print frequency (default: 10)')
+parser.add_argument('--log-freq', default=10, type=int,
+                    metavar='N', help='batch logging frequency (default: 10)')
 parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
                     help='path to latest checkpoint (default: none)')
 parser.add_argument('--pretrained', dest='pretrained', action='store_true',
@@ -53,8 +54,8 @@
 
 
 def main():
+    setup_default_logging()
     args = parser.parse_args()
-
     # might as well try to do something useful...
     args.pretrained = args.pretrained or not args.checkpoint
 
@@ -66,8 +67,8 @@ def main():
         pretrained=args.pretrained,
         checkpoint_path=args.checkpoint)
 
-    print('Model %s created, param count: %d' %
-          (args.model, sum([m.numel() for m in model.parameters()])))
+    logging.info('Model %s created, param count: %d' %
+                 (args.model, sum([m.numel() for m in model.parameters()])))
 
     config = resolve_data_config(model, args)
     model, test_time_pool = apply_test_time_pool(model, config, args)
@@ -105,9 +106,8 @@ def main():
             batch_time.update(time.time() - end)
             end = time.time()
 
-            if batch_idx % args.print_freq == 0:
-                print('Predict: [{0}/{1}]\t'
-                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
+            if batch_idx % args.log_freq == 0:
+                logging.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
                     batch_idx, len(loader), batch_time=batch_time))
 
     topk_ids = np.concatenate(topk_ids, axis=0).squeeze()

setup.py

@@ -19,21 +19,27 @@
     url='https://github.com/rwightman/pytorch-image-models',
     author='Ross Wightman',
     author_email='hello@rwightman.com',
-    classifiers=[  # Optional
+    classifiers=[
         # How mature is this project? Common values are
         #   3 - Alpha
         #   4 - Beta
         #   5 - Production/Stable
         'Development Status :: 3 - Alpha',
-        'Intended Audience :: Developers',
-        'Topic :: Software Development :: Build Tools',
-        'License :: OSI Approved :: Apache License',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
         'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
     ],
 
     # Note that this is a string of words separated by whitespace, not a list.
     keywords='pytorch pretrained models efficientnet mobilenetv3 mnasnet',
     packages=find_packages(exclude=['convert']),
-    install_requires=['torch', 'torchvision'],
+    install_requires=['torch >= 1.0', 'torchvision'],
     python_requires='>=3.6',
 )

timm/data/config.py

@@ -1,3 +1,4 @@
+import logging
 from .constants import *
 
 
@@ -56,9 +57,9 @@ def resolve_data_config(model, args, default_cfg={}, verbose=True):
         new_config['crop_pct'] = default_cfg['crop_pct']
 
     if verbose:
-        print('Data processing configuration for current model + dataset:')
+        logging.info('Data processing configuration for current model + dataset:')
         for n, v in new_config.items():
-            print('\t%s: %s' % (n, str(v)))
+            logging.info('\t%s: %s' % (n, str(v)))
 
     return new_config
 

timm/models/adaptive_avgmax_pool.py

@@ -82,7 +82,7 @@ def __init__(self, output_size=1, pool_type='avg'):
             self.pool = nn.AdaptiveMaxPool2d(output_size)
         else:
             if pool_type != 'avg':
-                print('Invalid pool type %s specified. Defaulting to average pooling.' % pool_type)
+                assert False, 'Invalid pool type: %s' % pool_type
             self.pool = nn.AdaptiveAvgPool2d(output_size)
 
     def forward(self, x):

timm/models/densenet.py

@@ -86,7 +86,6 @@ def densenet161(num_classes=1000, in_chans=3, pretrained=False, **kwargs):
     r"""Densenet-201 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
-    print(num_classes, in_chans, pretrained)
     default_cfg = default_cfgs['densenet161']
     model = DenseNet(num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24),
                      num_classes=num_classes, in_chans=in_chans, **kwargs)

timm/models/gen_efficientnet.py

@@ -17,6 +17,7 @@
 
 import math
 import re
+import logging
 from copy import deepcopy
 
 import torch
@@ -336,7 +337,7 @@ def _make_block(self, ba):
         ba['act_fn'] = ba['act_fn'] if ba['act_fn'] is not None else self.act_fn
         assert ba['act_fn'] is not None
         if self.verbose:
-            print('args:', ba)
+            logging.info('  Args: {}'.format(str(ba)))
         #  could replace this if with lambdas or functools binding if variety increases
         if bt == 'ir':
             ba['drop_connect_rate'] = self.drop_connect_rate
@@ -358,7 +359,7 @@ def _make_stack(self, stack_args):
         # each stack (stage) contains a list of block arguments
         for block_idx, ba in enumerate(stack_args):
             if self.verbose:
-                print('block', block_idx, end=', ')
+                logging.info(' Block: {}'.format(block_idx))
             if block_idx >= 1:
                 # only the first block in any stack/stage can have a stride > 1
                 ba['stride'] = 1
@@ -370,24 +371,22 @@ def __call__(self, in_chs, block_args):
         """ Build the blocks
         Args:
             in_chs: Number of input-channels passed to first block
-            arch_def: A list of lists, outer list defines stacks (or stages), inner
+            block_args: A list of lists, outer list defines stages, inner
                 list contains strings defining block configuration(s)
         Return:
              List of block stacks (each stack wrapped in nn.Sequential)
         """
         if self.verbose:
-            print('Building model trunk with %d stacks (stages)...' % len(block_args))
+            logging.info('Building model trunk with %d stages...' % len(block_args))
         self.in_chs = in_chs
         blocks = []
         # outer list of block_args defines the stacks ('stages' by some conventions)
         for stack_idx, stack in enumerate(block_args):
             if self.verbose:
-                print('stack', stack_idx)
+                logging.info('Stack: {}'.format(stack_idx))
             assert isinstance(stack, list)
             stack = self._make_stack(stack)
             blocks.append(stack)
-            if self.verbose:
-                print()
         return blocks
 
 

timm/models/helpers.py

@@ -1,6 +1,7 @@
 import torch
 import torch.utils.model_zoo as model_zoo
 import os
+import logging
 from collections import OrderedDict
 
 
@@ -21,9 +22,9 @@ def load_checkpoint(model, checkpoint_path, use_ema=False):
             model.load_state_dict(new_state_dict)
         else:
             model.load_state_dict(checkpoint)
-        print("=> Loaded {} from checkpoint '{}'".format(state_dict_key or 'weights', checkpoint_path))
+        logging.info("Loaded {} from checkpoint '{}'".format(state_dict_key or 'weights', checkpoint_path))
     else:
-        print("=> Error: No checkpoint found at '{}'".format(checkpoint_path))
+        logging.error("No checkpoint found at '{}'".format(checkpoint_path))
         raise FileNotFoundError()
 
 
@@ -40,27 +41,27 @@ def resume_checkpoint(model, checkpoint_path, start_epoch=None):
             if 'optimizer' in checkpoint:
                 optimizer_state = checkpoint['optimizer']
             start_epoch = checkpoint['epoch'] if start_epoch is None else start_epoch
-            print("=> Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, checkpoint['epoch']))
+            logging.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, checkpoint['epoch']))
         else:
             model.load_state_dict(checkpoint)
             start_epoch = 0 if start_epoch is None else start_epoch
-            print("=> Loaded checkpoint '{}'".format(checkpoint_path))
+            logging.info("Loaded checkpoint '{}'".format(checkpoint_path))
         return optimizer_state, start_epoch
     else:
-        print("=> Error: No checkpoint found at '{}'".format(checkpoint_path))
+        logging.error("No checkpoint found at '{}'".format(checkpoint_path))
         raise FileNotFoundError()
 
 
 def load_pretrained(model, default_cfg, num_classes=1000, in_chans=3, filter_fn=None):
     if 'url' not in default_cfg or not default_cfg['url']:
-        print("Warning: pretrained model URL is invalid, using random initialization.")
+        logging.warning("Pretrained model URL is invalid, using random initialization.")
         return
 
     state_dict = model_zoo.load_url(default_cfg['url'])
 
     if in_chans == 1:
         conv1_name = default_cfg['first_conv']
-        print('Converting first conv (%s) from 3 to 1 channel' % conv1_name)
+        logging.info('Converting first conv (%s) from 3 to 1 channel' % conv1_name)
         conv1_weight = state_dict[conv1_name + '.weight']
         state_dict[conv1_name + '.weight'] = conv1_weight.sum(dim=1, keepdim=True)
     elif in_chans != 3:

timm/models/test_time_pool.py

@@ -1,3 +1,4 @@
+import logging
 from torch import nn
 import torch.nn.functional as F
 from .adaptive_avgmax_pool import adaptive_avgmax_pool2d
@@ -31,8 +32,8 @@ def apply_test_time_pool(model, config, args):
     if not args.no_test_pool and \
             config['input_size'][-1] > model.default_cfg['input_size'][-1] and \
             config['input_size'][-2] > model.default_cfg['input_size'][-2]:
-        print('Target input size %s > pretrained default %s, using test time pooling' %
-              (str(config['input_size'][-2:]), str(model.default_cfg['input_size'][-2:])))
+        logging.info('Target input size %s > pretrained default %s, using test time pooling' %
+                     (str(config['input_size'][-2:]), str(model.default_cfg['input_size'][-2:])))
         model = TestTimePoolHead(model, original_pool=model.default_cfg['pool_size'])
         test_time_pool = True
     return model, test_time_pool

timm/scheduler/tanh_lr.py

@@ -50,7 +50,6 @@ def __init__(self,
         self.t_in_epochs = t_in_epochs
         if self.warmup_t:
             t_v = self.base_values if self.warmup_prefix else self._get_lr(self.warmup_t)
-            print(t_v)
             self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in t_v]
             super().update_groups(self.warmup_lr_init)
         else: