Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,325 Bytes
e85fecb |
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 |
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
import torchvision
from ...core import register
from .utils import IntermediateLayerGetter
__all__ = ["TorchVisionModel"]
@register()
class TorchVisionModel(torch.nn.Module):
def __init__(self, name, return_layers, weights=None, **kwargs) -> None:
super().__init__()
if weights is not None:
weights = getattr(torchvision.models.get_model_weights(name), weights)
model = torchvision.models.get_model(name, weights=weights, **kwargs)
# TODO hard code.
if hasattr(model, "features"):
model = IntermediateLayerGetter(model.features, return_layers)
else:
model = IntermediateLayerGetter(model, return_layers)
self.model = model
def forward(self, x):
return self.model(x)
# TorchVisionModel('swin_t', return_layers=['5', '7'])
# TorchVisionModel('resnet34', return_layers=['layer2','layer3', 'layer4'])
# TorchVisionModel:
# name: swin_t
# return_layers: ['5', '7']
# weights: DEFAULT
# model:
# type: TorchVisionModel
# name: resnet34
# return_layers: ['layer2','layer3', 'layer4']
# weights: DEFAULT
|