fastreid.modeling.backbones.osnet Namespace Reference

Classes
class	ChannelGate
class	Conv1x1
class	Conv1x1Linear
class	Conv3x3
class	ConvLayer
	Basic layers. More...
class	LightConv3x3
class	OSBlock
class	OSNet

Functions
	init_pretrained_weights (model, key='')
	build_osnet_backbone (cfg)

Variables
	logger = logging.getLogger(__name__)
dict	model_urls

Detailed Description

@author:  xingyu liao
@contact: sherlockliao01@gmail.com

Function Documentation

build_osnet_backbone()

fastreid.modeling.backbones.osnet.build_osnet_backbone

(

cfg

)

Create a OSNet instance from config.
Returns:
    OSNet: a :class:`OSNet` instance

Definition at line 486 of file osnet.py.

def build_osnet_backbone(cfg):
    """
    Create a OSNet instance from config.
    Returns:
        OSNet: a :class:`OSNet` instance
    """
 
    # fmt: off
    pretrain      = cfg.MODEL.BACKBONE.PRETRAIN
    pretrain_path = cfg.MODEL.BACKBONE.PRETRAIN_PATH
    with_ibn      = cfg.MODEL.BACKBONE.WITH_IBN
    bn_norm       = cfg.MODEL.BACKBONE.NORM
    depth         = cfg.MODEL.BACKBONE.DEPTH
    # fmt: on
 
    num_blocks_per_stage = [2, 2, 2]
    num_channels_per_stage = {
        "x1_0": [64, 256, 384, 512],
        "x0_75": [48, 192, 288, 384],
        "x0_5": [32, 128, 192, 256],
        "x0_25": [16, 64, 96, 128]}[depth]
    model = OSNet([OSBlock, OSBlock, OSBlock], num_blocks_per_stage, num_channels_per_stage,
                  bn_norm, IN=with_ibn)
 
    if pretrain:
        # Load pretrain path if specifically
        if pretrain_path:
            try:
                state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))
                logger.info(f"Loading pretrained model from {pretrain_path}")
                model.load_state_dict(state_dict)
            except FileNotFoundError as e:
                logger.info(f'{pretrain_path} is not found! Please check this path.')
                raise e
            except KeyError as e:
                logger.info("State dict keys error! Please check the state dict.")
                raise e
        else:
            if with_ibn:
                pretrain_key = "osnet_ibn_" + depth
            else:
                pretrain_key = "osnet_" + depth
 
            init_pretrained_weights(model, pretrain_key)
    return model

init_pretrained_weights()

fastreid.modeling.backbones.osnet.init_pretrained_weights	(		model,
			key = '' )

Initializes model with pretrained weights.

Layers that don't match with pretrained layers in name or size are kept unchanged.

Definition at line 402 of file osnet.py.

def init_pretrained_weights(model, key=''):
    """Initializes model with pretrained weights.
 
    Layers that don't match with pretrained layers in name or size are kept unchanged.
    """
    import os
    import errno
    import gdown
    from collections import OrderedDict
    import warnings
    import logging
 
    logger = logging.getLogger(__name__)
 
    def _get_torch_home():
        ENV_TORCH_HOME = 'TORCH_HOME'
        ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
        DEFAULT_CACHE_DIR = '~/.cache'
        torch_home = os.path.expanduser(
            os.getenv(
                ENV_TORCH_HOME,
                os.path.join(
                    os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
                )
            )
        )
        return torch_home
 
    torch_home = _get_torch_home()
    model_dir = os.path.join(torch_home, 'checkpoints')
    try:
        os.makedirs(model_dir)
    except OSError as e:
        if e.errno == errno.EEXIST:
            # Directory already exists, ignore.
            pass
        else:
            # Unexpected OSError, re-raise.
            raise
    filename = key + '_imagenet.pth'
    cached_file = os.path.join(model_dir, filename)
 
    if not os.path.exists(cached_file):
        if comm.is_main_process():
            gdown.download(model_urls[key], cached_file, quiet=False)
 
    comm.synchronize()
 
    state_dict = torch.load(cached_file, map_location=torch.device('cpu'))
    model_dict = model.state_dict()
    new_state_dict = OrderedDict()
    matched_layers, discarded_layers = [], []
 
    for k, v in state_dict.items():
        if k.startswith('module.'):
            k = k[7:]  # discard module.
 
        if k in model_dict and model_dict[k].size() == v.size():
            new_state_dict[k] = v
            matched_layers.append(k)
        else:
            discarded_layers.append(k)
 
    model_dict.update(new_state_dict)
    model.load_state_dict(model_dict)
 
    if len(matched_layers) == 0:
        warnings.warn(
            'The pretrained weights from "{}" cannot be loaded, '
            'please check the key names manually '
            '(** ignored and continue **)'.format(cached_file)
        )
    else:
        logger.info(
            'Successfully loaded imagenet pretrained weights from "{}"'.format(cached_file)
        )
        if len(discarded_layers) > 0:
            logger.info(
                '** The following layers are discarded '
                'due to unmatched keys or layer size: {}'.format(discarded_layers)
            )
 
 
@BACKBONE_REGISTRY.register()

Variable Documentation

logger

fastreid.modeling.backbones.osnet.logger = logging.getLogger(__name__)

Definition at line 19 of file osnet.py.

model_urls

dict fastreid.modeling.backbones.osnet.model_urls

Initial value:

=  {
    'osnet_x1_0':
        'https://drive.google.com/uc?id=1LaG1EJpHrxdAxKnSCJ_i0u-nbxSAeiFY',
    'osnet_x0_75':
        'https://drive.google.com/uc?id=1uwA9fElHOk3ZogwbeY5GkLI6QPTX70Hq',
    'osnet_x0_5':
        'https://drive.google.com/uc?id=16DGLbZukvVYgINws8u8deSaOqjybZ83i',
    'osnet_x0_25':
        'https://drive.google.com/uc?id=1rb8UN5ZzPKRc_xvtHlyDh-cSz88YX9hs',
    'osnet_ibn_x1_0':
        'https://drive.google.com/uc?id=1sr90V6irlYYDd4_4ISU2iruoRG8J__6l'
}

Definition at line 20 of file osnet.py.