trident
trident
Table Of Contents
trident
Table Of Contents

Source code for trident.models.pytorch_mobilenet

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import inspect
import math
import os
import uuid
from collections import *
from collections import deque
from copy import copy, deepcopy
from functools import partial
from itertools import repeat

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch._six import container_abcs
from torch.nn import init
from torch.nn.parameter import Parameter

from trident.backend.common import *
from trident.backend.pytorch_backend import to_numpy, to_tensor, Layer, Sequential
from trident.data.image_common import *
from trident.data.utils import download_model_from_google_drive
from trident.layers.pytorch_activations import get_activation, Identity
from trident.layers.pytorch_blocks import *
from trident.layers.pytorch_layers import *
from trident.layers.pytorch_normalizations import get_normalization
from trident.layers.pytorch_pooling import *
from trident.optims.pytorch_trainer import *

__all__ = ['MobileNet','MobileNetV2']

_session = get_session()
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
_epsilon=_session.epsilon
_trident_dir=_session.trident_dir


dirname = os.path.join(_trident_dir, 'models')
if not os.path.exists(dirname):
    try:
        os.makedirs(dirname)
    except OSError:
        # Except permission denied and potential race conditions
        # in multi-threaded environments.
        pass


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
    Args
        v:
        divisor:
        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 inverted_residual(in_filters,num_filters=64,strides=1,expansion = 4,name=''):
    mid_filters= int(round(in_filters * expansion))
    layers=[]
    if expansion!=1 :
        layers.append(Conv2d_Block((1,1),num_filters=mid_filters,strides=1,auto_pad=True,padding_mode='zero',normalization='batch',activation='relu6',name=name + '_{0}_conv'.format(len(layers))))

    layers.append(DepthwiseConv2d_Block((3, 3), depth_multiplier=1, strides=strides, auto_pad=True,padding_mode='zero', normalization='batch', activation='relu6', name=name + '_{0}_conv'.format(len(layers))))
    layers.append(Conv2d_Block((1, 1), num_filters=num_filters, strides=1, auto_pad=False, padding_mode='zero', normalization='batch', activation=None, name=name + '_{0}_conv'.format(len(layers))))
    if  strides == 1 and in_filters==num_filters:
        return ShortCut2d(Sequential(*layers), Identity(), activation=None)
    else:
        return Sequential(*layers)

[docs]def MobileNet( input_shape=(3, 224, 224), classes=1000, use_bias=False, width_mult=1.0,round_nearest=8, include_top=True, model_name='', **kwargs): input_filters = 32 last_filters = 1280 mobilenet=Sequential(name='mobilenet') inverted_residual_setting = [ # t, c, n, s [1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2], [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2], [6, 320, 1, 1], ] input_filters = _make_divisible(input_filters * width_mult, round_nearest) last_filters = _make_divisible(last_filters * max(1.0, width_mult), round_nearest) features = [] features.append(Conv2d_Block((3,3),num_filters=input_filters,strides=2,auto_pad=True,padding_mode='zero',normalization='batch',activation='relu6',name='first_layer')) for t, c, n, s in inverted_residual_setting: output_filters = _make_divisible(c * width_mult, round_nearest) for i in range(n): strides = s if i == 0 else 1 features.append( inverted_residual(input_filters,num_filters=output_filters, strides=strides, expansion=t,name='irb_{0}'.format(i))) input_filters = output_filters features.append(Conv2d_Block((1,1), last_filters,auto_pad=True,padding_mode='zero',normalization='batch',activation='relu6',name='last_layer')) mobilenet.add_module('features',Sequential(*features,name='features')) mobilenet.add_module('gap',GlobalAvgPool2d()) if include_top: mobilenet.add_module('drop', Dropout(0.2)) mobilenet.add_module('fc',Dense((classes),activation=None)) mobilenet.add_module('softmax', SoftMax(name='softmax')) model = ImageClassificationModel(input_shape=input_shape, output=mobilenet) model.signature=get_signature(model.model.forward) with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'imagenet_labels1.txt'), 'r', encoding='utf-8-sig') as f: labels = [l.rstrip() for l in f] model.class_names = labels model.preprocess_flow = [resize((224, 224), keep_aspect=True), normalize(0, 255), normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])] # model.summary() return model
[docs]def MobileNetV2(include_top=True, pretrained=True, input_shape=(3,224,224), classes=1000, **kwargs): if input_shape is not None and len(input_shape)==3: input_shape=tuple(input_shape) else: input_shape=(3, 224, 224) mob =MobileNet(input_shape=(3, 224, 224), classes=classes, use_bias=False, width_mult=1.0,round_nearest=8, include_top=include_top, model_name='mobilenet') if pretrained==True: download_model_from_google_drive('1ULenXTjOO5PdT3fHv6N8bPXEfoJAn5yL',dirname,'mobilenet_v2.pth') recovery_model=torch.load(os.path.join(dirname,'mobilenet_v2.pth')) recovery_model.eval() recovery_model.to(_device) if include_top==False: recovery_model.__delitem__(-1) else: if classes!=1000: new_fc = Dense(classes, activation=None, name='fc') new_fc.input_shape=recovery_model.fc.input_shape recovery_model.fc=new_fc mob.model=recovery_model return mob