Python Sequential.load_state_dict Examples

Example #1

File: build.py Project: alibaba-edu/Ranking-based-Instance-Selection

def build_model(cfg):
    backbone = build_backbone(cfg)
    head = build_head(cfg)

    model = Sequential(OrderedDict([("backbone", backbone), ("head", head)]))

    if cfg.MODEL.PRETRAIN == "imagenet":
        backbone_name = cfg.MODEL.BACKBONE.NAME
        if backbone_name.startswith('resnet50_replaceFIRST'):
            backbone_name = 'resnet50'
        pretrained_path = os.path.expanduser(
            cfg.MODEL.PRETRIANED_PATH[backbone_name])
        print(
            f">>> Loading imagenet pretrianed model from {pretrained_path}...")
        model.backbone.load_param(pretrained_path)
    elif os.path.exists(cfg.MODEL.PRETRAIN):
        print(f">>> Loading cfg model from {cfg.MODEL.PRETRAIN}...")
        ckp = torch.load(cfg.MODEL.PRETRAIN)
        model.load_state_dict(ckp["model"])
    elif cfg.MODEL.PRETRAIN == "scratch":
        print(">>> train from scratch")
    else:
        assert (
            False
        ), f"not imagenet pretrain and not exist model path: {cfg.MODEL.PRETRAIN}"
    return model

Example #2

File: build.py Project: wayne980/research-ms-loss

def build_model(cfg):
    backbone = build_backbone(cfg)
    head = build_head(cfg)

    model = Sequential(OrderedDict([('backbone', backbone), ('head', head)]))

    if cfg.MODEL.PRETRAIN == 'imagenet':
        print('Loading imagenet pretrianed model ...')
        pretrained_path = os.path.expanduser(
            cfg.MODEL.PRETRIANED_PATH[cfg.MODEL.BACKBONE.NAME])
        model.backbone.load_param(pretrained_path)
    elif os.path.exists(cfg.MODEL.PRETRAIN):
        ckp = torch.load(cfg.MODEL.PRETRAIN)
        model.load_state_dict(ckp['model_state_dict'])
    return model