File size: 1,070 Bytes
9bc4638 |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""Backbones from the TIMM library."""
from typing import List, Tuple
import torch
from timm.models import create_model
from torch import nn
class TimmBackbone(nn.Module):
def __init__(
self,
name: str,
features: Tuple[str, ...],
):
super().__init__()
out_indices = tuple(int(f[len("layer") :]) for f in features)
backbone = create_model(
name,
pretrained=True,
in_chans=3,
features_only=True,
out_indices=out_indices,
)
num_channels = backbone.feature_info.channels()
self.channel_list = num_channels[::-1]
self.body = backbone
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
xs = self.body(x)
out = []
for i, x in enumerate(xs):
out.append(x)
return out
|