add shufflenetv2 model

ShuffleNet/.gitignore

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# visual studio
+.vscode/
+
+.idea/
+results/*
+# models

ShuffleNet/LICENSE

+MIT License
+
+Copyright (c) 2018 Evgenii Zheltonozhskii
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ShuffleNet/README.md

+# ShuffleNetv2 in PyTorch
+
+An implementation of `ShuffleNetv2` in PyTorch. `ShuffleNetv2` is an efficient convolutional neural network architecture for mobile devices. For more information check the paper:
+[ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design](https://arxiv.org/abs/1807.11164)
+
+## Usage
+
+Clone the repo:
+```bash
+git clone https://github.com/Randl/ShuffleNetV2-pytorch
+pip install -r requirements.txt
+```
+
+Use the model defined in `model.py` to run ImageNet example:
+```bash
+python imagenet.py --dataroot "/path/to/imagenet/"
+```
+
+To continue training from checkpoint
+```bash
+python imagenet.py --dataroot "/path/to/imagenet/" --resume "/path/to/checkpoint/folder"
+```
+## Results
+
+
+For x0.5 model I achieved 0.4% lower top-1 accuracy than claimed.
+
+|Classification Checkpoint| MACs (M)   | Parameters (M)| Top-1 Accuracy| Top-5 Accuracy|  Claimed top-1|  Claimed top-5|
+|-------------------------|------------|---------------|---------------|---------------|---------------|---------------|
+|      [shufflenet_v2_0.5]|41          |1.37           |          59.86|          81.63|           60.3|              -|
+
+You can test it with
+```bash
+python imagenet.py --dataroot "/path/to/imagenet/" --resume "results/shufflenet_v2_0.5/model_best.pth.tar" -e --scaling 0.5
+```

ShuffleNet/clr.py

+# temporary file until https://github.com/pytorch/pytorch/pull/2016 is merged (hopefully 0.5)
+
+
+import numpy as np
+from torch.optim import Optimizer
+
+
+class CyclicLR(object):
+    """Sets the learning rate of each parameter group according to
+    cyclical learning rate policy (CLR). The policy cycles the learning
+    rate between two boundaries with a constant frequency, as detailed in
+    the paper `Cyclical Learning Rates for Training Neural Networks`_.
+    The distance between the two boundaries can be scaled on a per-iteration
+    or per-cycle basis.
+    Cyclical learning rate policy changes the learning rate after every batch.
+    `batch_step` should be called after a batch has been used for training.
+    To resume training, save `last_batch_iteration` and use it to instantiate `CycleLR`.
+    This class has three built-in policies, as put forth in the paper:
+    "triangular":
+        A basic triangular cycle w/ no amplitude scaling.
+    "triangular2":
+        A basic triangular cycle that scales initial amplitude by half each cycle.
+    "exp_range":
+        A cycle that scales initial amplitude by gamma**(cycle iterations) at each
+        cycle iteration.
+    This implementation was adapted from the github repo: `bckenstler/CLR`_
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        base_lr (float or list): Initial learning rate which is the
+            lower boundary in the cycle for eachparam groups.
+            Default: 0.001
+        max_lr (float or list): Upper boundaries in the cycle for
+            each parameter group. Functionally,
+            it defines the cycle amplitude (max_lr - base_lr).
+            The lr at any cycle is the sum of base_lr
+            and some scaling of the amplitude; therefore
+            max_lr may not actually be reached depending on
+            scaling function. Default: 0.006
+        step_size (int): Number of training iterations per
+            half cycle. Authors suggest setting step_size
+            2-8 x training iterations in epoch. Default: 2000
+        mode (str): One of {triangular, triangular2, exp_range}.
+            Values correspond to policies detailed above.
+            If scale_fn is not None, this argument is ignored.
+            Default: 'triangular'
+        gamma (float): Constant in 'exp_range' scaling function:
+            gamma**(cycle iterations)
+            Default: 1.0
+        scale_fn (function): Custom scaling policy defined by a single
+            argument lambda function, where
+            0 <= scale_fn(x) <= 1 for all x >= 0.
+            mode paramater is ignored
+            Default: None
+        scale_mode (str): {'cycle', 'iterations'}.
+            Defines whether scale_fn is evaluated on
+            cycle number or cycle iterations (training
+            iterations since start of cycle).
+            Default: 'cycle'
+        last_batch_iteration (int): The index of the last batch. Default: -1
+    Example:
+        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+        >>> scheduler = torch.optim.CyclicLR(optimizer)
+        >>> data_loader = torch.utils.data.DataLoader(...)
+        >>> for epoch in range(10):
+        >>>     for batch in data_loader:
+        >>>         scheduler.batch_step()
+        >>>         train_batch(...)
+    .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
+    .. _bckenstler/CLR: https://github.com/bckenstler/CLR
+    """
+
+    def __init__(self, optimizer, base_lr=1e-3, max_lr=6e-3,
+                 step_size=2000, mode='triangular', gamma=1.,
+                 scale_fn=None, scale_mode='cycle', last_batch_iteration=-1):
+
+        if not isinstance(optimizer, Optimizer):
+            raise TypeError('{} is not an Optimizer'.format(
+                type(optimizer).__name__))
+        self.optimizer = optimizer
+
+        if isinstance(base_lr, list) or isinstance(base_lr, tuple):
+            if len(base_lr) != len(optimizer.param_groups):
+                raise ValueError("expected {} base_lr, got {}".format(
+                    len(optimizer.param_groups), len(base_lr)))
+            self.base_lrs = list(base_lr)
+        else:
+            self.base_lrs = [base_lr] * len(optimizer.param_groups)
+
+        if isinstance(max_lr, list) or isinstance(max_lr, tuple):
+            if len(max_lr) != len(optimizer.param_groups):
+                raise ValueError("expected {} max_lr, got {}".format(
+                    len(optimizer.param_groups), len(max_lr)))
+            self.max_lrs = list(max_lr)
+        else:
+            self.max_lrs = [max_lr] * len(optimizer.param_groups)
+
+        self.step_size = step_size
+
+        if mode not in ['triangular', 'triangular2', 'exp_range'] \
+                and scale_fn is None:
+            raise ValueError('mode is invalid and scale_fn is None')
+
+        self.mode = mode
+        self.gamma = gamma
+
+        if scale_fn is None:
+            if self.mode == 'triangular':
+                self.scale_fn = self._triangular_scale_fn
+                self.scale_mode = 'cycle'
+            elif self.mode == 'triangular2':
+                self.scale_fn = self._triangular2_scale_fn
+                self.scale_mode = 'cycle'
+            elif self.mode == 'exp_range':
+                self.scale_fn = self._exp_range_scale_fn
+                self.scale_mode = 'iterations'
+        else:
+            self.scale_fn = scale_fn
+            self.scale_mode = scale_mode
+
+        self.batch_step(last_batch_iteration + 1)
+        self.last_batch_iteration = last_batch_iteration
+
+    def batch_step(self, batch_iteration=None):
+        if batch_iteration is None:
+            batch_iteration = self.last_batch_iteration + 1
+        self.last_batch_iteration = batch_iteration
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+
+    def _triangular_scale_fn(self, x):
+        return 1.
+
+    def _triangular2_scale_fn(self, x):
+        return 1 / (2. ** (x - 1))
+
+    def _exp_range_scale_fn(self, x):
+        return self.gamma ** (x)
+
+    def get_lr(self):
+        step_size = float(self.step_size)
+        cycle = np.floor(1 + self.last_batch_iteration / (2 * step_size))
+        x = np.abs(self.last_batch_iteration / step_size - 2 * cycle + 1)
+
+        lrs = []
+        param_lrs = zip(self.optimizer.param_groups, self.base_lrs, self.max_lrs)
+        for param_group, base_lr, max_lr in param_lrs:
+            base_height = (max_lr - base_lr) * np.maximum(0, (1 - x))
+            if self.scale_mode == 'cycle':
+                lr = base_lr + base_height * self.scale_fn(cycle)
+            else:
+                lr = base_lr + base_height * self.scale_fn(self.last_batch_iteration)
+            lrs.append(lr)
+        return lrs

ShuffleNet/data.py

+import os
+
+import torch
+import torch.nn.parallel
+import torch.optim
+import torch.utils.data
+from torchvision import datasets, transforms
+
+__imagenet_stats = {'mean': [0.485, 0.456, 0.406],
+                    'std': [0.229, 0.224, 0.225]}
+
+
+def inception_preproccess(input_size, normalize=__imagenet_stats):
+    return transforms.Compose([
+        transforms.RandomResizedCrop(input_size),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(**normalize)
+    ])
+
+
+def scale_crop(input_size, scale_size=None, normalize=__imagenet_stats):
+    t_list = [
+        transforms.CenterCrop(input_size),
+        transforms.ToTensor(),
+        transforms.Normalize(**normalize),
+    ]
+    if scale_size != input_size:
+        t_list = [transforms.Resize(scale_size)] + t_list
+
+    return transforms.Compose(t_list)
+
+
+def get_transform(augment=True, input_size=224):
+    normalize = __imagenet_stats
+    scale_size = int(input_size / 0.875)
+    if augment:
+        return inception_preproccess(input_size=input_size, normalize=normalize)
+    else:
+        return scale_crop(input_size=input_size, scale_size=scale_size, normalize=normalize)
+
+
+def get_loaders(dataroot, val_batch_size, train_batch_size, input_size, workers):
+    val_data = datasets.ImageFolder(root=os.path.join(dataroot, 'val'), transform=get_transform(False, input_size))
+    val_loader = torch.utils.data.DataLoader(val_data, batch_size=val_batch_size, shuffle=False, num_workers=workers,
+                                             pin_memory=True)
+
+    train_data = datasets.ImageFolder(root=os.path.join(dataroot, 'train'),
+                                      transform=get_transform(input_size=input_size))
+    train_loader = torch.utils.data.DataLoader(train_data, batch_size=train_batch_size, shuffle=True,
+                                               num_workers=workers, pin_memory=True)
+    return train_loader, val_loader

ShuffleNet/flops_benchmark.py

+#### https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.py
+import torch
+
+
+# ---- Public functions
+
+def add_flops_counting_methods(net_main_module):
+    """Adds flops counting functions to an existing model. After that
+    the flops count should be activated and the model should be run on an input
+    image.
+
+    Example:
+
+    fcn = add_flops_counting_methods(fcn)
+    fcn = fcn.cuda().train()
+    fcn.start_flops_count()
+
+
+    _ = fcn(batch)
+
+    fcn.compute_average_flops_cost() / 1e9 / 2 # Result in GFLOPs per image in batch
+
+    Important: dividing by 2 only works for resnet models -- see below for the details
+    of flops computation.
+
+    Attention: we are counting multiply-add as two flops in this work, because in
+    most resnet models convolutions are bias-free (BN layers act as bias there)
+    and it makes sense to count muliply and add as separate flops therefore.
+    This is why in the above example we divide by 2 in order to be consistent with
+    most modern benchmarks. For example in "Spatially Adaptive Computatin Time for Residual
+    Networks" by Figurnov et al multiply-add was counted as two flops.
+
+    This module computes the average flops which is necessary for dynamic networks which
+    have different number of executed layers. For static networks it is enough to run the network
+    once and get statistics (above example).
+
+    Implementation:
+    The module works by adding batch_count to the main module which tracks the sum
+    of all batch sizes that were run through the network.
+
+    Also each convolutional layer of the network tracks the overall number of flops
+    performed.
+
+    The parameters are updated with the help of registered hook-functions which
+    are being called each time the respective layer is executed.
+
+    Parameters
+    ----------
+    net_main_module : torch.nn.Module
+        Main module containing network
+
+    Returns
+    -------
+    net_main_module : torch.nn.Module
+        Updated main module with new methods/attributes that are used
+        to compute flops.
+    """
+
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    net_main_module.start_flops_count = start_flops_count.__get__(net_main_module)
+    net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module)
+    net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module)
+    net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(net_main_module)
+
+    net_main_module.reset_flops_count()
+
+    # Adding varialbles necessary for masked flops computation
+    net_main_module.apply(add_flops_mask_variable_or_reset)
+
+    return net_main_module
+
+
+def compute_average_flops_cost(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Returns current mean flops consumption per image.
+
+    """
+
+    batches_count = self.__batch_counter__
+
+    flops_sum = 0
+
+    for module in self.modules():
+
+        if isinstance(module, torch.nn.Conv2d):
+            flops_sum += module.__flops__
+
+    return flops_sum / batches_count
+
+
+def start_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Activates the computation of mean flops consumption per image.
+    Call it before you run the network.
+
+    """
+
+    add_batch_counter_hook_function(self)
+
+    self.apply(add_flops_counter_hook_function)
+
+
+def stop_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Stops computing the mean flops consumption per image.
+    Call whenever you want to pause the computation.
+
+    """
+
+    remove_batch_counter_hook_function(self)
+
+    self.apply(remove_flops_counter_hook_function)
+
+
+def reset_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Resets statistics computed so far.
+
+    """
+
+    add_batch_counter_variables_or_reset(self)
+
+    self.apply(add_flops_counter_variable_or_reset)
+
+
+def add_flops_mask(module, mask):
+    def add_flops_mask_func(module):
+        if isinstance(module, torch.nn.Conv2d):
+            module.__mask__ = mask
+
+    module.apply(add_flops_mask_func)
+
+
+def remove_flops_mask(module):
+    module.apply(add_flops_mask_variable_or_reset)
+
+
+# ---- Internal functions
+
+
+def conv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+    output_height, output_width = output.shape[2:]
+
+    kernel_height, kernel_width = conv_module.kernel_size
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    # We count multiply-add as 2 flops
+    conv_per_position_flops = 2 * kernel_height * kernel_width * in_channels * out_channels / groups
+
+    active_elements_count = batch_size * output_height * output_width
+
+    if conv_module.__mask__ is not None:
+        # (b, 1, h, w)
+        flops_mask = conv_module.__mask__.expand(batch_size, 1, output_height, output_width)
+        active_elements_count = flops_mask.sum()
+
+    overall_conv_flops = conv_per_position_flops * active_elements_count
+
+    bias_flops = 0
+
+    if conv_module.bias is not None:
+        bias_flops = out_channels * active_elements_count
+
+    overall_flops = overall_conv_flops + bias_flops
+
+    conv_module.__flops__ += overall_flops
+
+
+def batch_counter_hook(module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+
+    module.__batch_counter__ += batch_size
+
+
+def add_batch_counter_variables_or_reset(module):
+    module.__batch_counter__ = 0
+
+
+def add_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        return
+
+    handle = module.register_forward_hook(batch_counter_hook)
+    module.__batch_counter_handle__ = handle
+
+
+def remove_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        module.__batch_counter_handle__.remove()
+
+        del module.__batch_counter_handle__
+
+
+def add_flops_counter_variable_or_reset(module):
+    if isinstance(module, torch.nn.Conv2d):
+        module.__flops__ = 0
+
+
+def add_flops_counter_hook_function(module):
+    if isinstance(module, torch.nn.Conv2d):
+
+        if hasattr(module, '__flops_handle__'):
+            return
+
+        handle = module.register_forward_hook(conv_flops_counter_hook)
+        module.__flops_handle__ = handle
+
+
+def remove_flops_counter_hook_function(module):
+    if isinstance(module, torch.nn.Conv2d):
+
+        if hasattr(module, '__flops_handle__'):
+            module.__flops_handle__.remove()
+
+            del module.__flops_handle__
+
+
+# --- Masked flops counting
+
+
+# Also being run in the initialization
+def add_flops_mask_variable_or_reset(module):
+    if isinstance(module, torch.nn.Conv2d):
+        module.__mask__ = None
+
+
+def count_flops(model, batch_size, device, dtype, input_size, in_channels, *params):
+    net = model(*params)
+    # print(net)
+    net = add_flops_counting_methods(net)
+
+    net.to(device=device, dtype=dtype)
+    net = net.train()
+
+    batch = torch.randn(batch_size, in_channels, input_size, input_size).to(device=device, dtype=dtype)
+    net.start_flops_count()
+
+    _ = net(batch)
+    return net.compute_average_flops_cost() / 2  # Result in FLOPs

ShuffleNet/imagenet.py

+import argparse
+import csv
+import os
+import random
+import sys
+from datetime import datetime
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn.parallel
+import torch.optim
+import torch.utils.data
+from torch.optim.lr_scheduler import MultiStepLR
+from tqdm import trange
+
+import flops_benchmark
+from clr import CyclicLR
+from data import get_loaders
+from logger import CsvLogger
+from model import ShuffleNetV2
+from run import train, test, save_checkpoint, find_bounds_clr
+
+parser = argparse.ArgumentParser(description='ShuffleNetv2 training with PyTorch')
+parser.add_argument('--dataroot', required=True, metavar='PATH',
+                    help='Path to ImageNet train and val folders, preprocessed as described in '
+                         'https://github.com/facebook/fb.resnet.torch/blob/master/INSTALL.md#download-the-imagenet-dataset')
+parser.add_argument('--gpus', default=None, help='List of GPUs used for training - e.g 0,1,3')
+parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                    help='Number of data loading workers (default: 4)')
+parser.add_argument('--type', default='float32', help='Type of tensor: float32, float16, float64. Default: float32')
+
+# Optimization options
+parser.add_argument('--epochs', type=int, default=400, help='Number of epochs to train.')
+parser.add_argument('-b', '--batch-size', default=64, type=int, metavar='N', help='mini-batch size (default: 64)')
+parser.add_argument('--learning_rate', '-lr', type=float, default=0.01, help='The learning rate.')
+parser.add_argument('--momentum', '-m', type=float, default=0.9, help='Momentum.')
+parser.add_argument('--decay', '-d', type=float, default=4e-5, help='Weight decay (L2 penalty).')
+parser.add_argument('--gamma', type=float, default=0.1, help='LR is multiplied by gamma at scheduled epochs.')
+parser.add_argument('--schedule', type=int, nargs='+', default=[200, 300],
+                    help='Decrease learning rate at these epochs.')
+
+# CLR
+parser.add_argument('--clr', dest='clr', action='store_true', help='Use CLR')
+parser.add_argument('--min-lr', type=float, default=1e-5, help='Minimal LR for CLR.')
+parser.add_argument('--max-lr', type=float, default=1, help='Maximal LR for CLR.')
+parser.add_argument('--epochs-per-step', type=int, default=20,
+                    help='Number of epochs per step in CLR, recommended to be between 2 and 10.')
+parser.add_argument('--mode', default='triangular2', help='CLR mode. One of {triangular, triangular2, exp_range}')
+parser.add_argument('--find-clr', dest='find_clr', action='store_true',
+                    help='Run search for optimal LR in range (min_lr, max_lr)')
+
+# Checkpoints
+parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='Just evaluate model')
+parser.add_argument('--save', '-s', type=str, default='', help='Folder to save checkpoints.')
+parser.add_argument('--results_dir', metavar='RESULTS_DIR', default='./results', help='Directory to store results')
+parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
+parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')
+parser.add_argument('--log-interval', type=int, default=100, metavar='N', help='Number of batches between log messages')
+parser.add_argument('--seed', type=int, default=None, metavar='S', help='random seed (default: random)')
+
+# Architecture
+parser.add_argument('--scaling', type=float, default=1, metavar='SC', help='Scaling of ShuffleNet (default x1).')
+parser.add_argument('--input-size', type=int, default=224, metavar='I', help='Input size of ShuffleNet.')
+parser.add_argument('--c-tag', type=float, default=0.5, help="c' value")
+parser.add_argument('--SE', dest='SE', action='store_true', help='Use SE modules')
+parser.add_argument('--residual', dest='residual', action='store_true', help='Just residuals')
+parser.add_argument('--groups', type=int, default=2,  help='Groups per shuffle.')
+
+# https://arxiv.org/abs/1807.11164
+#g, SE, scale
+claimed_acc_top1 = {False: {0.5: 0.397, 1.: 0.306, 1.5: 0.274, 2.: 0.251}, True: {2.:0.346}}
+
+def main():
+    args = parser.parse_args()
+
+    if args.seed is None:
+        args.seed = random.randint(1, 10000)
+    print("Random Seed: ", args.seed)
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if args.gpus:
+        torch.cuda.manual_seed_all(args.seed)
+
+    time_stamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+    if args.evaluate:
+        args.results_dir = '/tmp'
+    if args.save is '':
+        args.save = time_stamp
+    save_path = os.path.join(args.results_dir, args.save)
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+
+    if args.gpus is not None:
+        args.gpus = [int(i) for i in args.gpus.split(',')]
+        device = 'cuda:' + str(args.gpus[0])
+        cudnn.benchmark = True
+    else:
+        device = 'cpu'
+
+    if args.type == 'float64':
+        dtype = torch.float64
+    elif args.type == 'float32':
+        dtype = torch.float32
+    elif args.type == 'float16':
+        dtype = torch.float16
+    else:
+        raise ValueError('Wrong type!')  # TODO int8
+
+    model = ShuffleNetV2(scale=args.scaling, c_tag=args.c_tag, SE=args.SE, residual=args.residual, groups=args.groups)
+    num_parameters = sum([l.nelement() for l in model.parameters()])
+    print(model)
+    print('number of parameters: {}'.format(num_parameters))
+    print('FLOPs: {}'.format(
+        flops_benchmark.count_flops(ShuffleNetV2,
+                                    args.batch_size // len(args.gpus) if args.gpus is not None else args.batch_size,
+                                    device, dtype, args.input_size, 3, args.scaling, 3,  args.c_tag, 1000, torch.nn.ReLU,
+                                    args.SE, args.residual, args.groups)))
+
+    train_loader, val_loader = get_loaders(args.dataroot, args.batch_size, args.batch_size, args.input_size,
+                                           args.workers)
+    # define loss function (criterion) and optimizer
+    criterion = torch.nn.CrossEntropyLoss()
+    if args.gpus is not None:
+        model = torch.nn.DataParallel(model, args.gpus)
+    model.to(device=device, dtype=dtype)
+    criterion.to(device=device, dtype=dtype)
+
+    optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.decay,
+                                nesterov=True)
+    if args.find_clr:
+        find_bounds_clr(model, train_loader, optimizer, criterion, device, dtype, min_lr=args.min_lr,
+                        max_lr=args.max_lr, step_size=args.epochs_per_step * len(train_loader), mode=args.mode,
+                        save_path=save_path)
+        return
+
+    if args.clr:
+        scheduler = CyclicLR(optimizer, base_lr=args.min_lr, max_lr=args.max_lr,
+                             step_size=args.epochs_per_step * len(train_loader), mode=args.mode)
+    else:
+        scheduler = MultiStepLR(optimizer, milestones=args.schedule, gamma=args.gamma)
+
+    best_test = 0
+
+    # optionally resume from a checkpoint
+    data = None
+    if args.resume:
+        if os.path.isfile(args.resume):
+            print("=> loading checkpoint '{}'".format(args.resume))
+            checkpoint = torch.load(args.resume, map_location=device)
+            args.start_epoch = checkpoint['epoch'] - 1
+            best_test = checkpoint['best_prec1']
+            model.load_state_dict(checkpoint['state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            print("=> loaded checkpoint '{}' (epoch {})"
+                  .format(args.resume, checkpoint['epoch']))
+        elif os.path.isdir(args.resume):
+            checkpoint_path = os.path.join(args.resume, 'checkpoint.pth.tar')
+            csv_path = os.path.join(args.resume, 'results.csv')
+            print("=> loading checkpoint '{}'".format(checkpoint_path))
+            checkpoint = torch.load(checkpoint_path, map_location=device)
+            args.start_epoch = checkpoint['epoch'] - 1
+            best_test = checkpoint['best_prec1']
+            model.load_state_dict(checkpoint['state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            print("=> loaded checkpoint '{}' (epoch {})".format(checkpoint_path, checkpoint['epoch']))
+            data = []
+            with open(csv_path) as csvfile:
+                reader = csv.DictReader(csvfile)
+                for row in reader:
+                    data.append(row)
+        else:
+            print("=> no checkpoint found at '{}'".format(args.resume))
+
+    if args.evaluate:
+        loss, top1, top5 = test(model, val_loader, criterion, device, dtype)  # TODO
+        return
+
+    csv_logger = CsvLogger(filepath=save_path, data=data)
+    csv_logger.save_params(sys.argv, args)
+
+    claimed_acc1 = None
+    claimed_acc5 = None
+    if args.SE in claimed_acc_top1:
+        if args.scaling in claimed_acc_top1[args.SE]:
+            claimed_acc1 = 1 - claimed_acc_top1[args.SE][args.scaling]
+            csv_logger.write_text('Claimed accuracy is {:.2f}% top-1'.format(claimed_acc1 * 100.))
+    train_network(args.start_epoch, args.epochs, scheduler, model, train_loader, val_loader, optimizer, criterion,
+                  device, dtype, args.batch_size, args.log_interval, csv_logger, save_path, claimed_acc1, claimed_acc5,
+                  best_test)
+
+
+def train_network(start_epoch, epochs, scheduler, model, train_loader, val_loader, optimizer, criterion, device, dtype,
+                  batch_size, log_interval, csv_logger, save_path, claimed_acc1, claimed_acc5, best_test):
+    for epoch in trange(start_epoch, epochs + 1):
+        if not isinstance(scheduler, CyclicLR):
+            scheduler.step()
+        train_loss, train_accuracy1, train_accuracy5, = train(model, train_loader, epoch, optimizer, criterion, device,
+                                                              dtype, batch_size, log_interval, scheduler)
+        test_loss, test_accuracy1, test_accuracy5 = test(model, val_loader, criterion, device, dtype)
+        csv_logger.write({'epoch': epoch + 1, 'val_error1': 1 - test_accuracy1, 'val_error5': 1 - test_accuracy5,
+                          'val_loss': test_loss, 'train_error1': 1 - train_accuracy1,
+                          'train_error5': 1 - train_accuracy5, 'train_loss': train_loss})
+        save_checkpoint({'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_prec1': best_test,
+                         'optimizer': optimizer.state_dict()}, test_accuracy1 > best_test, filepath=save_path)
+
+        csv_logger.plot_progress(claimed_acc1=claimed_acc1, claimed_acc5=claimed_acc5)
+
+        if test_accuracy1 > best_test:
+            best_test = test_accuracy1
+
+    csv_logger.write_text('Best accuracy is {:.2f}% top-1'.format(best_test * 100.))
+
+
+if __name__ == '__main__':
+    main()

ShuffleNet/logger.py

+import csv
+import os.path
+
+import matplotlib
+
+matplotlib.use('Agg')
+
+from matplotlib import pyplot as plt
+import numpy as np
+
+plt.switch_backend('agg')
+
+
+class CsvLogger:
+    def __init__(self, filepath='./', filename='results.csv', data=None):
+        self.log_path = filepath
+        self.log_name = filename
+        self.csv_path = os.path.join(self.log_path, self.log_name)
+        self.fieldsnames = ['epoch', 'val_error1', 'val_error5', 'val_loss', 'train_error1', 'train_error5',
+                            'train_loss']
+
+        with open(self.csv_path, 'w') as f:
+            writer = csv.DictWriter(f, fieldnames=self.fieldsnames)
+            writer.writeheader()
+
+        self.data = {}
+        for field in self.fieldsnames:
+            self.data[field] = []
+        if data is not None:
+            for d in data:
+                d_num = {}
+                for key in d:
+                    d_num[key] = float(d[key]) if key != 'epoch' else int(d[key])
+                self.write(d_num)
+
+    def write(self, data):
+        for k in self.data:
+            self.data[k].append(data[k])
+        with open(self.csv_path, 'a') as f:
+            writer = csv.DictWriter(f, fieldnames=self.fieldsnames)
+            writer.writerow(data)
+
+    def save_params(self, args, params):
+        with open(os.path.join(self.log_path, 'params.txt'), 'w') as f:
+            f.write('{}\n'.format(' '.join(args)))
+            f.write('{}\n'.format(params))
+
+    def write_text(self, text, print_t=True):
+        with open(os.path.join(self.log_path, 'params.txt'), 'a') as f:
+            f.write('{}\n'.format(text))
+        if print_t:
+            print(text)
+
+    def plot_progress_errk(self, claimed_acc=None, title='ShuffleNetv2', k=1):
+        tr_str = 'train_error{}'.format(k)
+        val_str = 'val_error{}'.format(k)
+        plt.figure(figsize=(9, 8), dpi=300)
+        plt.plot(self.data[tr_str], label='Training error')
+        plt.plot(self.data[val_str], label='Validation error')
+        if claimed_acc is not None:
+            plt.plot((0, len(self.data[tr_str])), (1 - claimed_acc, 1 - claimed_acc), 'k--',
+                     label='Claimed validation error ({:.2f}%)'.format(100. * (1 - claimed_acc)))
+        plt.plot((0, len(self.data[tr_str])),
+                 (np.min(self.data[val_str]), np.min(self.data[val_str])), 'r--',
+                 label='Best validation error ({:.2f}%)'.format(100. * np.min(self.data[val_str])))
+        plt.title('Top-{} error for {}'.format(k, title))
+        plt.xlabel('Epoch')
+        plt.ylabel('Error')
+        plt.legend()
+        plt.xlim(0, len(self.data[tr_str]) + 1)
+        plt.savefig(os.path.join(self.log_path, 'top{}.png'.format(k)))
+
+    def plot_progress_loss(self, title='ShuffleNetv2'):
+        plt.figure(figsize=(9, 8), dpi=300)
+        plt.plot(self.data['train_loss'], label='Training')
+        plt.plot(self.data['val_loss'], label='Validation')
+        plt.title(title)
+        plt.xlabel('Epoch')
+        plt.ylabel('Loss')
+        plt.legend()
+        plt.xlim(0, len(self.data['train_loss']) + 1)
+        plt.savefig(os.path.join(self.log_path, 'loss.png'))
+
+    def plot_progress(self, claimed_acc1=None, claimed_acc5=None, title='ShuffleNetv2'):
+        self.plot_progress_errk(claimed_acc1, title, 1)
+        self.plot_progress_errk(claimed_acc5, title, 5)
+        self.plot_progress_loss(title)
+        plt.close('all')

ShuffleNet/model.py

+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+def channel_shuffle(x, groups):
+    batchsize, num_channels, height, width = x.data.size()
+    assert (num_channels % groups == 0)
+    channels_per_group = num_channels // groups
+    # reshape
+    x = x.view(batchsize, groups, channels_per_group, height, width)
+
+    # transpose
+    # - contiguous() required if transpose() is used before view().
+    #   See https://github.com/pytorch/pytorch/issues/764
+    x = torch.transpose(x, 1, 2).contiguous()
+
+    # flatten
+    x = x.view(batchsize, -1, height, width)
+
+    return x
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+class BasicUnit(nn.Module):
+    def __init__(self, inplanes, outplanes, c_tag=0.5, activation=nn.ReLU, SE=False, residual=False, groups=2):
+        super(BasicUnit, self).__init__()
+        self.left_part = round(c_tag * inplanes)
+        self.right_part_in = inplanes - self.left_part
+        self.right_part_out = outplanes - self.left_part
+        self.conv1 = nn.Conv2d(self.right_part_in, self.right_part_out, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.right_part_out)
+        self.conv2 = nn.Conv2d(self.right_part_out, self.right_part_out, kernel_size=3, padding=1, bias=False,
+                               groups=self.right_part_out)
+        self.bn2 = nn.BatchNorm2d(self.right_part_out)
+        self.conv3 = nn.Conv2d(self.right_part_out, self.right_part_out, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.right_part_out)
+        self.activation = activation(inplace=True)
+
+        self.inplanes = inplanes
+        self.outplanes = outplanes
+        self.residual = residual
+        self.groups = groups
+        self.SE = SE
+        if self.SE:
+            self.SELayer = SELayer(self.right_part_out, 2)  # TODO
+
+    def forward(self, x):
+        left = x[:, :self.left_part, :, :]
+        right = x[:, self.left_part:, :, :]
+        out = self.conv1(right)
+        out = self.bn1(out)
+        out = self.activation(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        out = self.activation(out)
+
+        if self.SE:
+            out = self.SELayer(out)
+        if self.residual and self.inplanes == self.outplanes:
+            out += right
+
+        return channel_shuffle(torch.cat((left, out), 1), self.groups)
+
+
+class DownsampleUnit(nn.Module):
+    def __init__(self, inplanes, c_tag=0.5, activation=nn.ReLU, groups=2):
+        super(DownsampleUnit, self).__init__()
+
+        self.conv1r = nn.Conv2d(inplanes, inplanes, kernel_size=1, bias=False)
+        self.bn1r = nn.BatchNorm2d(inplanes)
+        self.conv2r = nn.Conv2d(inplanes, inplanes, kernel_size=3, stride=2, padding=1, bias=False, groups=inplanes)
+        self.bn2r = nn.BatchNorm2d(inplanes)
+        self.conv3r = nn.Conv2d(inplanes, inplanes, kernel_size=1, bias=False)
+        self.bn3r = nn.BatchNorm2d(inplanes)
+
+        self.conv1l = nn.Conv2d(inplanes, inplanes, kernel_size=3, stride=2, padding=1, bias=False, groups=inplanes)
+        self.bn1l = nn.BatchNorm2d(inplanes)
+        self.conv2l = nn.Conv2d(inplanes, inplanes, kernel_size=1, bias=False)
+        self.bn2l = nn.BatchNorm2d(inplanes)
+        self.activation = activation(inplace=True)
+
+        self.groups = groups
+        self.inplanes = inplanes
+
+    def forward(self, x):
+        out_r = self.conv1r(x)
+        out_r = self.bn1r(out_r)
+        out_r = self.activation(out_r)
+
+        out_r = self.conv2r(out_r)
+        out_r = self.bn2r(out_r)
+
+        out_r = self.conv3r(out_r)
+        out_r = self.bn3r(out_r)
+        out_r = self.activation(out_r)
+
+        out_l = self.conv1l(x)
+        out_l = self.bn1l(out_l)
+
+        out_l = self.conv2l(out_l)
+        out_l = self.bn2l(out_l)
+        out_l = self.activation(out_l)
+
+        return channel_shuffle(torch.cat((out_r, out_l), 1), self.groups)
+
+
+class ShuffleNetV2(nn.Module):
+    """ShuffleNetV2 implementation.
+    """
+
+    def __init__(self, scale=1.0, in_channels=3, c_tag=0.5, num_classes=1000, activation=nn.ReLU,
+                 SE=False, residual=False, groups=2):
+        """
+        ShuffleNetV2 constructor
+        :param scale:
+        :param in_channels:
+        :param c_tag:
+        :param num_classes:
+        :param activation:
+        :param SE:
+        :param residual:
+        :param groups:
+        """
+
+        super(ShuffleNetV2, self).__init__()
+
+        self.scale = scale
+        self.c_tag = c_tag
+        self.residual = residual
+        self.SE = SE
+        self.groups = groups
+
+        self.activation_type = activation
+        self.activation = activation(inplace=True)
+        self.num_classes = num_classes
+
+        self.num_of_channels = {0.5: [24, 48, 96, 192, 1024], 1: [24, 116, 232, 464, 1024],
+                                1.5: [24, 176, 352, 704, 1024], 2: [24, 244, 488, 976, 2048]}
+        self.c = [_make_divisible(chan, groups) for chan in self.num_of_channels[scale]]
+        self.n = [3, 8, 3]  # TODO: should be [3,7,3]
+        self.conv1 = nn.Conv2d(in_channels, self.c[0], kernel_size=3, bias=False, stride=2, padding=1)
+        self.bn1 = nn.BatchNorm2d(self.c[0])
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)
+        self.shuffles = self._make_shuffles()
+
+        self.conv_last = nn.Conv2d(self.c[-2], self.c[-1], kernel_size=1, bias=False)
+        self.bn_last = nn.BatchNorm2d(self.c[-1])
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(self.c[-1], self.num_classes)
+        self.init_params()
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.001)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def _make_stage(self, inplanes, outplanes, n, stage):
+        modules = OrderedDict()
+        stage_name = "ShuffleUnit{}".format(stage)
+
+        # First module is the only one utilizing stride
+        first_module = DownsampleUnit(inplanes=inplanes, activation=self.activation_type, c_tag=self.c_tag,
+                                      groups=self.groups)
+        modules["DownsampleUnit"] = first_module
+        second_module = BasicUnit(inplanes=inplanes * 2, outplanes=outplanes, activation=self.activation_type,
+                                  c_tag=self.c_tag, SE=self.SE, residual=self.residual, groups=self.groups)
+        modules[stage_name + "_{}".format(0)] = second_module
+        # add more LinearBottleneck depending on number of repeats
+        for i in range(n - 1):
+            name = stage_name + "_{}".format(i + 1)
+            module = BasicUnit(inplanes=outplanes, outplanes=outplanes, activation=self.activation_type,
+                               c_tag=self.c_tag, SE=self.SE, residual=self.residual, groups=self.groups)
+            modules[name] = module
+
+        return nn.Sequential(modules)
+
+    def _make_shuffles(self):
+        modules = OrderedDict()
+        stage_name = "ShuffleConvs"
+
+        for i in range(len(self.c) - 2):
+            name = stage_name + "_{}".format(i)
+            module = self._make_stage(inplanes=self.c[i], outplanes=self.c[i + 1], n=self.n[i], stage=i)
+            modules[name] = module
+
+        return nn.Sequential(modules)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.activation(x)
+        x = self.maxpool(x)
+
+        x = self.shuffles(x)
+        x = self.conv_last(x)
+        x = self.bn_last(x)
+        x = self.activation(x)
+
+        # average pooling layer
+        x = self.avgpool(x)
+
+        # flatten for input to fully-connected layer
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+
+
+if __name__ == "__main__":
+    """Testing
+    """
+    model1 = ShuffleNetV2()
+    print(model1)
+    model2 = ShuffleNetV2(scale=0.5, in_channels=3, c_tag=0.5, num_classes=1000, activation=nn.ReLU,
+                          SE=False, residual=False)
+    print(model2)
+    model3 = ShuffleNetV2(in_channels=2, num_classes=10)
+    print(model3)
+    x = torch.randn(1, 2, 224, 224)
+    print(model3(x))
+    model4 = ShuffleNetV2( num_classes=10, groups=3, c_tag=0.2)
+    print(model4)
+    model4_size = 769
+    x2 = torch.randn(1, 3, model4_size, model4_size, )
+    print(model4(x2))
+    model5 = ShuffleNetV2(scale=2.0,num_classes=10, SE=True, residual=True)
+    x3 = torch.randn(1, 3, 196, 196)
+    print(model5(x3))

ShuffleNet/requirements.txt

+torch>=0.4.0
+torchvision>=0.1.9
+tqdm>=4.19.4
+matplotlib
+numpy
\ No newline at end of file

ShuffleNet/run.py

+import os
+import shutil
+
+import matplotlib
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+import torch.utils.data
+from tqdm import tqdm, trange
+
+matplotlib.use('Agg')
+
+from matplotlib import pyplot as plt
+
+from clr import CyclicLR
+
+
+def train(model, loader, epoch, optimizer, criterion, device, dtype, batch_size, log_interval, scheduler):
+    model.train()
+    correct1, correct5 = 0, 0
+
+    for batch_idx, (data, target) in enumerate(tqdm(loader)):
+        if isinstance(scheduler, CyclicLR):
+            scheduler.batch_step()
+        data, target = data.to(device=device, dtype=dtype), target.to(device=device)
+
+        optimizer.zero_grad()
+        output = model(data)
+
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+        corr = correct(output, target, topk=(1, 5))
+        correct1 += corr[0]
+        correct5 += corr[1]
+
+        if batch_idx % log_interval == 0:
+            tqdm.write(
+                'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}. '
+                'Top-1 accuracy: {:.2f}%({:.2f}%). '
+                'Top-5 accuracy: {:.2f}%({:.2f}%).'.format(epoch, batch_idx, len(loader),
+                                                           100. * batch_idx / len(loader), loss.item(),
+                                                           100. * corr[0] / batch_size,
+                                                           100. * correct1 / (batch_size * (batch_idx + 1)),
+                                                           100. * corr[1] / batch_size,
+                                                           100. * correct5 / (batch_size * (batch_idx + 1))))
+    return loss.item(), correct1 / len(loader.dataset), correct5 / len(loader.dataset)
+
+
+def test(model, loader, criterion, device, dtype):
+    model.eval()
+    test_loss = 0
+    correct1, correct5 = 0, 0
+
+    for batch_idx, (data, target) in enumerate(tqdm(loader)):
+        data, target = data.to(device=device, dtype=dtype), target.to(device=device)
+        with torch.no_grad():
+            output = model(data)
+            test_loss += criterion(output, target).item()  # sum up batch loss
+            corr = correct(output, target, topk=(1, 5))
+        correct1 += corr[0]
+        correct5 += corr[1]
+
+    test_loss /= len(loader)
+
+    tqdm.write(
+        '\nTest set: Average loss: {:.4f}, Top1: {}/{} ({:.2f}%), '
+        'Top5: {}/{} ({:.2f}%)'.format(test_loss, int(correct1), len(loader.dataset),
+                                       100. * correct1 / len(loader.dataset), int(correct5),
+                                       len(loader.dataset), 100. * correct5 / len(loader.dataset)))
+    return test_loss, correct1 / len(loader.dataset), correct5 / len(loader.dataset)
+
+
+def correct(output, target, topk=(1,)):
+    """Computes the correct@k for the specified values of k"""
+    maxk = max(topk)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t().type_as(target)
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].reshape(-1).float().sum(0).item()
+        res.append(correct_k)
+    return res
+
+
+def save_checkpoint(state, is_best, filepath='./', filename='checkpoint.pth.tar'):
+    save_path = os.path.join(filepath, filename)
+    best_path = os.path.join(filepath, 'model_best.pth.tar')
+    torch.save(state, save_path)
+    if is_best:
+        shutil.copyfile(save_path, best_path)
+
+
+def find_bounds_clr(model, loader, optimizer, criterion, device, dtype, min_lr=8e-6, max_lr=8e-5, step_size=2000,
+                    mode='triangular', save_path='.'):
+    model.train()
+    correct1, correct5 = 0, 0
+    scheduler = CyclicLR(optimizer, base_lr=min_lr, max_lr=max_lr, step_size=step_size, mode=mode)
+    epoch_count = step_size // len(loader)  # Assuming step_size is multiple of batch per epoch
+    accuracy = []
+    for _ in trange(epoch_count):
+        for batch_idx, (data, target) in enumerate(tqdm(loader)):
+            if scheduler is not None:
+                scheduler.batch_step()
+            data, target = data.to(device=device, dtype=dtype), target.to(device=device)
+
+            optimizer.zero_grad()
+            output = model(data)
+
+            loss = criterion(output, target)
+            loss.backward()
+            optimizer.step()
+
+            corr = correct(output, target)
+            accuracy.append(corr[0] / data.shape[0])
+
+    lrs = np.linspace(min_lr, max_lr, step_size)
+    plt.plot(lrs, accuracy)
+    plt.show()
+    plt.savefig(os.path.join(save_path, 'find_bounds_clr.png'))
+    np.save(os.path.join(save_path, 'acc.npy'), accuracy)
+    return

ShuffleNetV2Driver.py

+
+from fitlog import FitLog
+from ShuffleNet.model import ShuffleNetV2
+from iframe_feeder import iframe_feeder
+from torch.utils.data import DataLoader
+import torch
+from tqdm import tqdm, trange
+from torch.optim.lr_scheduler import MultiStepLR
+
+
+
+class ShuffleNetV2Driver:
+    def __init__(self, featd, feati, normd, normi, lab):
+        self.batchsize = 256
+        self.lr = 0.01
+        self.momentum = 0.9
+        self.decay = 4e-5
+        self.gamma = 0.1
+        self.schedule = [200, 300]
+
+        self.feeder = iframe_feeder(featd, feati, lab, normd, normi)
+        self.feeder.set_mode('train')
+
+        self.model = ShuffleNetV2(num_classes=2, scale=0.5, SE=True, residual=True)#torchvision.models.ShuffleNetV2(num_classes=2)#
+        self.fitlog = FitLog()
+        self.detail_log = FitLog(prefix='dt_')
+        self.device = self.get_device() #torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print("device info: " + str(self.device))
+
+        self.optimizer = torch.optim.SGD(self.model.parameters(),\
+            self.lr, momentum=self.momentum, weight_decay=self.decay,\
+            nesterov=True)
+        self.scheduler = MultiStepLR(self.optimizer, self.schedule, self.gamma)
+        self.criterion = torch.nn.CrossEntropyLoss()
+        self.loader = DataLoader(self.feeder, batch_size=self.batchsize, shuffle=True, num_workers=0)
+
+        self.print_interval = 25
+        self.n_epoch = 200
+
+        self.model.to(self.device)
+
+    def get_device(self):
+        device_mlu = None
+        device_gpu = None
+        try:
+            # device_mlu = torch.device('mlu')
+            device_gpu = torch.device('cuda')
+        except Exception as err:
+            print(err)
+        if device_mlu:
+            self.fitlog.append('mlu', True, True)
+            return device_mlu
+        elif device_gpu:
+            self.fitlog.append('cuda', True, True)
+            return device_gpu
+        else:
+            self.fitlog.append('cpu', True, True)
+            return torch.device('cpu')
+
+    def train(self, epoch):
+        self.feeder.set_mode('train')
+        self.model.train()
+        nbatch = len(self.loader)
+
+        for batch_idx, (feats, labs) in enumerate(tqdm(self.loader)):
+            feats = feats.to(self.device)
+            labs = labs.to(self.device)
+            self.optimizer.zero_grad()
+            res = self.model(feats)
+            loss = self.criterion(res, labs)
+            loss.backward()
+            self.optimizer.step()
+            
+            if batch_idx % self.print_interval == 0:
+                xstr = "Train: epoch: {} batch: {}/{}, loss: {:.6f}".format(epoch, batch_idx, nbatch, loss)
+                tqdm.write(xstr)
+                self.fitlog.append(xstr, True, True)
+                
+    def validate(self):
+        self.feeder.set_mode('test')
+        self.model.eval()
+        loss_val = 0
+        n_correct = 0
+        n_total = 0
+
+        for batch_idx, (feats, labs) in enumerate(tqdm(self.loader)):
+            feats = feats.to(self.device)
+            labs = labs.to(self.device)
+            with torch.no_grad():
+                res = self.model(feats)
+                loss_val += self.criterion(res, labs).item()
+                _, pred = res.max(1)
+                n_correct += pred.eq(labs).sum().item()
+                n_total += labs.shape[0]
+        
+        loss_val = loss_val / len(self.loader)
+        acc = n_correct / n_total * 100
+
+        self.detail_log.append(str(labs.tolist()))
+        self.detail_log.append(str(pred.tolist()))
+
+        xstr = "Validation: avg loss: {:.4f}, avg acc: {:.4f}%".format(loss_val, acc)
+        tqdm.write(xstr)
+        self.fitlog.append(xstr, True, True)
+
+    def finish(self):
+        self.feeder.finish()
+        self.fitlog.close()
+        self.detail_log.close()
+
+    def run(self):
+        for i in range(1, self.n_epoch + 1):
+            self.train(i)
+            self.validate()
+        
+
+if __name__ == "__main__":
+    datafolder = "data/"
+
+    driver = ShuffleNetV2Driver(datafolder + "s2_ftimgd",\
+                                datafolder + "s2_ftimgi",\
+                                datafolder + "s2_normd",\
+                                datafolder + "s2_normi",\
+                                datafolder + "s2_label.json")
+    driver.run()
+    driver.finish()
+

analysis.py

+import sys
+import path_config
+sys.path.append(path_config.get_pese_path())
+from ANSAFF.py.tool import signal_detection
+
+def analysis4file(file):
+    labs = []
+    preds = []
+    fh = open(file, 'r')
+    i = 0
+    for line in fh:
+        xline = line[1:len(line)-1]
+        xline = xline.split(',')
+        xline = [float(i) for i in xline]
+        if i % 2 == 0:
+            labs.append(xline[0])
+        else:
+            preds.append(xline[0])
+        i += 1
+    fh.close()
+    tp,fp,tn,fn = signal_detection(labs, preds)
+    precision = tp/(tp+fp)
+    recall = tp/(tp+fn)
+    f1 = 2*precision*recall/(precision+recall)
+    print("precision: " + str(precision) + ", recall: " + str(recall) + ", f1: " + str(f1))