| | import torch |
| | from mmengine.model import BaseModule |
| | from torch import nn |
| |
|
| | from mmseg.registry import MODELS |
| | import os |
| | _DINOV2_MMSEG_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) |
| | _DINOV2_TORCHHUB_DIR = os.path.join(_DINOV2_MMSEG_ROOT, 'torchhub', 'facebookresearch_dinov2_main') |
| |
|
| |
|
| | @MODELS.register_module() |
| | class DINOv2(nn.Module): |
| | """Use DINOv2 pre-trained models |
| | """ |
| |
|
| | def __init__(self, version='large', freeze=False, load_from=None): |
| | super().__init__() |
| | |
| | if version == 'large': |
| | self.dinov2 = torch.hub.load(_DINOV2_TORCHHUB_DIR, 'dinov2_vitl14', source='local', pretrained=False) |
| | else: |
| | raise NotImplementedError |
| |
|
| | if load_from is not None: |
| | if load_from.split('/')[-1] == 'depth_anything_vitl14.pth': |
| | print(load_from) |
| | d = torch.load(load_from, map_location='cpu') |
| | new_d = {} |
| | for key, value in d.items(): |
| | if 'pretrained' in key: |
| | new_d[key.replace('pretrained.', '')] = value |
| | self.dinov2.load_state_dict(new_d) |
| | else: |
| | print(load_from) |
| | all_d = torch.load(load_from, map_location='cpu') |
| | d = all_d['state_dict'] |
| | new_d = {} |
| | for key, value in d.items(): |
| | if 'backbone.dinov2' in key: |
| | new_d[key.replace('backbone.dinov2.', '')] = value |
| | self.dinov2.load_state_dict(new_d) |
| | |
| | self.freeze = freeze |
| | |
| | def forward(self, inputs): |
| | B, _, h, w = inputs.shape |
| | |
| | if self.freeze: |
| | with torch.no_grad(): |
| | features = self.dinov2.get_intermediate_layers(inputs, 4) |
| | else: |
| | features = self.dinov2.get_intermediate_layers(inputs, 4) |
| | |
| | outs = [] |
| | for feature in features: |
| | C = feature.shape[-1] |
| | feature = feature.permute(0, 2, 1).reshape(B, C, h // 14, w // 14).contiguous() |
| | outs.append(feature) |
| | |
| | return outs |
| |
|
