File size: 2,770 Bytes
2d07fab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
import torch.nn as nn
from torchvision.ops.misc import FrozenBatchNorm2d
from torchvision.models import resnet, detection, segmentation
import timm
# https://detectron2.readthedocs.io/en/latest/modules/layers.html#detectron2.layers.FrozenBatchNorm2d.convert_frozen_batchnorm
@torch.no_grad()
def convert_frozen_batchnorm(module):
bn_module = (
nn.modules.batchnorm.BatchNorm2d,
nn.modules.batchnorm.SyncBatchNorm
)
res = module
if isinstance(module, bn_module):
res = FrozenBatchNorm2d(module.num_features)
if module.affine:
res.weight.data = module.weight.data.clone().detach()
res.bias.data = module.bias.data.clone().detach()
res.running_mean.data = module.running_mean.data
res.running_var.data = module.running_var.data
res.eps = module.eps
else:
for name, child in module.named_children():
new_child = convert_frozen_batchnorm(child)
if new_child is not child:
res.add_module(name, new_child)
return res
def get_backbone(backbone, pretrained=True):
if backbone in ('resnet18', 'resnet34', 'resnet50', 'resnet101'):
# pretrained on ImageNet for classification
model = resnet.__dict__[backbone](
pretrained=pretrained, norm_layer=FrozenBatchNorm2d
)
elif backbone == 'resnet50d':
# pretrained on COCO for detection
model = convert_frozen_batchnorm(
detection.fasterrcnn_resnet50_fpn(pretrained=pretrained).backbone.body
)
elif backbone == 'resnet50s':
# pretrained on COCO for segmentation
model = convert_frozen_batchnorm(
segmentation.deeplabv3_resnet50(pretrained=pretrained).backbone
)
elif backbone == 'resnet101s':
# pretrained on COCO for segmentation
model = convert_frozen_batchnorm(
segmentation.deeplabv3_resnet101(pretrained=pretrained).backbone
)
elif backbone in ('cspdarknet53', 'efficientnet-b0', 'efficientnet-b3'):
# model = convert_frozen_batchnorm(
# timm.create_model(
# backbone.replace('-', '_'),
# pretrained=True,
# features_only=True,
# #out_indices=(1, 2, 3, 4)
# )
# )
model = convert_frozen_batchnorm(
timm.create_model(
backbone.replace('-', '_'),
pretrained=pretrained,
num_classes=0,
global_pool=''
)
)
else:
raise RuntimeError(f'{backbone} is not a valid backbone')
# empty cache (dealloc modules other than the backbone)
torch.cuda.empty_cache()
return model
|