trident
Table Of Contents
trident
Table Of Contents

Source code for trident.models.pytorch_arcfacenet

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 trident.backend.common import *
from trident.backend.pytorch_backend import to_numpy, to_tensor, Layer, Sequential,  summary
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, Relu, PRelu
from trident.layers.pytorch_blocks import *
from trident.layers.pytorch_layers import *
from trident.layers.pytorch_normalizations import get_normalization, BatchNorm2d, BatchNorm
from trident.layers.pytorch_pooling import *
from trident.optims.pytorch_trainer import *

__all__ = ['SEResNet_IR','BottleNeck_IR_SE','BottleNeck_IR']

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

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






[docs]def BottleNeck_IR(num_filters, strides,keep_filter=True): blocks = OrderedDict() blocks['res_layer'] = Sequential(BatchNorm2d(), Conv2d_Block((3, 3), num_filters=num_filters, strides=1, auto_pad=True, use_bias=False, activation=PRelu(num_filters)), Conv2d_Block((3, 3), num_filters, strides=strides, use_bias=False, activation=None, normalization='batch')) if keep_filter: blocks['shortcut_layer']=MaxPool2d(1, strides=strides,name='shortcut_layer') else: blocks['shortcut_layer'] = Conv2d_Block((1, 1), num_filters, strides=strides, use_bias=False, activation=None, normalization='batch') return ShortCut2d(blocks,mode='add')
[docs]def BottleNeck_IR_SE( num_filters, strides,keep_filter=True): blocks = OrderedDict() blocks['res_layer'] = Sequential(BatchNorm2d(), Conv2d_Block((3, 3), num_filters=num_filters, strides=1, auto_pad=True, use_bias=False, activation=PRelu(num_filters)), Conv2d_Block((3, 3), num_filters, strides=strides, use_bias=False, activation=None, normalization='batch'), SqueezeExcite(num_filters//16,num_filters),name='res_layer') if keep_filter: blocks['shortcut_layer'] = MaxPool2d(1, strides=strides, name='shortcut_layer') else: blocks['shortcut_layer'] =Conv2d_Block((1, 1), num_filters, strides=strides, use_bias=False, activation=None, normalization='batch',name='shortcut_layer') return ShortCut2d(blocks,mode='add')
def get_block(Bottleneck, out_channel, num_units, strides=2,keep_filter=True): blocks=[Bottleneck(out_channel, strides,keep_filter)] for i in range(num_units - 1): blocks.append(Bottleneck(out_channel, 1,True)) return blocks
[docs]def SEResNet_IR(num_layers=50,Bottleneck=BottleNeck_IR_SE,drop_ratio=0.4,feature_dim=128): blocks=OrderedDict() blocks['input_layer']=Conv2d_Block((3,3),64,strides=1,auto_pad=True,use_bias=False,activation=PRelu(64),normalization='batch',name='input_layer') blocks['body']=Sequential( get_block(Bottleneck, out_channel=64, num_units=3,keep_filter=True)+ get_block(Bottleneck, out_channel=128, num_units=4,keep_filter=False)+ get_block(Bottleneck, out_channel=256, num_units=14,keep_filter=False)+ get_block(Bottleneck, out_channel=512, num_units=3,keep_filter=False) ) blocks['output_layer']=Sequential( BatchNorm2d(), Dropout(drop_ratio), Flatten(), Dense(feature_dim), BatchNorm(), name='output_layer' ) return Sequential(blocks).to(_device)