Load standard model, like vgg16, resnet18, Args: model_name: e.g., vgg16, inception, resnet18, ... ctg_num: e.g., 1000 use_sync_bn: True/False
(model_name, ctg_num, use_sync_bn)
| 48 | |
| 49 | |
| 50 | def load_model(model_name, ctg_num, use_sync_bn): |
| 51 | """Load standard model, like vgg16, resnet18, |
| 52 | |
| 53 | Args: |
| 54 | model_name: e.g., vgg16, inception, resnet18, ... |
| 55 | ctg_num: e.g., 1000 |
| 56 | use_sync_bn: True/False |
| 57 | """ |
| 58 | if model_name == 'convnext': |
| 59 | model = convnext_base(num_classes=ctg_num) |
| 60 | model_path = 'pre_model/convnext_base_1k_384.pth' |
| 61 | check_point = torch.load(model_path, map_location='cpu')['model'] |
| 62 | check_point.pop('head.weight') |
| 63 | check_point.pop('head.bias') |
| 64 | model.load_state_dict(check_point, strict=False) |
| 65 | |
| 66 | model = augment_inputs_network(model) |
| 67 | |
| 68 | elif model_name == 'replknet': |
| 69 | model = create_RepLKNet31B(num_classes=ctg_num, use_sync_bn=use_sync_bn) |
| 70 | model_path = 'pre_model/RepLKNet-31B_ImageNet-1K_384.pth' |
| 71 | check_point = torch.load(model_path) |
| 72 | check_point.pop('head.weight') |
| 73 | check_point.pop('head.bias') |
| 74 | model.load_state_dict(check_point, strict=False) |
| 75 | |
| 76 | model = augment_inputs_network(model) |
| 77 | |
| 78 | elif model_name == 'all': |
| 79 | model = final_model() |
| 80 | |
| 81 | print("model_name", model_name) |
| 82 | |
| 83 | return model |
| 84 |
no test coverage detected