Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| from timm.models import efficientnet, convnext | |
| def build_backbone(model_name, pretrained): | |
| model = getattr(Backbones, model_name)(pretrained=pretrained) | |
| return model | |
| class Backbones(object): | |
| def efficientnet_b3_p(pretrained): | |
| # channels: 24, 12, 40, 120, 384 | |
| # for test, pretrained can be set to False | |
| model = efficientnet.efficientnet_b3_pruned(pretrained=pretrained, features_only=True) | |
| ''' | |
| # pre-downloaded weights | |
| cp_path = os.path.join('checkpoints', 'effnetb3_pruned-59ecf72d.pth') | |
| state_dict = torch.load(cp_path, map_location=torch.device('cpu')) | |
| model.load_state_dict(state_dict=state_dict, strict=False)''' | |
| return model | |