JUGGHM commited on Apr 24, 2024

Commit

0b1902c

verified ·

1 Parent(s): 36b8fb6

Delete mono

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

mono/configs/HourglassDecoder/convlarge.0.3_150.py +0 -25
mono/configs/HourglassDecoder/test_kitti_convlarge.0.3_150.py +0 -25
mono/configs/HourglassDecoder/test_nyu_convlarge.0.3_150.py +0 -25
mono/configs/HourglassDecoder/vit.raft5.large.py +0 -33
mono/configs/HourglassDecoder/vit.raft5.small.py +0 -33
mono/configs/__init__.py +0 -1
mono/configs/_base_/_data_base_.py +0 -13
mono/configs/_base_/datasets/_data_base_.py +0 -12
mono/configs/_base_/default_runtime.py +0 -4
mono/configs/_base_/models/backbones/convnext_large.py +0 -16
mono/configs/_base_/models/backbones/dino_vit_large.py +0 -7
mono/configs/_base_/models/backbones/dino_vit_large_reg.py +0 -7
mono/configs/_base_/models/backbones/dino_vit_small_reg.py +0 -7
mono/configs/_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py +0 -10
mono/configs/_base_/models/encoder_decoder/dino_vit_large.dpt_raft.py +0 -20
mono/configs/_base_/models/encoder_decoder/dino_vit_large_reg.dpt_raft.py +0 -19
mono/configs/_base_/models/encoder_decoder/dino_vit_small_reg.dpt_raft.py +0 -19
mono/model/__init__.py +0 -5
mono/model/__pycache__/__init__.cpython-39.pyc +0 -0
mono/model/__pycache__/monodepth_model.cpython-39.pyc +0 -0
mono/model/backbones/ConvNeXt.py +0 -271
mono/model/backbones/ViT_DINO.py +0 -1504
mono/model/backbones/ViT_DINO_reg.py +0 -1293
mono/model/backbones/__init__.py +0 -11
mono/model/backbones/__pycache__/ConvNeXt.cpython-39.pyc +0 -0
mono/model/backbones/__pycache__/__init__.cpython-39.pyc +0 -0
mono/model/decode_heads/HourGlassDecoder.py +0 -274
mono/model/decode_heads/RAFTDepthNormalDPTDecoder5.py +0 -1033
mono/model/decode_heads/__init__.py +0 -4
mono/model/decode_heads/__pycache__/HourGlassDecoder.cpython-39.pyc +0 -0
mono/model/decode_heads/__pycache__/__init__.cpython-39.pyc +0 -0
mono/model/model_pipelines/__base_model__.py +0 -20
mono/model/model_pipelines/__init__.py +0 -6
mono/model/model_pipelines/__pycache__/__base_model__.cpython-39.pyc +0 -0
mono/model/model_pipelines/__pycache__/__init__.cpython-39.pyc +0 -0
mono/model/model_pipelines/__pycache__/dense_pipeline.cpython-39.pyc +0 -0
mono/model/model_pipelines/dense_pipeline.py +0 -16
mono/model/monodepth_model.py +0 -37
mono/tools/test_scale_cano.py +0 -158
mono/utils/__init__.py +0 -1
mono/utils/__pycache__/__init__.cpython-39.pyc +0 -0
mono/utils/__pycache__/avg_meter.cpython-39.pyc +0 -0
mono/utils/__pycache__/comm.cpython-39.pyc +0 -0
mono/utils/__pycache__/custom_data.cpython-39.pyc +0 -0
mono/utils/__pycache__/do_test.cpython-39.pyc +0 -0
mono/utils/__pycache__/logger.cpython-39.pyc +0 -0
mono/utils/__pycache__/mldb.cpython-39.pyc +0 -0
mono/utils/__pycache__/running.cpython-39.pyc +0 -0
mono/utils/__pycache__/transform.cpython-39.pyc +0 -0
mono/utils/__pycache__/unproj_pcd.cpython-39.pyc +0 -0

mono/configs/HourglassDecoder/convlarge.0.3_150.py DELETED Viewed

@@ -1,25 +0,0 @@
-_base_=[
-       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
-       '../_base_/datasets/_data_base_.py',
-       '../_base_/default_runtime.py',
-       ]
-model = dict(
-    backbone=dict(
-        pretrained=False,
-    )
-)
-# configs of the canonical space
-data_basic=dict(
-    canonical_space = dict(
-        img_size=(512, 960),
-        focal_length=1000.0,
-    ),
-    depth_range=(0, 1),
-    depth_normalize=(0.3, 150),
-    crop_size = (544, 1216),
-)
-batchsize_per_gpu = 2
-thread_per_gpu = 4

mono/configs/HourglassDecoder/test_kitti_convlarge.0.3_150.py DELETED Viewed

@@ -1,25 +0,0 @@
-_base_=[
-       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
-       '../_base_/datasets/_data_base_.py',
-       '../_base_/default_runtime.py',
-       ]
-model = dict(
-    backbone=dict(
-        pretrained=False,
-    )
-)
-# configs of the canonical space
-data_basic=dict(
-    canonical_space = dict(
-        img_size=(512, 960),
-        focal_length=1000.0,
-    ),
-    depth_range=(0, 1),
-    depth_normalize=(0.3, 150),
-    crop_size = (512, 1088),
-)
-batchsize_per_gpu = 2
-thread_per_gpu = 4

mono/configs/HourglassDecoder/test_nyu_convlarge.0.3_150.py DELETED Viewed

@@ -1,25 +0,0 @@
-_base_=[
-       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
-       '../_base_/datasets/_data_base_.py',
-       '../_base_/default_runtime.py',
-       ]
-model = dict(
-    backbone=dict(
-        pretrained=False,
-    )
-)
-# configs of the canonical space
-data_basic=dict(
-    canonical_space = dict(
-        img_size=(512, 960),
-        focal_length=1000.0,
-    ),
-    depth_range=(0, 1),
-    depth_normalize=(0.3, 150),
-    crop_size = (480, 1216),
-)
-batchsize_per_gpu = 2
-thread_per_gpu = 4

mono/configs/HourglassDecoder/vit.raft5.large.py DELETED Viewed

@@ -1,33 +0,0 @@
-_base_=[
-       '../_base_/models/encoder_decoder/dino_vit_large_reg.dpt_raft.py',
-       '../_base_/datasets/_data_base_.py',
-       '../_base_/default_runtime.py',
-       ]
-import numpy as np
-model=dict(
-    decode_head=dict(
-        type='RAFTDepthNormalDPT5',
-        iters=8,
-        n_downsample=2,
-        detach=False,
-    )
-)
-max_value = 200
-# configs of the canonical space
-data_basic=dict(
-    canonical_space = dict(
-        # img_size=(540, 960),
-        focal_length=1000.0,
-    ),
-    depth_range=(0, 1),
-    depth_normalize=(0.1, max_value),
-    crop_size = (616, 1064),  # %28 = 0
-     clip_depth_range=(0.1, 200),
-    vit_size=(616,1064)
-)
-batchsize_per_gpu = 1
-thread_per_gpu = 1

mono/configs/HourglassDecoder/vit.raft5.small.py DELETED Viewed

@@ -1,33 +0,0 @@
-_base_=[
-       '../_base_/models/encoder_decoder/dino_vit_small_reg.dpt_raft.py',
-       '../_base_/datasets/_data_base_.py',
-       '../_base_/default_runtime.py',
-       ]
-import numpy as np
-model=dict(
-    decode_head=dict(
-        type='RAFTDepthNormalDPT5',
-        iters=4,
-        n_downsample=2,
-        detach=False,
-    )
-)
-max_value = 200
-# configs of the canonical space
-data_basic=dict(
-    canonical_space = dict(
-        # img_size=(540, 960),
-        focal_length=1000.0,
-    ),
-    depth_range=(0, 1),
-    depth_normalize=(0.1, max_value),
-    crop_size = (616, 1064),  # %28 = 0
-     clip_depth_range=(0.1, 200),
-    vit_size=(616,1064)
-)
-batchsize_per_gpu = 1
-thread_per_gpu = 1

mono/configs/__init__.py DELETED Viewed

	@@ -1 +0,0 @@
1	-

mono/configs/_base_/_data_base_.py DELETED Viewed

@@ -1,13 +0,0 @@
-# canonical camera setting and basic data setting
-# we set it same as the E300 camera (crop version)
-#
-data_basic=dict(
-    canonical_space = dict(
-        img_size=(540, 960),
-        focal_length=1196.0,
-    ),
-    depth_range=(0.9, 150),
-    depth_normalize=(0.006, 1.001),
-    crop_size = (512, 960),
-    clip_depth_range=(0.9, 150),
-)

mono/configs/_base_/datasets/_data_base_.py DELETED Viewed

@@ -1,12 +0,0 @@
-# canonical camera setting and basic data setting
-#
-data_basic=dict(
-    canonical_space = dict(
-        img_size=(540, 960),
-        focal_length=1196.0,
-    ),
-    depth_range=(0.9, 150),
-    depth_normalize=(0.006, 1.001),
-    crop_size = (512, 960),
-    clip_depth_range=(0.9, 150),
-)

mono/configs/_base_/default_runtime.py DELETED Viewed

@@ -1,4 +0,0 @@
-load_from = None
-cudnn_benchmark = True
-test_metrics = ['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3','rmse_log', 'log10', 'sq_rel']

mono/configs/_base_/models/backbones/convnext_large.py DELETED Viewed

@@ -1,16 +0,0 @@
-#_base_ = ['./_model_base_.py',]
-#'https://download.openmmlab.com/mmclassification/v0/convnext/downstream/convnext-large_3rdparty_in21k_20220301-e6e0ea0a.pth'
-model = dict(
-    #type='EncoderDecoderAuxi',
-    backbone=dict(
-        type='convnext_large',
-        pretrained=True,
-        in_22k=True,
-        out_indices=[0, 1, 2, 3],
-        drop_path_rate=0.4,
-        layer_scale_init_value=1.0,
-        checkpoint='data/pretrained_weight_repo/convnext/convnext_large_22k_1k_384.pth',
-        prefix='backbones.',
-        out_channels=[192, 384, 768, 1536]),
-    )

mono/configs/_base_/models/backbones/dino_vit_large.py DELETED Viewed

@@ -1,7 +0,0 @@
-model = dict(
-    backbone=dict(
-        type='vit_large',
-        prefix='backbones.',
-        out_channels=[1024, 1024, 1024, 1024],
-        drop_path_rate = 0.0),
-    )

mono/configs/_base_/models/backbones/dino_vit_large_reg.py DELETED Viewed

@@ -1,7 +0,0 @@
-model = dict(
-    backbone=dict(
-        type='vit_large_reg',
-        prefix='backbones.',
-        out_channels=[1024, 1024, 1024, 1024],
-        drop_path_rate = 0.0),
-    )

mono/configs/_base_/models/backbones/dino_vit_small_reg.py DELETED Viewed

@@ -1,7 +0,0 @@
-model = dict(
-    backbone=dict(
-        type='vit_small_reg',
-        prefix='backbones.',
-        out_channels=[384, 384, 384, 384],
-        drop_path_rate = 0.0),
-    )

mono/configs/_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py DELETED Viewed

@@ -1,10 +0,0 @@
-# model settings
-_base_ = ['../backbones/convnext_large.py',]
-model = dict(
-    type='DensePredModel',
-    decode_head=dict(
-        type='HourglassDecoder',
-        in_channels=[192, 384, 768, 1536],
-        decoder_channel=[128, 128, 256, 512],
-        prefix='decode_heads.'),
-)

mono/configs/_base_/models/encoder_decoder/dino_vit_large.dpt_raft.py DELETED Viewed

@@ -1,20 +0,0 @@
-# model settings
-_base_ = ['../backbones/dino_vit_large.py']
-model = dict(
-    type='DensePredModel',
-    decode_head=dict(
-        type='RAFTDepthDPT',
-        in_channels=[1024, 1024, 1024, 1024],
-        use_cls_token=True,
-        feature_channels = [256, 512, 1024, 1024], # [2/7, 1/7, 1/14, 1/14]
-        decoder_channels = [128, 256, 512, 1024, 1024], # [4/7, 2/7, 1/7, 1/14, 1/14]
-        up_scale = 7,
-        hidden_channels=[128, 128, 128, 128], # [x_4, x_8, x_16, x_32] [192, 384, 768, 1536]
-        n_gru_layers=3,
-        n_downsample=2,
-        iters=12,
-        slow_fast_gru=True,
-        corr_radius=4,
-        corr_levels=4,
-        prefix='decode_heads.'),
-)

mono/configs/_base_/models/encoder_decoder/dino_vit_large_reg.dpt_raft.py DELETED Viewed

@@ -1,19 +0,0 @@
-# model settings
-_base_ = ['../backbones/dino_vit_large_reg.py']
-model = dict(
-    type='DensePredModel',
-    decode_head=dict(
-        type='RAFTDepthDPT',
-        in_channels=[1024, 1024, 1024, 1024],
-        use_cls_token=True,
-        feature_channels = [256, 512, 1024, 1024], # [2/7, 1/7, 1/14, 1/14]
-        decoder_channels = [128, 256, 512, 1024, 1024], # [4/7, 2/7, 1/7, 1/14, 1/14]
-        up_scale = 7,
-        hidden_channels=[128, 128, 128, 128], # [x_4, x_8, x_16, x_32] [192, 384, 768, 1536]
-        n_gru_layers=3,
-        n_downsample=2,
-        iters=3,
-        slow_fast_gru=True,
-        num_register_tokens=4,
-        prefix='decode_heads.'),
-)

mono/configs/_base_/models/encoder_decoder/dino_vit_small_reg.dpt_raft.py DELETED Viewed

@@ -1,19 +0,0 @@
-# model settings
-_base_ = ['../backbones/dino_vit_small_reg.py']
-model = dict(
-    type='DensePredModel',
-    decode_head=dict(
-        type='RAFTDepthDPT',
-        in_channels=[384, 384, 384, 384],
-        use_cls_token=True,
-        feature_channels = [96, 192, 384, 768], # [2/7, 1/7, 1/14, 1/14]
-        decoder_channels = [48, 96, 192, 384, 384], # [-, 1/4, 1/7, 1/14, 1/14]
-        up_scale = 7,
-        hidden_channels=[48, 48, 48, 48], # [x_4, x_8, x_16, x_32] [1/4, 1/7, 1/14, -]
-        n_gru_layers=3,
-        n_downsample=2,
-        iters=3,
-        slow_fast_gru=True,
-        num_register_tokens=4,
-        prefix='decode_heads.'),
-)

mono/model/__init__.py DELETED Viewed

@@ -1,5 +0,0 @@
-from .monodepth_model import DepthModel
-# from .__base_model__ import BaseDepthModel
-__all__ = ['DepthModel', 'BaseDepthModel']

mono/model/__pycache__/__init__.cpython-39.pyc DELETED Viewed

Binary file (250 Bytes)

mono/model/__pycache__/monodepth_model.cpython-39.pyc DELETED Viewed

Binary file (1.62 kB)

mono/model/backbones/ConvNeXt.py DELETED Viewed

@@ -1,271 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from timm.models.layers import trunc_normal_, DropPath
-from timm.models.registry import register_model
-class Block(nn.Module):
-    r""" ConvNeXt Block. There are two equivalent implementations:
-    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
-    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
-    We use (2) as we find it slightly faster in PyTorch
-    Args:
-        dim (int): Number of input channels.
-        drop_path (float): Stochastic depth rate. Default: 0.0
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-    """
-    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
-        super().__init__()
-        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
-        self.norm = LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Linear(4 * dim, dim)
-        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
-                                    requires_grad=True) if layer_scale_init_value > 0 else None
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-    def forward(self, x):
-        input = x
-        x = self.dwconv(x)
-        x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-        if self.gamma is not None:
-            x = self.gamma * x
-        x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
-        x = input + self.drop_path(x)
-        return x
-class ConvNeXt(nn.Module):
-    r""" ConvNeXt
-        A PyTorch impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
-    Args:
-        in_chans (int): Number of input image channels. Default: 3
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-    def __init__(self, in_chans=3, num_classes=1000,
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0.,
-                 layer_scale_init_value=1e-6, head_init_scale=1.,
-                 **kwargs,):
-        super().__init__()
-        self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
-        stem = nn.Sequential(
-            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
-            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
-        )
-        self.downsample_layers.append(stem)
-        for i in range(3):
-            downsample_layer = nn.Sequential(
-                    LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
-                    nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
-            )
-            self.downsample_layers.append(downsample_layer)
-        self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(4):
-            stage = nn.Sequential(
-                *[Block(dim=dims[i], drop_path=dp_rates[cur + j],
-                layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-        #self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
-        #self.head = nn.Linear(dims[-1], num_classes)
-        self.apply(self._init_weights)
-        #self.head.weight.data.mul_(head_init_scale)
-        #self.head.bias.data.mul_(head_init_scale)
-    def _init_weights(self, m):
-        if isinstance(m, (nn.Conv2d, nn.Linear)):
-            trunc_normal_(m.weight, std=.02)
-            nn.init.constant_(m.bias, 0)
-    def forward_features(self, x):
-        features = []
-        for i in range(4):
-            x = self.downsample_layers[i](x)
-            x = self.stages[i](x)
-            features.append(x)
-        return features # global average pooling, (N, C, H, W) -> (N, C)
-    def forward(self, x):
-        #x = self.forward_features(x)
-        #x = self.head(x)
-        features = self.forward_features(x)
-        return features
-class LayerNorm(nn.Module):
-    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
-    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
-    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
-    with shape (batch_size, channels, height, width).
-    """
-    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(normalized_shape))
-        self.bias = nn.Parameter(torch.zeros(normalized_shape))
-        self.eps = eps
-        self.data_format = data_format
-        if self.data_format not in ["channels_last", "channels_first"]:
-            raise NotImplementedError
-        self.normalized_shape = (normalized_shape, )
-    def forward(self, x):
-        if self.data_format == "channels_last":
-            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
-        elif self.data_format == "channels_first":
-            u = x.mean(1, keepdim=True)
-            s = (x - u).pow(2).mean(1, keepdim=True)
-            x = (x - u) / torch.sqrt(s + self.eps)
-            x = self.weight[:, None, None] * x + self.bias[:, None, None]
-            return x
-model_urls = {
-    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
-    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
-    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
-    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
-    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
-    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
-    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
-    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
-    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
-}
-def convnext_tiny(pretrained=True,in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
-    if pretrained:
-        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
-        #url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
-        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
-        model_dict = model.state_dict()
-        pretrained_dict = {}
-        unmatched_pretrained_dict = {}
-        for k, v in checkpoint['model'].items():
-            if k in model_dict:
-                pretrained_dict[k] = v
-            else:
-                unmatched_pretrained_dict[k] = v
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-        print(
-            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
-            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
-        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
-    return model
-def convnext_small(pretrained=True,in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
-    if pretrained:
-        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
-        #url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
-        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
-        model_dict = model.state_dict()
-        pretrained_dict = {}
-        unmatched_pretrained_dict = {}
-        for k, v in checkpoint['model'].items():
-            if k in model_dict:
-                pretrained_dict[k] = v
-            else:
-                unmatched_pretrained_dict[k] = v
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-        print(
-            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
-            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
-        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
-    return model
-def convnext_base(pretrained=True, in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
-    if pretrained:
-        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
-        #url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
-        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
-        model_dict = model.state_dict()
-        pretrained_dict = {}
-        unmatched_pretrained_dict = {}
-        for k, v in checkpoint['model'].items():
-            if k in model_dict:
-                pretrained_dict[k] = v
-            else:
-                unmatched_pretrained_dict[k] = v
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-        print(
-            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
-            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
-        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
-    return model
-def convnext_large(pretrained=True, in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
-    if pretrained:
-        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
-        #url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
-        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
-        model_dict = model.state_dict()
-        pretrained_dict = {}
-        unmatched_pretrained_dict = {}
-        for k, v in checkpoint['model'].items():
-            if k in model_dict:
-                pretrained_dict[k] = v
-            else:
-                unmatched_pretrained_dict[k] = v
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-        print(
-            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
-            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
-        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
-    return model
-def convnext_xlarge(pretrained=True, in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
-    if pretrained:
-        assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
-        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
-        #url = model_urls['convnext_xlarge_22k']
-        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
-        model_dict = model.state_dict()
-        pretrained_dict = {}
-        unmatched_pretrained_dict = {}
-        for k, v in checkpoint['model'].items():
-            if k in model_dict:
-                pretrained_dict[k] = v
-            else:
-                unmatched_pretrained_dict[k] = v
-        model_dict.update(pretrained_dict)
-        model.load_state_dict(model_dict)
-        print(
-            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
-            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
-        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
-    return model
-if __name__ == '__main__':
-    import torch
-    model = convnext_base(True, in_22k=False).cuda()
-    rgb = torch.rand((2, 3, 256, 256)).cuda()
-    out = model(rgb)
-    print(len(out))
-    for i, ft in enumerate(out):
-        print(i, ft.shape)

mono/model/backbones/ViT_DINO.py DELETED Viewed

@@ -1,1504 +0,0 @@
-# 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.
-# References:
-#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
-from functools import partial
-import math
-import logging
-from typing import Sequence, Tuple, Union, Callable, Optional, Dict, Any, List
-import torch
-import torch.nn as nn
-from torch import Tensor
-import torch.utils.checkpoint
-from torch.nn.init import trunc_normal_
-#from dinov2.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
-logger = logging.getLogger("dinov2")
-class ConvBlock(nn.Module):
-    def __init__(self, channels):
-        super(ConvBlock, self).__init__()
-        self.act = nn.ReLU(inplace=True)
-        self.conv1 = nn.Conv2d(
-            channels,
-            channels,
-            kernel_size=3,
-            stride=1,
-            padding=1
-        )
-        self.norm1 = nn.BatchNorm2d(channels)
-        self.conv2 = nn.Conv2d(
-            channels,
-            channels,
-            kernel_size=3,
-            stride=1,
-            padding=1
-        )
-        self.norm2 = nn.BatchNorm2d(channels)
-    def forward(self, x):
-        out = self.norm1(x)
-        out = self.act(out)
-        out = self.conv1(out)
-        out = self.norm2(out)
-        out = self.act(out)
-        out = self.conv2(out)
-        return x + out
-def make_2tuple(x):
-    if isinstance(x, tuple):
-        assert len(x) == 2
-        return x
-    assert isinstance(x, int)
-    return (x, x)
-def drop_path(x, drop_prob: float = 0.0, training: bool = False):
-    if drop_prob == 0.0 or not training:
-        return x
-    keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
-    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
-    if keep_prob > 0.0:
-        random_tensor.div_(keep_prob)
-    output = x * random_tensor
-    return output
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-    def __init__(self, drop_prob=None):
-        super(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training)
-class LayerScale(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        init_values: Union[float, Tensor] = 1e-5,
-        inplace: bool = False,
-    ) -> None:
-        super().__init__()
-        self.inplace = inplace
-        self.gamma = nn.Parameter(init_values * torch.ones(dim))
-    def forward(self, x: Tensor) -> Tensor:
-        return x.mul_(self.gamma) if self.inplace else x * self.gamma
-class PatchEmbed(nn.Module):
-    """
-    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
-    Args:
-        img_size: Image size.
-        patch_size: Patch token size.
-        in_chans: Number of input image channels.
-        embed_dim: Number of linear projection output channels.
-        norm_layer: Normalization layer.
-    """
-    def __init__(
-        self,
-        img_size: Union[int, Tuple[int, int]] = 224,
-        patch_size: Union[int, Tuple[int, int]] = 16,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-        norm_layer: Optional[Callable] = None,
-        flatten_embedding: bool = True,
-    ) -> None:
-        super().__init__()
-        image_HW = make_2tuple(img_size)
-        patch_HW = make_2tuple(patch_size)
-        patch_grid_size = (
-            image_HW[0] // patch_HW[0],
-            image_HW[1] // patch_HW[1],
-        )
-        self.img_size = image_HW
-        self.patch_size = patch_HW
-        self.patches_resolution = patch_grid_size
-        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-        self.flatten_embedding = flatten_embedding
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-    def forward(self, x: Tensor) -> Tensor:
-        _, _, H, W = x.shape
-        patch_H, patch_W = self.patch_size
-        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
-        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
-        x = self.proj(x)  # B C H W
-        H, W = x.size(2), x.size(3)
-        x = x.flatten(2).transpose(1, 2)  # B HW C
-        x = self.norm(x)
-        if not self.flatten_embedding:
-            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
-        return x
-    def flops(self) -> float:
-        Ho, Wo = self.patches_resolution
-        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
-        if self.norm is not None:
-            flops += Ho * Wo * self.embed_dim
-        return flops
-class Mlp(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = nn.GELU,
-        drop: float = 0.0,
-        bias: bool = True,
-    ) -> None:
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
-        self.drop = nn.Dropout(drop)
-    def forward(self, x: Tensor) -> Tensor:
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-class SwiGLUFFN(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = None,
-        drop: float = 0.0,
-        bias: bool = True,
-    ) -> None:
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
-        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
-    def forward(self, x: Tensor) -> Tensor:
-        x12 = self.w12(x)
-        x1, x2 = x12.chunk(2, dim=-1)
-        hidden = F.silu(x1) * x2
-        return self.w3(hidden)
-try:
-    from xformers.ops import SwiGLU
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    SwiGLU = SwiGLUFFN
-    XFORMERS_AVAILABLE = False
-class SwiGLUFFNFused(SwiGLU):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = None,
-        drop: float = 0.0,
-        bias: bool = True,
-    ) -> None:
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
-        super().__init__(
-            in_features=in_features,
-            hidden_features=hidden_features,
-            out_features=out_features,
-            bias=bias,
-        )
-try:
-    from xformers.ops import memory_efficient_attention, unbind, fmha
-    from xformers.components.attention import ScaledDotProduct
-    from xformers.components import MultiHeadDispatch
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    logger.warning("xFormers not available")
-    XFORMERS_AVAILABLE = False
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int = 8,
-        qkv_bias: bool = False,
-        proj_bias: bool = True,
-        attn_drop: float = 0.0,
-        proj_drop: float = 0.0,
-        window_size: int = 0,
-    ) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = head_dim**-0.5
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim, bias=proj_bias)
-        self.proj_drop = nn.Dropout(proj_drop)
-        #if not self.training:
-        #
-        # self.attn = ScaledDotProduct()
-            #self.attn = MultiHeadDispatch(dim_model=EMB, residual_dropout=DROPOUT, num_heads=HEADS, attention=attn)
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
-        attn = q @ k.transpose(-2, -1)
-        if attn_bias is not None:
-            attn = attn + attn_bias[:, :, :N]
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-class MemEffAttention(Attention):
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
-        #if True:
-            assert attn_bias is None, "xFormers is required for nested tensors usage"
-            return super().forward(x, attn_bias)
-        B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
-        q, k, v = unbind(qkv, 2)
-        if attn_bias is not None:
-            x = memory_efficient_attention(q, k, v, attn_bias=attn_bias[:, :, :N])
-        else:
-            x = memory_efficient_attention(q, k, v)
-        x = x.reshape([B, N, C])
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-try:
-    from xformers.ops import fmha
-    from xformers.ops import scaled_index_add, index_select_cat
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    logger.warning("xFormers not available")
-    XFORMERS_AVAILABLE = False
-class Block(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        qkv_bias: bool = False,
-        proj_bias: bool = True,
-        ffn_bias: bool = True,
-        drop: float = 0.0,
-        attn_drop: float = 0.0,
-        init_values = None,
-        drop_path: float = 0.0,
-        act_layer: Callable[..., nn.Module] = nn.GELU,
-        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
-        attn_class: Callable[..., nn.Module] = Attention,
-        ffn_layer: Callable[..., nn.Module] = Mlp,
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.norm1 = norm_layer(dim)
-        self.attn = attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-        )
-        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = ffn_layer(
-            in_features=dim,
-            hidden_features=mlp_hidden_dim,
-            act_layer=act_layer,
-            drop=drop,
-            bias=ffn_bias,
-        )
-        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.sample_drop_ratio = drop_path
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        def attn_residual_func(x: Tensor, attn_bias) -> Tensor:
-            return self.ls1(self.attn(self.norm1(x), attn_bias))
-        def ffn_residual_func(x: Tensor) -> Tensor:
-            return self.ls2(self.mlp(self.norm2(x)))
-        if self.training and self.sample_drop_ratio > 0.1:
-            # the overhead is compensated only for a drop path rate larger than 0.1
-            x = drop_add_residual_stochastic_depth(
-                x,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                attn_bias=attn_bias
-            )
-            x = drop_add_residual_stochastic_depth(
-                x,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-            )
-        elif self.training and self.sample_drop_ratio > 0.0:
-            x = x + self.drop_path1(attn_residual_func(x, attn_bias))
-            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
-        else:
-            x = x + attn_residual_func(x, attn_bias)
-            x = x + ffn_residual_func(x)
-        return x
-def drop_add_residual_stochastic_depth(
-    x: Tensor,
-    residual_func: Callable[[Tensor], Tensor],
-    sample_drop_ratio: float = 0.0, attn_bias=None
-) -> Tensor:
-    # 1) extract subset using permutation
-    b, n, d = x.shape
-    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
-    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
-    x_subset = x[brange]
-    # 2) apply residual_func to get residual
-    residual = residual_func(x_subset, attn_bias)
-    x_flat = x.flatten(1)
-    residual = residual.flatten(1)
-    residual_scale_factor = b / sample_subset_size
-    # 3) add the residual
-    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
-    return x_plus_residual.view_as(x)
-def get_branges_scales(x, sample_drop_ratio=0.0):
-    b, n, d = x.shape
-    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
-    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
-    residual_scale_factor = b / sample_subset_size
-    return brange, residual_scale_factor
-def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
-    if scaling_vector is None:
-        x_flat = x.flatten(1)
-        residual = residual.flatten(1)
-        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
-    else:
-        x_plus_residual = scaled_index_add(
-            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
-        )
-    return x_plus_residual
-attn_bias_cache: Dict[Tuple, Any] = {}
-def get_attn_bias_and_cat(x_list, branges=None):
-    """
-    this will perform the index select, cat the tensors, and provide the attn_bias from cache
-    """
-    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
-    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
-    if all_shapes not in attn_bias_cache.keys():
-        seqlens = []
-        for b, x in zip(batch_sizes, x_list):
-            for _ in range(b):
-                seqlens.append(x.shape[1])
-        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
-        attn_bias._batch_sizes = batch_sizes
-        attn_bias_cache[all_shapes] = attn_bias
-    if branges is not None:
-        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
-    else:
-        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
-        cat_tensors = torch.cat(tensors_bs1, dim=1)
-    return attn_bias_cache[all_shapes], cat_tensors
-def drop_add_residual_stochastic_depth_list(
-    x_list: List[Tensor],
-    residual_func: Callable[[Tensor, Any], Tensor],
-    sample_drop_ratio: float = 0.0,
-    scaling_vector=None,
-) -> Tensor:
-    # 1) generate random set of indices for dropping samples in the batch
-    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
-    branges = [s[0] for s in branges_scales]
-    residual_scale_factors = [s[1] for s in branges_scales]
-    # 2) get attention bias and index+concat the tensors
-    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
-    # 3) apply residual_func to get residual, and split the result
-    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
-    outputs = []
-    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
-        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
-    return outputs
-class NestedTensorBlock(Block):
-    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
-        """
-        x_list contains a list of tensors to nest together and run
-        """
-        assert isinstance(self.attn, MemEffAttention)
-        if self.training and self.sample_drop_ratio > 0.0:
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.attn(self.norm1(x), attn_bias=attn_bias)
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.mlp(self.norm2(x))
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
-            )
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
-            )
-            return x_list
-        else:
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls2(self.mlp(self.norm2(x)))
-            attn_bias, x = get_attn_bias_and_cat(x_list)
-            x = x + attn_residual_func(x, attn_bias=attn_bias)
-            x = x + ffn_residual_func(x)
-            return attn_bias.split(x)
-    def forward(self, x_or_x_list, attn_bias=None):
-        if isinstance(x_or_x_list, Tensor):
-            return super().forward(x_or_x_list, attn_bias)
-        elif isinstance(x_or_x_list, list):
-            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
-            return self.forward_nested(x_or_x_list)
-        else:
-            raise AssertionError
-def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
-    if not depth_first and include_root:
-        fn(module=module, name=name)
-    for child_name, child_module in module.named_children():
-        child_name = ".".join((name, child_name)) if name else child_name
-        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
-    if depth_first and include_root:
-        fn(module=module, name=name)
-    return module
-class BlockChunk(nn.ModuleList):
-    def forward(self, x, others=None):
-        for b in self:
-            if others == None:
-                x = b(x)
-            else:
-                x = b(x, others)
-        return x
-class DinoVisionTransformer(nn.Module):
-    def __init__(
-        self,
-        img_size=224,
-        patch_size=16,
-        in_chans=3,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        ffn_bias=True,
-        proj_bias=True,
-        drop_path_rate=0.0,
-        drop_path_uniform=False,
-        #init_values=None,  # for layerscale: None or 0 => no layerscale
-        init_values=1e-5,  # for layerscale: None or 0 => no layerscale
-        embed_layer=PatchEmbed,
-        act_layer=nn.GELU,
-        block_fn=NestedTensorBlock,
-        ffn_layer="mlp",
-        block_chunks=1,
-        window_size=37,
-        **kwargs
-    ):
-        """
-        Args:
-            img_size (int, tuple): input image size
-            patch_size (int, tuple): patch size
-            in_chans (int): number of input channels
-            embed_dim (int): embedding dimension
-            depth (int): depth of transformer
-            num_heads (int): number of attention heads
-            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
-            qkv_bias (bool): enable bias for qkv if True
-            proj_bias (bool): enable bias for proj in attn if True
-            ffn_bias (bool): enable bias for ffn if True
-            drop_path_rate (float): stochastic depth rate
-            drop_path_uniform (bool): apply uniform drop rate across blocks
-            weight_init (str): weight init scheme
-            init_values (float): layer-scale init values
-            embed_layer (nn.Module): patch embedding layer
-            act_layer (nn.Module): MLP activation layer
-            block_fn (nn.Module): transformer block class
-            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
-            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
-        """
-        super().__init__()
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
-        self.num_tokens = 1
-        self.n_blocks = depth
-        self.num_heads = num_heads
-        self.patch_size = patch_size
-        self.window_size = window_size
-        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
-        num_patches = self.patch_embed.num_patches
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
-        if drop_path_uniform is True:
-            dpr = [drop_path_rate] * depth
-        else:
-            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
-        if ffn_layer == "mlp":
-            logger.info("using MLP layer as FFN")
-            ffn_layer = Mlp
-        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
-            logger.info("using SwiGLU layer as FFN")
-            ffn_layer = SwiGLUFFNFused
-        elif ffn_layer == "identity":
-            logger.info("using Identity layer as FFN")
-            def f(*args, **kwargs):
-                return nn.Identity()
-            ffn_layer = f
-        else:
-            raise NotImplementedError
-        blocks_list = [
-            block_fn(
-                dim=embed_dim,
-                num_heads=num_heads,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                proj_bias=proj_bias,
-                ffn_bias=ffn_bias,
-                drop_path=dpr[i],
-                norm_layer=norm_layer,
-                act_layer=act_layer,
-                ffn_layer=ffn_layer,
-                init_values=init_values,
-            )
-            for i in range(depth)
-        ]
-        if block_chunks > 0:
-            self.chunked_blocks = True
-            chunked_blocks = []
-            chunksize = depth // block_chunks
-            for i in range(0, depth, chunksize):
-                # this is to keep the block index consistent if we chunk the block list
-                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
-            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
-        else:
-            self.chunked_blocks = False
-            self.blocks = nn.ModuleList(blocks_list)
-        self.norm = norm_layer(embed_dim)
-        self.head = nn.Identity()
-        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
-        self.init_weights()
-    def init_weights(self):
-        trunc_normal_(self.pos_embed, std=0.02)
-        nn.init.normal_(self.cls_token, std=1e-6)
-        named_apply(init_weights_vit_timm, self)
-    def interpolate_pos_encoding(self, x, w, h):
-        previous_dtype = x.dtype
-        npatch = x.shape[1] - 1
-        N = self.pos_embed.shape[1] - 1
-        if npatch == N and w == h:
-            return self.pos_embed
-        pos_embed = self.pos_embed.float()
-        class_pos_embed = pos_embed[:, 0]
-        patch_pos_embed = pos_embed[:, 1:]
-        dim = x.shape[-1]
-        w0 = w // self.patch_size
-        h0 = h // self.patch_size
-        # we add a small number to avoid floating point error in the interpolation
-        # see discussion at https://github.com/facebookresearch/dino/issues/8
-        w0, h0 = w0 + 0.1, h0 + 0.1
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
-            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
-            mode="bicubic",
-        )
-        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
-    def prepare_tokens_with_masks(self, x, masks=None):
-        B, nc, w, h = x.shape
-        x = self.patch_embed(x)
-        if masks is not None:
-            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
-        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
-        x = x + self.interpolate_pos_encoding(x, w, h)
-        return x
-    def forward_features_list(self, x_list, masks_list):
-        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
-        for blk in self.blocks:
-            x = blk(x)
-        all_x = x
-        output = []
-        for x, masks in zip(all_x, masks_list):
-            x_norm = self.norm(x)
-            output.append(
-                {
-                    "x_norm_clstoken": x_norm[:, 0],
-                    "x_norm_patchtokens": x_norm[:, 1:],
-                    "x_prenorm": x,
-                    "masks": masks,
-                }
-            )
-        return output
-    def forward_features(self, x, masks=None):
-        if isinstance(x, list):
-            return self.forward_features_list(x, masks)
-        B, C, H, W = x.size()
-        pad_h = (self.patch_size - H % self.patch_size)
-        pad_w = (self.patch_size - W % self.patch_size)
-        if pad_h == self.patch_size:
-            pad_h = 0
-        if pad_w == self.patch_size:
-            pad_w = 0
-        #x = nn.functional.pad(x, (pad_h//2, pad_h-pad_h//2, pad_w//2, pad_w-pad_w//2))
-        if pad_h + pad_w > 0:
-            x = torch.nn.functional.interpolate(x, (H+pad_h, W+pad_w), mode='bilinear')
-        x = self.prepare_tokens_with_masks(x, masks)
-        features = []
-        for blk in self.blocks:
-            x = blk(x)
-        # for idx in range(len(self.blocks[0])):
-        #     x = self.blocks[0][idx](x)
-        #     if (idx + 1) % (len(self.blocks[0]) // 4) == 0:
-        #         features.append(x)
-        #return [features, (B, (H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size, H, W)]
-        x_norm = self.norm(x)
-        # return {
-        #     "x_norm_clstoken": x_norm[:, 0],
-        #     "x_norm_patchtokens": x_norm[:, 1:],
-        #     "x_prenorm": x,
-        #     "masks": masks,
-        # }
-        features = []
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        return [features, (B, (H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size, H, W)]
-    def _get_intermediate_layers_not_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        # If n is an int, take the n last blocks. If it's a list, take them
-        output, total_block_len = [], len(self.blocks)
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for i, blk in enumerate(self.blocks):
-            x = blk(x)
-            if i in blocks_to_take:
-                output.append(x)
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def _get_intermediate_layers_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        output, i, total_block_len = [], 0, len(self.blocks[-1])
-        # If n is an int, take the n last blocks. If it's a list, take them
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for block_chunk in self.blocks:
-            for blk in block_chunk[i:]:  # Passing the nn.Identity()
-                x = blk(x)
-                if i in blocks_to_take:
-                    output.append(x)
-                i += 1
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def get_intermediate_layers(
-        self,
-        x: torch.Tensor,
-        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
-        reshape: bool = False,
-        return_class_token: bool = False,
-        norm=True,
-    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
-        if self.chunked_blocks:
-            outputs = self._get_intermediate_layers_chunked(x, n)
-        else:
-            outputs = self._get_intermediate_layers_not_chunked(x, n)
-        if norm:
-            outputs = [self.norm(out) for out in outputs]
-        class_tokens = [out[:, 0] for out in outputs]
-        outputs = [out[:, 1:] for out in outputs]
-        if reshape:
-            B, _, w, h = x.shape
-            outputs = [
-                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
-                for out in outputs
-            ]
-        if return_class_token:
-            return tuple(zip(outputs, class_tokens))
-        return tuple(outputs)
-    def forward(self, *args, is_training=False, **kwargs):
-        ret = self.forward_features(*args, **kwargs)
-        return ret
-        # if is_training:
-        #     return ret
-        # else:
-        #     return self.head(ret["x_norm_clstoken"])
-class PosConv(nn.Module):
-    # PEG  from https://arxiv.org/abs/2102.10882
-    def __init__(self, in_chans, embed_dim=768, stride=1):
-        super(PosConv, self).__init__()
-        self.proj = nn.Sequential(
-            nn.Conv2d(in_chans, embed_dim, 37, stride, 18, bias=True, groups=embed_dim),
-        )
-        self.stride = stride
-    def forward(self, x, size):
-        B, N, C = x.shape
-        cnn_feat_token = x.transpose(1, 2).view(B, C, *size)
-        x = self.proj(cnn_feat_token)
-        if self.stride == 1:
-            x += cnn_feat_token
-        x = x.flatten(2).transpose(1, 2)
-        return x
-    #def no_weight_decay(self):
-        #return ['proj.%d.weight' % i for i in range(4)]
-class DinoWindowVisionTransformer(nn.Module):
-    def __init__(
-        self,
-        img_size=224,
-        patch_size=16,
-        in_chans=3,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        ffn_bias=True,
-        proj_bias=True,
-        drop_path_rate=0.0,
-        drop_path_uniform=False,
-        #init_values=None,  # for layerscale: None or 0 => no layerscale
-        init_values=1e-5,  # for layerscale: None or 0 => no layerscale
-        embed_layer=PatchEmbed,
-        act_layer=nn.GELU,
-        block_fn=NestedTensorBlock,
-        ffn_layer="mlp",
-        block_chunks=1,
-        window_size=7,
-        **kwargs
-    ):
-        """
-        Args:
-            img_size (int, tuple): input image size
-            patch_size (int, tuple): patch size
-            in_chans (int): number of input channels
-            embed_dim (int): embedding dimension
-            depth (int): depth of transformer
-            num_heads (int): number of attention heads
-            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
-            qkv_bias (bool): enable bias for qkv if True
-            proj_bias (bool): enable bias for proj in attn if True
-            ffn_bias (bool): enable bias for ffn if True
-            drop_path_rate (float): stochastic depth rate
-            drop_path_uniform (bool): apply uniform drop rate across blocks
-            weight_init (str): weight init scheme
-            init_values (float): layer-scale init values
-            embed_layer (nn.Module): patch embedding layer
-            act_layer (nn.Module): MLP activation layer
-            block_fn (nn.Module): transformer block class
-            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
-            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
-        """
-        super().__init__()
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
-        self.num_tokens = 1
-        self.n_blocks = depth
-        self.num_heads = num_heads
-        self.patch_size = patch_size
-        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
-        num_patches = self.patch_embed.num_patches
-        #self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        #self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
-        self.pos_conv = PosConv(self.embed_dim, self.embed_dim)
-        self.window_size = window_size
-        #self.conv_block = nn.ModuleList([ConvBlock(embed_dim) for i in range(4)])
-        #self.conv_block = nn.ModuleList([nn.Identity() for i in range(4)])
-        if drop_path_uniform is True:
-            dpr = [drop_path_rate] * depth
-        else:
-            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
-        if ffn_layer == "mlp":
-            logger.info("using MLP layer as FFN")
-            ffn_layer = Mlp
-        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
-            logger.info("using SwiGLU layer as FFN")
-            ffn_layer = SwiGLUFFNFused
-        elif ffn_layer == "identity":
-            logger.info("using Identity layer as FFN")
-            def f(*args, **kwargs):
-                return nn.Identity()
-            ffn_layer = f
-        else:
-            raise NotImplementedError
-        blocks_list = [
-            block_fn(
-                dim=embed_dim,
-                num_heads=num_heads,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                proj_bias=proj_bias,
-                ffn_bias=ffn_bias,
-                drop_path=dpr[i],
-                norm_layer=norm_layer,
-                act_layer=act_layer,
-                ffn_layer=ffn_layer,
-                init_values=init_values,
-            )
-            for i in range(depth)
-        ]
-        if block_chunks > 0:
-            self.chunked_blocks = True
-            chunked_blocks = []
-            chunksize = depth // block_chunks
-            for i in range(0, depth, chunksize):
-                # this is to keep the block index consistent if we chunk the block list
-                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
-            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
-        else:
-            self.chunked_blocks = False
-            self.blocks = nn.ModuleList(blocks_list)
-        self.norm = norm_layer(embed_dim)
-        self.head = nn.Identity()
-        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
-        self.nh = -1
-        self.nw = -1
-        try:
-            H = cfg.data_basic['crop_size'][0]
-            W = cfg.data_basic['crop_size'][1]
-            pad_h = (self.patch_size - H % self.patch_size)
-            pad_w = (self.patch_size - W % self.patch_size)
-            if pad_h == self.patch_size:
-                pad_h = 0
-            if pad_w == self.patch_size:
-                pad_w = 0
-            self.nh = (H + pad_h) // self.patch_size
-            self.nw = (W + pad_w) // self.patch_size
-            self.prepare_attn_bias((self.nh, self.nw))
-        except:
-            pass
-        self.init_weights()
-        self.total_step = 10000 # For PE -> GPE transfer
-        self.start_step = 2000
-        self.current_step = 20000
-    def init_weights(self):
-        #trunc_normal_(self.pos_embed, std=0.02)
-        #nn.init.normal_(self.cls_token, std=1e-6)
-        named_apply(init_weights_vit_timm, self)
-        for i in range(4):
-            try:
-                nn.init.constant_(self.conv_block[i].conv2.weight, 0.0)
-            except:
-                pass
-    def interpolate_pos_encoding(self, x, w, h):
-        previous_dtype = x.dtype
-        #npatch = x.shape[1] - 1
-        #N = self.pos_embed.shape[1] - 1
-        npatch = x.shape[1]
-        N = self.pos_embed.shape[1]
-        if npatch == N and w == h:
-            return self.pos_embed
-        pos_embed = self.pos_embed.float()
-        #class_pos_embed = pos_embed[:, 0]
-        #patch_pos_embed = pos_embed[:, 1:]
-        patch_pos_embed = pos_embed
-        dim = x.shape[-1]
-        w0 = w // self.patch_size
-        h0 = h // self.patch_size
-        # we add a small number to avoid floating point error in the interpolation
-        # see discussion at https://github.com/facebookresearch/dino/issues/8
-        w0, h0 = w0 + 0.1, h0 + 0.1
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
-            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
-            mode="bicubic",
-        )
-        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return patch_pos_embed.to(previous_dtype)
-        #return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
-    def window_partition(self, x: torch.Tensor, window_size: int, hw: Tuple[int, int], conv_feature=False) -> Tuple[torch.Tensor, Tuple[int, int]]:
-        """
-        Partition into non-overlapping windows with padding if needed.
-        Args:
-            x (tensor): input tokens with [B, H, W, C].
-            window_size (int): window size.
-        Returns:
-            windows: windows after partition with [B * num_windows, window_size, window_size, C].
-            (Hp, Wp): padded height and width before partition
-        """
-        if conv_feature == False:
-            B, N, C = x.shape
-            H, W = hw[0], hw[1]
-            x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
-            windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size * window_size, C)
-        else:
-            B, C, H, W = x.shape
-            x = x.view(B, C, H // window_size, window_size, W // window_size, window_size)
-            windows = x.permute(0, 2, 4, 3, 5, 1).contiguous().view(-1, window_size * window_size, C)
-        #y = torch.cat((x_cls, windows), dim=1)
-        return windows   #, (Hp, Wp)
-    def window_unpartition(self,
-        windows: torch.Tensor, window_size: int, hw: Tuple[int, int], conv_feature=False
-    ) -> torch.Tensor:
-        """
-        Window unpartition into original sequences and removing padding.
-        Args:
-            windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].
-            window_size (int): window size.
-            pad_hw (Tuple): padded height and width (Hp, Wp).
-            hw (Tuple): original height and width (H, W) before padding.
-        Returns:
-            x: unpartitioned sequences with [B, H, W, C].
-        """
-        H, W = hw
-        B = windows.shape[0] // (H * W // window_size // window_size)
-        x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
-        if conv_feature == False:
-            x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp * Wp, -1)
-        else:
-            C = windows.shape[-1]
-            x = x.permute(0, 5, 1, 3, 2, 4).contiguous().view(B, C, H, W)
-        # if Hp > H or Wp > W:
-        #     x = x[:, :H, :W, :].contiguous()
-        return x
-    def prepare_tokens_with_masks(self, x, masks=None, step=-1):
-        B, nc, w, h = x.shape
-        x = self.patch_embed(x)
-        if masks is not None:
-            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
-        #x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
-        if step == -1:
-            step = self.current_step
-        else:
-            self.current_step = step
-        if step < self.start_step:
-            coef = 0.0
-        elif step < self.total_step:
-            coef = (step - self.start_step) / (self.total_step - self.start_step)
-        else:
-            coef = 1.0
-        x = x + (1 - coef) * self.interpolate_pos_encoding(x, w, h) + coef * self.pos_conv(x, (self.nh, self.nw))
-        return x
-    def prepare_attn_bias(self, shape):
-        window_size = self.window_size
-        if window_size <= 0:
-            return
-        import xformers.components.attention.attention_patterns as AP
-        nh, nw = shape
-        radius = (window_size-1)//2
-        mask_ori = AP.local_2d_pattern(nh, nw, distance = radius + 0.1, p=torch.inf).cuda()
-        pad = (8 - (nh * nw) % 8)
-        if pad == 8:
-            pad = 0
-        mask_pad = nn.functional.pad(mask_ori, (0, pad)).contiguous()
-        if pad > 0:
-            mask = mask_pad[:, :-pad].view(nh, nw, nh, nw)
-        else:
-            mask = mask_pad[:, :].view(nh, nw, nh, nw)
-        # angle
-        mask[:radius+1,  :radius+1,  :window_size,  :window_size] = True
-        mask[:radius+1,  -radius-1:, :window_size,  -window_size:] = True
-        mask[-radius-1:, :radius+1,  -window_size:, :window_size] = True
-        mask[-radius-1:, -radius-1:, -window_size:, -window_size:] = True
-        # edge
-        mask[radius+1:-radius-1,  :radius+1,  :,  :] = mask[radius+1:-radius-1,  radius:radius+1,    :,  :]
-        mask[radius+1:-radius-1,  -radius-1:, :,  :] = mask[radius+1:-radius-1,  -radius-1:-radius,  :,  :]
-        mask[:radius+1,   radius+1:-radius-1, :,  :] = mask[radius:radius+1,   radius+1:-radius-1,   :,  :]
-        mask[-radius-1:,  radius+1:-radius-1, :,  :] = mask[-radius-1:-radius, radius+1:-radius-1,   :,  :]
-        mask = mask.view(nh*nw, nh*nw)
-        bias_pad = torch.log(mask_pad)
-        #bias = bias_pad[:, :-pad]
-        self.register_buffer('attn_bias', bias_pad)
-        return bias_pad
-    def forward_features_list(self, x_list, masks_list):
-        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
-        for blk in self.blocks:
-            x = blk(x)
-        all_x = x
-        output = []
-        for x, masks in zip(all_x, masks_list):
-            x_norm = self.norm(x)
-            output.append(
-                {
-                    "x_norm_clstoken": x_norm[:, 0],
-                    "x_norm_patchtokens": x_norm[:, 1:],
-                    "x_prenorm": x,
-                    "masks": masks,
-                }
-            )
-        return output
-    def forward_features(self, x, masks=None, **kwargs):
-        if isinstance(x, list):
-            return self.forward_features_list(x, masks)
-        B, C, H, W = x.size()
-        pad_h = (self.patch_size - H % self.patch_size)
-        pad_w = (self.patch_size - W % self.patch_size)
-        if pad_h == self.patch_size:
-            pad_h = 0
-        if pad_w == self.patch_size:
-            pad_w = 0
-        #x = nn.functional.pad(x, (pad_h//2, pad_h-pad_h//2, pad_w//2, pad_w-pad_w//2))
-        if pad_h + pad_w > 0:
-            x = torch.nn.functional.interpolate(x, (H+pad_h, W+pad_w), mode='bilinear')
-        nh = (H+pad_h)//self.patch_size
-        nw = (W+pad_w)//self.patch_size
-        if self.window_size > 0:
-            if nh == self.nh and nw == self.nw:
-                attn_bias = self.attn_bias
-            else:
-                attn_bias = self.prepare_attn_bias(((H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size))
-                self.nh = nh
-                self.nw = nw
-            attn_bias = attn_bias.unsqueeze(0).repeat(B * self.num_heads, 1, 1)
-        else:
-            attn_bias = None
-        x = self.prepare_tokens_with_masks(x, masks)
-        #x = self.patch_embed(x)
-        features = []
-        #x = self.window_partition(x, self.window_size, (H // self.patch_size, W // self.patch_size))
-        for blk in self.blocks:
-            x = blk(x, attn_bias)
-        #x = self.window_unpartition(x, self.window_size, (H // self.patch_size, W // self.patch_size))
-        # for idx in range(len(self.blocks[0])):
-        #     x = self.blocks[0][idx](x, attn_bias)
-        #     if (idx + 1) % (len(self.blocks[0]) // 4) == 0:
-        #         x = self.window_unpartition(x, self.window_size, (H // self.patch_size, W // self.patch_size), conv_feature=True)
-        #         x = self.conv_block[idx // (len(self.blocks[0]) // 4)](x)
-        #         if idx + 1 != len(self.blocks[0]):
-        #             x = self.window_partition(x, self.window_size, (H // self.patch_size, W // self.patch_size), conv_feature=True)
-        #         else:
-        #             b, c, h, w = x.size()
-        #             x = x.permute(0, 2, 3, 1).contiguous().view(b, h, w, c)
-                #features.append(x)
-        #return [features, (B, (H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size, H, W)]
-        x_norm = self.norm(x)
-        # return {
-        #     "x_norm_clstoken": x_norm[:, 0],
-        #     "x_norm_patchtokens": x_norm[:, 1:],
-        #     "x_prenorm": x,
-        #     "masks": masks,
-        # }
-        features = []
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        return [features, (B, (H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size, H, W)]
-    def _get_intermediate_layers_not_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        # If n is an int, take the n last blocks. If it's a list, take them
-        output, total_block_len = [], len(self.blocks)
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for i, blk in enumerate(self.blocks):
-            x = blk(x)
-            if i in blocks_to_take:
-                output.append(x)
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def _get_intermediate_layers_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        output, i, total_block_len = [], 0, len(self.blocks[-1])
-        # If n is an int, take the n last blocks. If it's a list, take them
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for block_chunk in self.blocks:
-            for blk in block_chunk[i:]:  # Passing the nn.Identity()
-                x = blk(x)
-                if i in blocks_to_take:
-                    output.append(x)
-                i += 1
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def get_intermediate_layers(
-        self,
-        x: torch.Tensor,
-        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
-        reshape: bool = False,
-        return_class_token: bool = False,
-        norm=True,
-    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
-        if self.chunked_blocks:
-            outputs = self._get_intermediate_layers_chunked(x, n)
-        else:
-            outputs = self._get_intermediate_layers_not_chunked(x, n)
-        if norm:
-            outputs = [self.norm(out) for out in outputs]
-        class_tokens = [out[:, 0] for out in outputs]
-        outputs = [out[:, 1:] for out in outputs]
-        if reshape:
-            B, _, w, h = x.shape
-            outputs = [
-                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
-                for out in outputs
-            ]
-        if return_class_token:
-            return tuple(zip(outputs, class_tokens))
-        return tuple(outputs)
-    def forward(self, *args, is_training=False, **kwargs):
-        ret = self.forward_features(*args, **kwargs)
-        return ret
-        # if is_training:
-        #     return ret
-        # else:
-        #     return self.head(ret["x_norm_clstoken"])
-def init_weights_vit_timm(module: nn.Module, name: str = ""):
-    """ViT weight initialization, original timm impl (for reproducibility)"""
-    if isinstance(module, nn.Linear):
-        trunc_normal_(module.weight, std=0.02)
-        if module.bias is not None:
-            nn.init.zeros_(module.bias)
-def vit_small(patch_size=14, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=384,
-        depth=12,
-        num_heads=6,
-        mlp_ratio=4,
-        block_fn=partial(NestedTensorBlock, attn_class=MemEffAttention),
-        **kwargs,
-    )
-    return model
-def vit_base(patch_size=14, **kwargs):
-    model = DinoWindowVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4,
-        block_fn=partial(NestedTensorBlock, attn_class=MemEffAttention),
-        **kwargs,
-    )
-    return model
-def vit_large(patch_size=14, checkpoint=None, **kwargs):
-    model = DinoVisionTransformer(
-        img_size = 518,
-        patch_size=patch_size,
-        embed_dim=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4,
-        block_fn=partial(NestedTensorBlock, attn_class=MemEffAttention),
-        **kwargs,
-    )
-    if checkpoint is not None:
-        with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
-        try:
-            model.load_state_dict(state_dict, strict=True)
-        except:
-            new_state_dict = {}
-            for key, value in state_dict.items():
-                if 'blocks' in key:
-                    key_new = 'blocks.0' + key[len('blocks'):]
-                else:
-                    key_new = key
-                new_state_dict[key_new] = value
-            model.load_state_dict(new_state_dict, strict=True)
-        #del model.norm
-        del model.mask_token
-    return model
-    # model = DinoWindowVisionTransformer(
-    #     img_size = 518,
-    #     patch_size=patch_size,
-    #     embed_dim=1024,
-    #     depth=24,
-    #     num_heads=16,
-    #     mlp_ratio=4,
-    #     block_fn=partial(NestedTensorBlock, attn_class=MemEffAttention),
-    #     window_size=37,
-    #     **kwargs,
-    # )
-    # if checkpoint is not None:
-    #     with open(checkpoint, "rb") as f:
-    #         state_dict = torch.load(f)
-    #     try:
-    #         model.load_state_dict(state_dict, strict=True)
-    #     except:
-    #         new_state_dict = {}
-    #         for key, value in state_dict.items():
-    #             if 'blocks' in key:
-    #                 key_new = 'blocks.0' + key[len('blocks'):]
-    #             else:
-    #                 key_new = key
-    #             if 'pos_embed' in key:
-    #                 value = value[:, 1:, :]
-    #             new_state_dict[key_new] = value
-    #         model.load_state_dict(new_state_dict, strict=False)
-    #     #del model.norm
-    #     del model.mask_token
-    return model
-def vit_giant2(patch_size=16, **kwargs):
-    """
-    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
-    """
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=1536,
-        depth=40,
-        num_heads=24,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        **kwargs,
-    )
-    return model
-if __name__ == '__main__':
-    try:
-        from mmcv.utils import Config
-    except:
-        from mmengine import Config
-    #rgb = torch.rand((2, 3, 518, 518)).cuda()
-    #cfg.data_basic['crop_size']['0']
-    #cfg.data_basic['crop_size']['1']
-    cfg = Config.fromfile('/cpfs01/user/mu.hu/monodepth/mono/configs/HourglassDecoder/pub12.convlarge.0.3_150.py')
-    #rgb = torch.arange(0, 2*3*1036*1036, 1).cuda().float().view(2, 3, 1036, 1036)
-    rgb = torch.zeros(1, 3, 1400, 1680).cuda()
-    model = vit_large(checkpoint="/cpfs02/shared/public/custom/group_local_map/yvan/pretrained_weight_repo/vit/dinov2_vitl14_pretrain.pth", kwarg=cfg).cuda()
-    #import timm
-    #model2 = timm.models.vision_transformer.vit_large_patch14_dinov2().cuda()
-    #timm.models.load_checkpoint(model2, '/cpfs02/shared/public/yvan/pretrained_weight_repo/vit/dinov2_vitl14_pretrain.pth', filter_fn=timm.models.vision_transformer.checkpoint_filter_fn)
-    out1 = model(rgb)
-    #out2 = model2(rgb)
-    temp = 0
-# import time
-# window_size = 37
-# def prepare_window_masks(shape):
-#     if window_size <= 0:
-#         return None
-#     import xformers.components.attention.attention_patterns as AP
-#     B, nh, nw, _, _ = shape
-#     radius = (window_size-1)//2
-#     #time0 = time.time()
-#     d = AP.local_nd_distance(nh, nw, distance = radius + 0.1, p=torch.inf).cuda()
-#     #mask = AP.local_2d_pattern(nh, nw, distance = radius + 0.1, p=torch.inf).cuda()
-#     # mask = mask.view(nh, nw, nh, nw)
-#     # #time1 = time.time() - time0
-#     # # angle
-#     # mask[:radius+1,  :radius+1,  :window_size,  :window_size] = True
-#     # mask[:radius+1,  -radius-1:, :window_size,  -window_size:] = True
-#     # mask[-radius-1:, :radius+1,  -window_size:, :window_size] = True
-#     # mask[-radius-1:, -radius-1:, -window_size:, -window_size:] = True
-#     # time2 = time.time() - time0 - time1
-#     # # edge
-#     # mask[radius+1:-radius-1,  :radius+1,  :,  :] = mask[radius+1:-radius-1,  radius:radius+1,    :,  :]
-#     # mask[radius+1:-radius-1,  -radius-1:, :,  :] = mask[radius+1:-radius-1,  -radius-1:-radius,  :,  :]
-#     # mask[:radius+1,   radius+1:-radius-1, :,  :] = mask[radius:radius+1,   radius+1:-radius-1,   :,  :]
-#     # mask[-radius-1:,  radius+1:-radius-1, :,  :] = mask[-radius-1:-radius, radius+1:-radius-1,   :,  :]
-#     # time3 = time.time() - time0 - time2
-#     # print(time1, time2, time3)
-# #     return mask.view(nw*nw, nh*nw).unsqueeze(0).repeat(B, 1)
-# shape = (1, 55, 55, None, None)
-# mask = prepare_window_masks(shape)
-# # temp = 1

mono/model/backbones/ViT_DINO_reg.py DELETED Viewed

@@ -1,1293 +0,0 @@
-# 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.
-# References:
-#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
-from functools import partial
-import math
-import logging
-from typing import Sequence, Tuple, Union, Callable, Optional, Dict, Any, List
-import torch
-import torch.nn as nn
-from torch import Tensor
-import torch.utils.checkpoint
-from torch.nn.init import trunc_normal_
-import torch.nn.init
-import torch.nn.functional as F
-#from dinov2.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
-logger = logging.getLogger("dinov2")
-# SSF finetuning originally by dongzelian
-def init_ssf_scale_shift(dim):
-    scale = nn.Parameter(torch.ones(dim))
-    shift = nn.Parameter(torch.zeros(dim))
-    nn.init.normal_(scale, mean=1, std=.02)
-    nn.init.normal_(shift, std=.02)
-    return scale, shift
-def ssf_ada(x, scale, shift):
-    assert scale.shape == shift.shape
-    if x.shape[-1] == scale.shape[0]:
-        return x * scale + shift
-    elif x.shape[1] == scale.shape[0]:
-        return x * scale.view(1, -1, 1, 1) + shift.view(1, -1, 1, 1)
-    else:
-        raise ValueError('the input tensor shape does not match the shape of the scale factor.')
-# LoRA finetuning originally by edwardjhu
-class LoRALayer():
-    def __init__(
-        self,
-        r: int,
-        lora_alpha: int,
-        lora_dropout: float,
-        merge_weights: bool,
-    ):
-        self.r = r
-        self.lora_alpha = lora_alpha
-        # Optional dropout
-        if lora_dropout > 0.:
-            self.lora_dropout = nn.Dropout(p=lora_dropout)
-        else:
-            self.lora_dropout = lambda x: x
-        # Mark the weight as unmerged
-        self.merged = False
-        self.merge_weights = merge_weights
-class LoRALinear(nn.Linear, LoRALayer):
-    # LoRA implemented in a dense layer
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        r: int = 0,
-        lora_alpha: int = 1,
-        lora_dropout: float = 0.,
-        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
-        merge_weights: bool = True,
-        **kwargs
-    ):
-        nn.Linear.__init__(self, in_features, out_features, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
-                           merge_weights=merge_weights)
-        self.fan_in_fan_out = fan_in_fan_out
-        # Actual trainable parameters
-        if r > 0:
-            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
-            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
-            self.scaling = self.lora_alpha / self.r
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False
-        self.reset_parameters()
-        if fan_in_fan_out:
-            self.weight.data = self.weight.data.transpose(0, 1)
-    def reset_parameters(self):
-        #nn.Linear.reset_parameters(self)
-        if hasattr(self, 'lora_A'):
-            # initialize B the same way as the default for nn.Linear and A to zero
-            # this is different than what is described in the paper but should not affect performance
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-    # def train(self, mode: bool = True):
-    #     def T(w):
-    #         return w.transpose(0, 1) if self.fan_in_fan_out else w
-    #     nn.Linear.train(self, mode)
-    #     if mode:
-    #         if self.merge_weights and self.merged:
-    #             # Make sure that the weights are not merged
-    #             if self.r > 0:
-    #                 self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
-    #             self.merged = False
-    #     else:
-    #         if self.merge_weights and not self.merged:
-    #             # Merge the weights and mark it
-    #             if self.r > 0:
-    #                 self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
-    #             self.merged = True
-    def forward(self, x: torch.Tensor):
-        def T(w):
-            return w.transpose(0, 1) if self.fan_in_fan_out else w
-        if self.r > 0 and not self.merged:
-            result = F.linear(x, T(self.weight), bias=self.bias)
-            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
-            return result
-        else:
-            return F.linear(x, T(self.weight), bias=self.bias)
-def make_2tuple(x):
-    if isinstance(x, tuple):
-        assert len(x) == 2
-        return x
-    assert isinstance(x, int)
-    return (x, x)
-def drop_path(x, drop_prob: float = 0.0, training: bool = False):
-    if drop_prob == 0.0 or not training:
-        return x
-    keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
-    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
-    if keep_prob > 0.0:
-        random_tensor.div_(keep_prob)
-    output = x * random_tensor
-    return output
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-    def __init__(self, drop_prob=None):
-        super(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training)
-class LayerScale(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        init_values: Union[float, Tensor] = 1e-5,
-        inplace: bool = False,
-    ) -> None:
-        super().__init__()
-        self.inplace = inplace
-        self.gamma = nn.Parameter(init_values * torch.ones(dim))
-    def forward(self, x: Tensor) -> Tensor:
-        return x.mul_(self.gamma) if self.inplace else x * self.gamma
-class PatchEmbed(nn.Module):
-    """
-    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
-    Args:
-        img_size: Image size.
-        patch_size: Patch token size.
-        in_chans: Number of input image channels.
-        embed_dim: Number of linear projection output channels.
-        norm_layer: Normalization layer.
-    """
-    def __init__(
-        self,
-        img_size: Union[int, Tuple[int, int]] = 224,
-        patch_size: Union[int, Tuple[int, int]] = 16,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-        norm_layer: Optional[Callable] = None,
-        flatten_embedding: bool = True,
-        tuning_mode: Optional[str] = None
-    ) -> None:
-        super().__init__()
-        image_HW = make_2tuple(img_size)
-        patch_HW = make_2tuple(patch_size)
-        patch_grid_size = (
-            image_HW[0] // patch_HW[0],
-            image_HW[1] // patch_HW[1],
-        )
-        self.img_size = image_HW
-        self.patch_size = patch_HW
-        self.patches_resolution = patch_grid_size
-        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-        self.flatten_embedding = flatten_embedding
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(embed_dim)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-    def forward(self, x: Tensor) -> Tensor:
-        _, _, H, W = x.shape
-        patch_H, patch_W = self.patch_size
-        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
-        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
-        x = self.proj(x)  # B C H W
-        H, W = x.size(2), x.size(3)
-        x = x.flatten(2).transpose(1, 2)  # B HW C
-        x = self.norm(x)
-        if self.tuning_mode == 'ssf':
-            x = ssf_ada(x, self.ssf_scale_1, self.ssf_shift_1)
-        if not self.flatten_embedding:
-            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
-        return x
-    def flops(self) -> float:
-        Ho, Wo = self.patches_resolution
-        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
-        if self.norm is not None:
-            flops += Ho * Wo * self.embed_dim
-        return flops
-class Mlp(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = nn.GELU,
-        drop: float = 0.0,
-        bias: bool = True,
-        tuning_mode: Optional[int] = None
-    ) -> None:
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
-        self.drop = nn.Dropout(drop)
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(hidden_features)
-                self.ssf_scale_2, self.ssf_shift_2 = init_ssf_scale_shift(out_features)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-    def forward(self, x: Tensor) -> Tensor:
-        x = self.fc1(x)
-        if self.tuning_mode == 'ssf':
-            x = ssf_ada(x, self.ssf_scale_1, self.ssf_shift_1)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        if self.tuning_mode == 'ssf':
-            x = ssf_ada(x, self.ssf_scale_2, self.ssf_shift_2)
-        x = self.drop(x)
-        return x
-class SwiGLUFFN(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = None,
-        drop: float = 0.0,
-        bias: bool = True,
-        tuning_mode: Optional[int] = None
-    ) -> None:
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
-        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(2 * hidden_features)
-                self.ssf_scale_2, self.ssf_shift_2 = init_ssf_scale_shift(out_features)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-    def forward(self, x: Tensor) -> Tensor:
-        x12 = self.w12(x)
-        if self.tuning_mode == 'ssf':
-            x12 = ssf_ada(x12, self.ssf_scale_1, self.ssf_shift_1)
-        x1, x2 = x12.chunk(2, dim=-1)
-        hidden = F.silu(x1) * x2
-        out = self.w3(hidden)
-        if self.tuning_mode == 'ssf':
-            out = ssf_ada(out, self.ssf_scale_2, self.ssf_scale_2)
-        return out
-try:
-    from xformers.ops import SwiGLU
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    SwiGLU = SwiGLUFFN
-    XFORMERS_AVAILABLE = False
-class SwiGLUFFNFused(SwiGLU):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-        act_layer: Callable[..., nn.Module] = None,
-        drop: float = 0.0,
-        bias: bool = True,
-    ) -> None:
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
-        super().__init__(
-            in_features=in_features,
-            hidden_features=hidden_features,
-            out_features=out_features,
-            bias=bias,
-        )
-try:
-    from xformers.ops import memory_efficient_attention, unbind, fmha
-    from xformers.components.attention import ScaledDotProduct
-    from xformers.components import MultiHeadDispatch
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    logger.warning("xFormers not available")
-    XFORMERS_AVAILABLE = False
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int = 8,
-        qkv_bias: bool = False,
-        proj_bias: bool = True,
-        attn_drop: float = 0.0,
-        proj_drop: float = 0.0,
-        window_size: int = 0,
-        tuning_mode: Optional[int] = None
-    ) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = head_dim**-0.5
-        if tuning_mode == 'lora':
-            self.tuning_mode = tuning_mode
-            self.qkv = LoRALinear(dim, dim * 3, bias=qkv_bias, r=8)
-        else:
-            self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        if tuning_mode == 'lora':
-            self.tuning_mode = tuning_mode
-            self.proj = LoRALinear(dim, dim, bias=proj_bias, r=8)
-        else:
-            self.proj = nn.Linear(dim, dim, bias=proj_bias)
-        self.proj_drop = nn.Dropout(proj_drop)
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(dim * 3)
-                self.ssf_scale_2, self.ssf_shift_2 = init_ssf_scale_shift(dim)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-        #if not self.training:
-        #
-        # self.attn = ScaledDotProduct()
-            #self.attn = MultiHeadDispatch(dim_model=EMB, residual_dropout=DROPOUT, num_heads=HEADS, attention=attn)
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        B, N, C = x.shape
-        if self.tuning_mode == 'ssf':
-            qkv = ssf_ada(self.qkv(x), self.ssf_scale_1, self.ssf_shift_1).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        else:
-            qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
-        attn = q @ k.transpose(-2, -1)
-        if attn_bias is not None:
-            attn = attn + attn_bias[:, :, :N]
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
-        x = self.proj(x)
-        if self.tuning_mode == 'ssf':
-            x = ssf_ada(x, self.ssf_scale_2, self.ssf_shift_2)
-        x = self.proj_drop(x)
-        return x
-class MemEffAttention(Attention):
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
-        #if True:
-            assert attn_bias is None, "xFormers is required for nested tensors usage"
-            return super().forward(x, attn_bias)
-        B, N, C = x.shape
-        if self.tuning_mode == 'ssf':
-            qkv = ssf_ada(self.qkv(x), self.ssf_scale_1, self.ssf_shift_1).reshape(B, N, 3, self.num_heads, C // self.num_heads)
-        else:
-            qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
-        q, k, v = unbind(qkv, 2)
-        if attn_bias is not None:
-            x = memory_efficient_attention(q, k, v, attn_bias=attn_bias[:, :, :N])
-        else:
-            x = memory_efficient_attention(q, k, v)
-        x = x.reshape([B, N, C])
-        x = self.proj(x)
-        if self.tuning_mode == 'ssf':
-            x = ssf_ada(x, self.ssf_scale_2, self.ssf_shift_2)
-        x = self.proj_drop(x)
-        return x
-try:
-    from xformers.ops import fmha
-    from xformers.ops import scaled_index_add, index_select_cat
-    #import numpy.bool
-    XFORMERS_AVAILABLE = True
-except ImportError:
-    logger.warning("xFormers not available")
-    XFORMERS_AVAILABLE = False
-class Block(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        qkv_bias: bool = False,
-        proj_bias: bool = True,
-        ffn_bias: bool = True,
-        drop: float = 0.0,
-        attn_drop: float = 0.0,
-        init_values = None,
-        drop_path: float = 0.0,
-        act_layer: Callable[..., nn.Module] = nn.GELU,
-        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
-        attn_class: Callable[..., nn.Module] = Attention,
-        ffn_layer: Callable[..., nn.Module] = Mlp,
-        tuning_mode: Optional[int] = None
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.norm1 = norm_layer(dim)
-        self.attn = attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-            tuning_mode=tuning_mode
-        )
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(dim)
-                self.ssf_scale_2, self.ssf_shift_2 = init_ssf_scale_shift(dim)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = ffn_layer(
-            in_features=dim,
-            hidden_features=mlp_hidden_dim,
-            act_layer=act_layer,
-            drop=drop,
-            bias=ffn_bias,
-        )
-        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.sample_drop_ratio = drop_path
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        def attn_residual_func(x: Tensor, attn_bias) -> Tensor:
-            if self.tuning_mode == 'ssf':
-                return self.ls1(self.attn(ssf_ada(self.norm1(x), self.ssf_scale_1, self.ssf_shift_1), attn_bias))
-            else:
-                return self.ls1(self.attn(self.norm1(x), attn_bias))
-        def ffn_residual_func(x: Tensor) -> Tensor:
-            if self.tuning_mode == 'ssf':
-                return self.ls2(self.mlp(ssf_ada(self.norm2(x), self.ssf_scale_2, self.ssf_shift_2)))
-            else:
-                return self.ls2(self.mlp(self.norm2(x)))
-        if self.training and self.sample_drop_ratio > 0.1:
-            # the overhead is compensated only for a drop path rate larger than 0.1
-            x = drop_add_residual_stochastic_depth(
-                x,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                attn_bias=attn_bias
-            )
-            x = drop_add_residual_stochastic_depth(
-                x,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-            )
-        elif self.training and self.sample_drop_ratio > 0.0:
-            x = x + self.drop_path1(attn_residual_func(x, attn_bias))
-            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
-        else:
-            x = x + attn_residual_func(x, attn_bias)
-            x = x + ffn_residual_func(x)
-        return x
-def drop_add_residual_stochastic_depth(
-    x: Tensor,
-    residual_func: Callable[[Tensor], Tensor],
-    sample_drop_ratio: float = 0.0, attn_bias=None
-) -> Tensor:
-    # 1) extract subset using permutation
-    b, n, d = x.shape
-    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
-    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
-    x_subset = x[brange]
-    # 2) apply residual_func to get residual
-    residual = residual_func(x_subset, attn_bias)
-    x_flat = x.flatten(1)
-    residual = residual.flatten(1)
-    residual_scale_factor = b / sample_subset_size
-    # 3) add the residual
-    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
-    return x_plus_residual.view_as(x)
-def get_branges_scales(x, sample_drop_ratio=0.0):
-    b, n, d = x.shape
-    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
-    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
-    residual_scale_factor = b / sample_subset_size
-    return brange, residual_scale_factor
-def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
-    if scaling_vector is None:
-        x_flat = x.flatten(1)
-        residual = residual.flatten(1)
-        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
-    else:
-        x_plus_residual = scaled_index_add(
-            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
-        )
-    return x_plus_residual
-attn_bias_cache: Dict[Tuple, Any] = {}
-def get_attn_bias_and_cat(x_list, branges=None):
-    """
-    this will perform the index select, cat the tensors, and provide the attn_bias from cache
-    """
-    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
-    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
-    if all_shapes not in attn_bias_cache.keys():
-        seqlens = []
-        for b, x in zip(batch_sizes, x_list):
-            for _ in range(b):
-                seqlens.append(x.shape[1])
-        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
-        attn_bias._batch_sizes = batch_sizes
-        attn_bias_cache[all_shapes] = attn_bias
-    if branges is not None:
-        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
-    else:
-        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
-        cat_tensors = torch.cat(tensors_bs1, dim=1)
-    return attn_bias_cache[all_shapes], cat_tensors
-def drop_add_residual_stochastic_depth_list(
-    x_list: List[Tensor],
-    residual_func: Callable[[Tensor, Any], Tensor],
-    sample_drop_ratio: float = 0.0,
-    scaling_vector=None,
-) -> Tensor:
-    # 1) generate random set of indices for dropping samples in the batch
-    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
-    branges = [s[0] for s in branges_scales]
-    residual_scale_factors = [s[1] for s in branges_scales]
-    # 2) get attention bias and index+concat the tensors
-    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
-    # 3) apply residual_func to get residual, and split the result
-    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
-    outputs = []
-    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
-        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
-    return outputs
-class NestedTensorBlock(Block):
-    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
-        """
-        x_list contains a list of tensors to nest together and run
-        """
-        assert isinstance(self.attn, MemEffAttention)
-        if self.training and self.sample_drop_ratio > 0.0:
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.attn(self.norm1(x), attn_bias=attn_bias)
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.mlp(self.norm2(x))
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
-            )
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
-            )
-            return x_list
-        else:
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls2(self.mlp(self.norm2(x)))
-            attn_bias, x = get_attn_bias_and_cat(x_list)
-            x = x + attn_residual_func(x, attn_bias=attn_bias)
-            x = x + ffn_residual_func(x)
-            return attn_bias.split(x)
-    def forward(self, x_or_x_list, attn_bias=None):
-        if isinstance(x_or_x_list, Tensor):
-            return super().forward(x_or_x_list, attn_bias)
-        elif isinstance(x_or_x_list, list):
-            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
-            return self.forward_nested(x_or_x_list)
-        else:
-            raise AssertionError
-def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
-    if not depth_first and include_root:
-        fn(module=module, name=name)
-    for child_name, child_module in module.named_children():
-        child_name = ".".join((name, child_name)) if name else child_name
-        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
-    if depth_first and include_root:
-        fn(module=module, name=name)
-    return module
-class BlockChunk(nn.ModuleList):
-    def forward(self, x, others=None):
-        for b in self:
-            if others == None:
-                x = b(x)
-            else:
-                x = b(x, others)
-        return x
-class DinoVisionTransformer(nn.Module):
-    def __init__(
-        self,
-        img_size=518,
-        patch_size=16,
-        in_chans=3,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        ffn_bias=True,
-        proj_bias=True,
-        drop_path_rate=0.0,
-        drop_path_uniform=False,
-        init_values=1e-5,  # for layerscale: None or 0 => no layerscale
-        embed_layer=PatchEmbed,
-        act_layer=nn.GELU,
-        block_fn=Block,
-        ffn_layer="mlp",
-        block_chunks=1,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        tuning_mode=None,
-        **kwargs
-    ):
-        """
-        Args:
-            img_size (int, tuple): input image size
-            patch_size (int, tuple): patch size
-            in_chans (int): number of input channels
-            embed_dim (int): embedding dimension
-            depth (int): depth of transformer
-            num_heads (int): number of attention heads
-            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
-            qkv_bias (bool): enable bias for qkv if True
-            proj_bias (bool): enable bias for proj in attn if True
-            ffn_bias (bool): enable bias for ffn if True
-            drop_path_rate (float): stochastic depth rate
-            drop_path_uniform (bool): apply uniform drop rate across blocks
-            weight_init (str): weight init scheme
-            init_values (float): layer-scale init values
-            embed_layer (nn.Module): patch embedding layer
-            act_layer (nn.Module): MLP activation layer
-            block_fn (nn.Module): transformer block class
-            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
-            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
-            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
-            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
-            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
-        """
-        super().__init__()
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
-        self.num_tokens = 1
-        self.n_blocks = depth
-        self.num_heads = num_heads
-        self.patch_size = patch_size
-        self.num_register_tokens = num_register_tokens
-        self.interpolate_antialias = interpolate_antialias
-        self.interpolate_offset = interpolate_offset
-        if tuning_mode != None:
-            self.tuning_mode = tuning_mode
-            if tuning_mode == 'ssf':
-                self.ssf_scale_1, self.ssf_shift_1 = init_ssf_scale_shift(embed_dim)
-            else:
-                pass
-                #raise NotImplementedError()
-        else:
-            self.tuning_mode = None
-        tuning_mode_list = [tuning_mode] * depth
-        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, tuning_mode=tuning_mode)
-        num_patches = self.patch_embed.num_patches
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
-        assert num_register_tokens >= 0
-        self.register_tokens = (
-            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
-        )
-        if drop_path_uniform is True:
-            dpr = [drop_path_rate] * depth
-        else:
-            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
-        if ffn_layer == "mlp":
-            logger.info("using MLP layer as FFN")
-            ffn_layer = Mlp
-        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
-            logger.info("using SwiGLU layer as FFN")
-            ffn_layer = SwiGLUFFNFused
-        elif ffn_layer == "identity":
-            logger.info("using Identity layer as FFN")
-            def f(*args, **kwargs):
-                return nn.Identity()
-            ffn_layer = f
-        else:
-            raise NotImplementedError
-        blocks_list = [
-            block_fn(
-                dim=embed_dim,
-                num_heads=num_heads,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                proj_bias=proj_bias,
-                ffn_bias=ffn_bias,
-                drop_path=dpr[i],
-                norm_layer=norm_layer,
-                act_layer=act_layer,
-                ffn_layer=ffn_layer,
-                init_values=init_values,
-                tuning_mode=tuning_mode_list[i]
-            )
-            for i in range(depth)
-        ]
-        if block_chunks > 0:
-            self.chunked_blocks = True
-            chunked_blocks = []
-            chunksize = depth // block_chunks
-            for i in range(0, depth, chunksize):
-                # this is to keep the block index consistent if we chunk the block list
-                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
-            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
-        else:
-            self.chunked_blocks = False
-            self.blocks = nn.ModuleList(blocks_list)
-        self.norm = norm_layer(embed_dim)
-        self.head = nn.Identity()
-        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
-        self.init_weights()
-    def init_weights(self):
-        trunc_normal_(self.pos_embed, std=0.02)
-        nn.init.normal_(self.cls_token, std=1e-6)
-        if self.register_tokens is not None:
-            nn.init.normal_(self.register_tokens, std=1e-6)
-        named_apply(init_weights_vit_timm, self)
-    def interpolate_pos_encoding(self, x, w, h):
-        previous_dtype = x.dtype
-        npatch = x.shape[1] - 1
-        N = self.pos_embed.shape[1] - 1
-        if npatch == N and w == h:
-            return self.pos_embed
-        pos_embed = self.pos_embed.float()
-        class_pos_embed = pos_embed[:, 0]
-        patch_pos_embed = pos_embed[:, 1:]
-        dim = x.shape[-1]
-        w0 = w // self.patch_size
-        h0 = h // self.patch_size
-        # we add a small number to avoid floating point error in the interpolation
-        # see discussion at https://github.com/facebookresearch/dino/issues/8
-        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
-        sqrt_N = math.sqrt(N)
-        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
-            scale_factor=(sx, sy),
-            mode="bicubic",
-            antialias=self.interpolate_antialias,
-        )
-        assert int(w0) == patch_pos_embed.shape[-2]
-        assert int(h0) == patch_pos_embed.shape[-1]
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
-    def prepare_tokens_with_masks(self, x, masks=None):
-        B, nc, w, h = x.shape
-        x = self.patch_embed(x)
-        if masks is not None:
-            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
-        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
-        x = x + self.interpolate_pos_encoding(x, w, h)
-        if self.register_tokens is not None:
-            x = torch.cat(
-                (
-                    x[:, :1],
-                    self.register_tokens.expand(x.shape[0], -1, -1),
-                    x[:, 1:],
-                ),
-                dim=1,
-            )
-        return x
-    def forward_features_list(self, x_list, masks_list):
-        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
-        for blk in self.blocks:
-            x = blk(x)
-        all_x = x
-        output = []
-        for x, masks in zip(all_x, masks_list):
-            x_norm = self.norm(x)
-            output.append(
-                {
-                    "x_norm_clstoken": x_norm[:, 0],
-                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
-                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
-                    "x_prenorm": x,
-                    "masks": masks,
-                }
-            )
-        return output
-    def forward_features(self, x, masks=None):
-        if isinstance(x, list):
-            return self.forward_features_list(x, masks)
-        B, C, H, W = x.size()
-        pad_h = (self.patch_size - H % self.patch_size)
-        pad_w = (self.patch_size - W % self.patch_size)
-        if pad_h == self.patch_size:
-            pad_h = 0
-        if pad_w == self.patch_size:
-            pad_w = 0
-        #x = nn.functional.pad(x, (pad_h//2, pad_h-pad_h//2, pad_w//2, pad_w-pad_w//2))
-        if pad_h + pad_w > 0:
-            x = torch.nn.functional.interpolate(x, (H+pad_h, W+pad_w), mode='bilinear')
-        x = self.prepare_tokens_with_masks(x, masks)
-        for blk in self.blocks:
-            x = blk(x)
-        x_norm = self.norm(x)
-        if self.tuning_mode == 'ssf':
-            x_norm = ssf_ada(x_norm, self.ssf_scale_1, self.ssf_shift_1)
-        # return {
-        #     "x_norm_clstoken": x_norm[:, 0],
-        #     "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
-        #     "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
-        #     "x_prenorm": x,
-        #     "masks": masks,
-        # }
-        features = []
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        features.append(x_norm)
-        return [features, (B, (H+pad_h)//self.patch_size, (W+pad_w)//self.patch_size, H, W, self.num_register_tokens)]
-    def _get_intermediate_layers_not_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        # If n is an int, take the n last blocks. If it's a list, take them
-        output, total_block_len = [], len(self.blocks)
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for i, blk in enumerate(self.blocks):
-            x = blk(x)
-            if i in blocks_to_take:
-                output.append(x)
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def _get_intermediate_layers_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        output, i, total_block_len = [], 0, len(self.blocks[-1])
-        # If n is an int, take the n last blocks. If it's a list, take them
-        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        for block_chunk in self.blocks:
-            for blk in block_chunk[i:]:  # Passing the nn.Identity()
-                x = blk(x)
-                if i in blocks_to_take:
-                    output.append(x)
-                i += 1
-        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        return output
-    def get_intermediate_layers(
-        self,
-        x: torch.Tensor,
-        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
-        reshape: bool = False,
-        return_class_token: bool = False,
-        norm=True,
-    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
-        if self.chunked_blocks:
-            outputs = self._get_intermediate_layers_chunked(x, n)
-        else:
-            outputs = self._get_intermediate_layers_not_chunked(x, n)
-        if norm:
-            outputs = [self.norm(out) for out in outputs]
-        class_tokens = [out[:, 0] for out in outputs]
-        outputs = [out[:, 1:] for out in outputs]
-        if reshape:
-            B, _, w, h = x.shape
-            outputs = [
-                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
-                for out in outputs
-            ]
-        if return_class_token:
-            return tuple(zip(outputs, class_tokens))
-        return tuple(outputs)
-    def forward(self, *args, is_training=False, **kwargs):
-        ret = self.forward_features(*args, **kwargs)
-        return ret
-        # if is_training:
-        #     return ret
-        # else:
-        #     return self.head(ret["x_norm_clstoken"])
-def init_weights_vit_timm(module: nn.Module, name: str = ""):
-    """ViT weight initialization, original timm impl (for reproducibility)"""
-    if isinstance(module, nn.Linear):
-        trunc_normal_(module.weight, std=0.02)
-        if module.bias is not None:
-            nn.init.zeros_(module.bias)
-def load_ckpt_dino(checkpoint, model):
-    if checkpoint is not None:
-        try:
-            with open(checkpoint, "rb") as f:
-                state_dict = torch.load(f)
-        except:
-            print('NO pretrained imagenet ckpt available! Check your path!')
-            del model.mask_token
-            return
-        try:
-            model.load_state_dict(state_dict, strict=True)
-        except:
-            new_state_dict = {}
-            for key, value in state_dict.items():
-                if 'blocks' in key:
-                    key_new = 'blocks.0' + key[len('blocks'):]
-                else:
-                    key_new = key
-                new_state_dict[key_new] = value
-            model.load_state_dict(new_state_dict, strict=True)
-        del model.mask_token
-        return
-    else:
-        return
-def vit_small(patch_size=14, num_register_tokens=0, checkpoint=None, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=384,
-        depth=12,
-        num_heads=6,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        **kwargs,
-    )
-    load_ckpt_dino(checkpoint, model)
-    return model
-def vit_base(patch_size=14, num_register_tokens=0, checkpoint=None, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        **kwargs,
-    )
-    return model
-def vit_large(patch_size=14, num_register_tokens=0, checkpoint=None, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        **kwargs,
-    )
-    if checkpoint is not None:
-        with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
-        try:
-            model.load_state_dict(state_dict, strict=True)
-        except:
-            new_state_dict = {}
-            for key, value in state_dict.items():
-                if 'blocks' in key:
-                    key_new = 'blocks.0' + key[len('blocks'):]
-                else:
-                    key_new = key
-                new_state_dict[key_new] = value
-            model.load_state_dict(new_state_dict, strict=True)
-        del model.mask_token
-    return model
-def vit_giant2(patch_size=14, num_register_tokens=0, checkpoint=None, **kwargs):
-    """
-    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
-    """
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=1536,
-        depth=40,
-        num_heads=24,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        ffn_layer='swiglu',
-        **kwargs,
-    )
-    return model
-def vit_small_reg(patch_size=14, num_register_tokens=4, checkpoint=None, tuning_mode=None, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=384,
-        depth=12,
-        num_heads=6,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        tuning_mode=tuning_mode,
-        **kwargs,
-    )
-    load_ckpt_dino(checkpoint, model)
-    return model
-def vit_base_reg(patch_size=14, num_register_tokens=4, checkpoint=None, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        **kwargs,
-    )
-    load_ckpt_dino(checkpoint, model)
-    return model
-def vit_large_reg(patch_size=14, num_register_tokens=4, checkpoint=None, tuning_mode=None, **kwargs):
-    model = DinoVisionTransformer(
-        img_size = 518,
-        patch_size=patch_size,
-        embed_dim=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        tuning_mode=tuning_mode,
-        **kwargs,
-    )
-    load_ckpt_dino(checkpoint, model)
-    return model
-def vit_giant2_reg(patch_size=14, num_register_tokens=4, checkpoint=None, tuning_mode=None, **kwargs):
-    """
-    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
-    """
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=1536,
-        depth=40,
-        num_heads=24,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        ffn_layer='swiglu',
-        tuning_mode=tuning_mode,
-        **kwargs,
-    )
-    load_ckpt_dino(checkpoint, model)
-    return model
-if __name__ == '__main__':
-    try:
-        from mmcv.utils import Config
-    except:
-        from mmengine import Config
-    #rgb = torch.rand((2, 3, 518, 518)).cuda()
-    #cfg.data_basic['crop_size']['0']
-    #cfg.data_basic['crop_size']['1']
-    cfg = Config.fromfile('/opt/ml/project/mu.hu/projects/monodepth_vit/mono/configs/RAFTDecoder/vit.raft5.large.kitti.py')
-    #rgb = torch.arange(0, 2*3*1036*1036, 1).cuda().float().view(2, 3, 1036, 1036)
-    rgb = torch.zeros(1, 3, 616, 1064).cuda()
-    cfg['tuning_mode'] = 'ssf'
-    #model = vit_large_reg(checkpoint="/cpfs02/shared/public/groups/local_map/yvan/pretrained_weight_repo/vit/dinov2_vitl14_reg4_pretrain.pth", kwarg=cfg).cuda()
-    model = vit_large_reg(tuning_mode='ssf').cuda()
-    #import timm
-    #model2 = timm.models.vision_transformer.vit_large_patch14_dinov2().cuda()
-    #timm.models.load_checkpoint(model2, '/cpfs02/shared/public/yvan/pretrained_weight_repo/vit/dinov2_vitl14_pretrain.pth', filter_fn=timm.models.vision_transformer.checkpoint_filter_fn)
-    out1 = model(rgb)
-    #out2 = model2(rgb)
-    temp = 0

mono/model/backbones/__init__.py DELETED Viewed

@@ -1,11 +0,0 @@
-from .ConvNeXt import convnext_xlarge
-from .ConvNeXt import convnext_small
-from .ConvNeXt import convnext_base
-from .ConvNeXt import convnext_large
-from .ConvNeXt import convnext_tiny
-from .ViT_DINO import vit_large
-from .ViT_DINO_reg import vit_small_reg, vit_large_reg
-__all__ = [
-    'convnext_xlarge', 'convnext_small', 'convnext_base', 'convnext_large', 'convnext_tiny', 'vit_small_reg', 'vit_large_reg'
-]

mono/model/backbones/__pycache__/ConvNeXt.cpython-39.pyc DELETED Viewed

Binary file (9.37 kB)

mono/model/backbones/__pycache__/__init__.cpython-39.pyc DELETED Viewed

Binary file (410 Bytes)

mono/model/decode_heads/HourGlassDecoder.py DELETED Viewed

@@ -1,274 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-import math
-import torch.nn.functional as F
-def compute_depth_expectation(prob, depth_values):
-    depth_values = depth_values.view(*depth_values.shape, 1, 1)
-    depth = torch.sum(prob * depth_values, 1)
-    return depth
-class ConvBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=3):
-        super(ConvBlock, self).__init__()
-        if kernel_size == 3:
-            self.conv = nn.Sequential(
-                nn.ReflectionPad2d(1),
-                nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
-            )
-        elif kernel_size == 1:
-            self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)
-        self.nonlin = nn.ELU(inplace=True)
-    def forward(self, x):
-        out = self.conv(x)
-        out = self.nonlin(out)
-        return out
-class ConvBlock_double(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=3):
-        super(ConvBlock_double, self).__init__()
-        if kernel_size == 3:
-            self.conv = nn.Sequential(
-                nn.ReflectionPad2d(1),
-                nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
-            )
-        elif kernel_size == 1:
-            self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)
-        self.nonlin = nn.ELU(inplace=True)
-        self.conv_2 = nn.Conv2d(out_channels, out_channels, 1, padding=0, stride=1)
-        self.nonlin_2 =nn.ELU(inplace=True)
-    def forward(self, x):
-        out = self.conv(x)
-        out = self.nonlin(out)
-        out = self.conv_2(out)
-        out = self.nonlin_2(out)
-        return out
-class DecoderFeature(nn.Module):
-    def __init__(self, feat_channels, num_ch_dec=[64, 64, 128, 256]):
-        super(DecoderFeature, self).__init__()
-        self.num_ch_dec = num_ch_dec
-        self.feat_channels = feat_channels
-        self.upconv_3_0 = ConvBlock(self.feat_channels[3], self.num_ch_dec[3], kernel_size=1)
-        self.upconv_3_1 = ConvBlock_double(
-            self.feat_channels[2] + self.num_ch_dec[3],
-            self.num_ch_dec[3],
-            kernel_size=1)
-        self.upconv_2_0 = ConvBlock(self.num_ch_dec[3], self.num_ch_dec[2], kernel_size=3)
-        self.upconv_2_1 = ConvBlock_double(
-            self.feat_channels[1] + self.num_ch_dec[2],
-            self.num_ch_dec[2],
-            kernel_size=3)
-        self.upconv_1_0 = ConvBlock(self.num_ch_dec[2], self.num_ch_dec[1], kernel_size=3)
-        self.upconv_1_1 = ConvBlock_double(
-            self.feat_channels[0] + self.num_ch_dec[1],
-            self.num_ch_dec[1],
-            kernel_size=3)
-        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
-    def forward(self, ref_feature):
-        x = ref_feature[3]
-        x = self.upconv_3_0(x)
-        x = torch.cat((self.upsample(x), ref_feature[2]), 1)
-        x = self.upconv_3_1(x)
-        x = self.upconv_2_0(x)
-        x = torch.cat((self.upsample(x), ref_feature[1]), 1)
-        x = self.upconv_2_1(x)
-        x = self.upconv_1_0(x)
-        x = torch.cat((self.upsample(x), ref_feature[0]), 1)
-        x = self.upconv_1_1(x)
-        return x
-class UNet(nn.Module):
-    def __init__(self, inp_ch=32, output_chal=1, down_sample_times=3, channel_mode='v0'):
-        super(UNet, self).__init__()
-        basic_block = ConvBnReLU
-        num_depth = 128
-        self.conv0 = basic_block(inp_ch, num_depth)
-        if channel_mode == 'v0':
-            channels = [num_depth, num_depth//2, num_depth//4, num_depth//8, num_depth // 8]
-        elif channel_mode == 'v1':
-            channels = [num_depth, num_depth, num_depth, num_depth, num_depth, num_depth]
-        self.down_sample_times = down_sample_times
-        for i in range(down_sample_times):
-            setattr(
-                self, 'conv_%d' % i,
-                nn.Sequential(
-                    basic_block(channels[i], channels[i+1], stride=2),
-                    basic_block(channels[i+1], channels[i+1])
-                )
-            )
-        for i in range(down_sample_times-1,-1,-1):
-            setattr(self, 'deconv_%d' % i,
-                    nn.Sequential(
-                        nn.ConvTranspose2d(
-                            channels[i+1],
-                            channels[i],
-                            kernel_size=3,
-                            padding=1,
-                            output_padding=1,
-                            stride=2,
-                            bias=False),
-                        nn.BatchNorm2d(channels[i]),
-                        nn.ReLU(inplace=True)
-                    )
-                )
-            self.prob = nn.Conv2d(num_depth, output_chal, 1, stride=1, padding=0)
-    def forward(self, x):
-        features = {}
-        conv0 = self.conv0(x)
-        x = conv0
-        features[0] = conv0
-        for i in range(self.down_sample_times):
-            x = getattr(self, 'conv_%d' % i)(x)
-            features[i+1] = x
-        for i in range(self.down_sample_times-1,-1,-1):
-            x = features[i] + getattr(self, 'deconv_%d' % i)(x)
-        x = self.prob(x)
-        return x
-class ConvBnReLU(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
-        super(ConvBnReLU, self).__init__()
-        self.conv = nn.Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=pad,
-            bias=False
-        )
-        self.bn = nn.BatchNorm2d(out_channels)
-    def forward(self, x):
-        return F.relu(self.bn(self.conv(x)), inplace=True)
-class HourglassDecoder(nn.Module):
-    def __init__(self, cfg):
-        super(HourglassDecoder, self).__init__()
-        self.inchannels = cfg.model.decode_head.in_channels #  [256, 512, 1024, 2048]
-        self.decoder_channels = cfg.model.decode_head.decoder_channel # [64, 64, 128, 256]
-        self.min_val = cfg.data_basic.depth_normalize[0]
-        self.max_val = cfg.data_basic.depth_normalize[1]
-        self.num_ch_dec = self.decoder_channels # [64, 64, 128, 256]
-        self.num_depth_regressor_anchor = 512
-        self.feat_channels = self.inchannels
-        unet_in_channel = self.num_ch_dec[1]
-        unet_out_channel = 256
-        self.decoder_mono = DecoderFeature(self.feat_channels, self.num_ch_dec)
-        self.conv_out_2 = UNet(inp_ch=unet_in_channel,
-                               output_chal=unet_out_channel + 1,
-                               down_sample_times=3,
-                               channel_mode='v0',
-                               )
-        self.depth_regressor_2 = nn.Sequential(
-            nn.Conv2d(unet_out_channel,
-                      self.num_depth_regressor_anchor,
-                      kernel_size=3,
-                      padding=1,
-                ),
-            nn.BatchNorm2d(self.num_depth_regressor_anchor),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(
-                self.num_depth_regressor_anchor,
-                self.num_depth_regressor_anchor,
-                kernel_size=1,
-            )
-        )
-        self.residual_channel = 16
-        self.conv_up_2 = nn.Sequential(
-            nn.Conv2d(1 + 2 + unet_out_channel, self.residual_channel, 3, padding=1),
-            nn.BatchNorm2d(self.residual_channel),
-            nn.ReLU(),
-            nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
-            nn.Upsample(scale_factor=4),
-            nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(self.residual_channel, 1, 1, padding=0),
-        )
-    def get_bins(self, bins_num):
-        depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device='cuda')
-        depth_bins_vec = torch.exp(depth_bins_vec)
-        return depth_bins_vec
-    def register_depth_expectation_anchor(self, bins_num, B):
-        depth_bins_vec = self.get_bins(bins_num)
-        depth_bins_vec = depth_bins_vec.unsqueeze(0).repeat(B, 1)
-        self.register_buffer('depth_expectation_anchor', depth_bins_vec, persistent=False)
-    def upsample(self, x, scale_factor=2):
-        return F.interpolate(x, scale_factor=scale_factor, mode='nearest')
-    def regress_depth_2(self, feature_map_d):
-        prob = self.depth_regressor_2(feature_map_d).softmax(dim=1)
-        B = prob.shape[0]
-        if "depth_expectation_anchor" not in self._buffers:
-            self.register_depth_expectation_anchor(self.num_depth_regressor_anchor, B)
-        d = compute_depth_expectation(
-            prob,
-            self.depth_expectation_anchor[:B, ...]
-        ).unsqueeze(1)
-        return d
-    def create_mesh_grid(self, height, width, batch, device="cuda", set_buffer=True):
-        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=device),
-                               torch.arange(0, width, dtype=torch.float32, device=device)], indexing='ij')
-        meshgrid = torch.stack((x, y))
-        meshgrid = meshgrid.unsqueeze(0).repeat(batch, 1, 1, 1)
-        return meshgrid
-    def forward(self, features_mono, **kwargs):
-        '''
-        trans_ref2src: list of transformation matrix from the reference view to source view. [B, 4, 4]
-        inv_intrinsic_pool: list of inverse intrinsic matrix.
-        features_mono: features of reference and source views. [[ref_f1, ref_f2, ref_f3, ref_f4],[src1_f1, src1_f2, src1_f3, src1_f4], ...].
-        '''
-        outputs = {}
-        # get encoder feature of the reference view
-        ref_feat = features_mono
-        feature_map_mono = self.decoder_mono(ref_feat)
-        feature_map_mono_pred = self.conv_out_2(feature_map_mono)
-        confidence_map_2 = feature_map_mono_pred[:, -1:, :, :]
-        feature_map_d_2 = feature_map_mono_pred[:, :-1, :, :]
-        depth_pred_2 = self.regress_depth_2(feature_map_d_2)
-        B, _, H, W = depth_pred_2.shape
-        meshgrid = self.create_mesh_grid(H, W, B)
-        depth_pred_mono = self.upsample(depth_pred_2, scale_factor=4) + 1e-1 * \
-            self.conv_up_2(
-                torch.cat((depth_pred_2, meshgrid[:B, ...], feature_map_d_2), 1)
-            )
-        confidence_map_mono = self.upsample(confidence_map_2, scale_factor=4)
-        outputs=dict(
-            prediction=depth_pred_mono,
-            confidence=confidence_map_mono,
-            pred_logit=None,
-        )
-        return outputs

mono/model/decode_heads/RAFTDepthNormalDPTDecoder5.py DELETED Viewed

@@ -1,1033 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-import math
-import torch.nn.functional as F
-# LORA finetuning originally by edwardjhu
-class LoRALayer():
-    def __init__(
-        self,
-        r: int,
-        lora_alpha: int,
-        lora_dropout: float,
-        merge_weights: bool,
-    ):
-        self.r = r
-        self.lora_alpha = lora_alpha
-        # Optional dropout
-        if lora_dropout > 0.:
-            self.lora_dropout = nn.Dropout(p=lora_dropout)
-        else:
-            self.lora_dropout = lambda x: x
-        # Mark the weight as unmerged
-        self.merged = False
-        self.merge_weights = merge_weights
-class LoRALinear(nn.Linear, LoRALayer):
-    # LoRA implemented in a dense layer
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        r: int = 0,
-        lora_alpha: int = 1,
-        lora_dropout: float = 0.,
-        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
-        merge_weights: bool = True,
-        **kwargs
-    ):
-        nn.Linear.__init__(self, in_features, out_features, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
-                           merge_weights=merge_weights)
-        self.fan_in_fan_out = fan_in_fan_out
-        # Actual trainable parameters
-        if r > 0:
-            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
-            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
-            self.scaling = self.lora_alpha / self.r
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False
-        self.reset_parameters()
-        if fan_in_fan_out:
-            self.weight.data = self.weight.data.transpose(0, 1)
-    def reset_parameters(self):
-        #nn.Linear.reset_parameters(self)
-        if hasattr(self, 'lora_A'):
-            # initialize B the same way as the default for nn.Linear and A to zero
-            # this is different than what is described in the paper but should not affect performance
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-    # def train(self, mode: bool = True):
-    #     def T(w):
-    #         return w.transpose(0, 1) if self.fan_in_fan_out else w
-    #     nn.Linear.train(self, mode)
-    #     if mode:
-    #         if self.merge_weights and self.merged:
-    #             # Make sure that the weights are not merged
-    #             if self.r > 0:
-    #                 self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
-    #             self.merged = False
-    #     else:
-    #         if self.merge_weights and not self.merged:
-    #             # Merge the weights and mark it
-    #             if self.r > 0:
-    #                 self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
-    #             self.merged = True
-    def forward(self, x: torch.Tensor):
-        def T(w):
-            return w.transpose(0, 1) if self.fan_in_fan_out else w
-        if self.r > 0 and not self.merged:
-            result = F.linear(x, T(self.weight), bias=self.bias)
-            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
-            return result
-        else:
-            return F.linear(x, T(self.weight), bias=self.bias)
-class ConvLoRA(nn.Conv2d, LoRALayer):
-    def __init__(self, in_channels, out_channels, kernel_size, r=0, lora_alpha=1, lora_dropout=0., merge_weights=True, **kwargs):
-        #self.conv = conv_module(in_channels, out_channels, kernel_size, **kwargs)
-        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
-        assert isinstance(kernel_size, int)
-        # Actual trainable parameters
-        if r > 0:
-            self.lora_A = nn.Parameter(
-                self.weight.new_zeros((r * kernel_size, in_channels * kernel_size))
-            )
-            self.lora_B = nn.Parameter(
-              self.weight.new_zeros((out_channels//self.groups*kernel_size, r*kernel_size))
-            )
-            self.scaling = self.lora_alpha / self.r
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False
-        self.reset_parameters()
-        self.merged = False
-    def reset_parameters(self):
-        #self.conv.reset_parameters()
-        if hasattr(self, 'lora_A'):
-            # initialize A the same way as the default for nn.Linear and B to zero
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-    # def train(self, mode=True):
-    #     super(ConvLoRA, self).train(mode)
-    #     if mode:
-    #         if self.merge_weights and self.merged:
-    #             if self.r > 0:
-    #                 # Make sure that the weights are not merged
-    #                 self.conv.weight.data -= (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
-    #             self.merged = False
-    #     else:
-    #         if self.merge_weights and not self.merged:
-    #             if self.r > 0:
-    #                 # Merge the weights and mark it
-    #                 self.conv.weight.data += (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
-    #             self.merged = True
-    def forward(self, x):
-        if self.r > 0 and not self.merged:
-            # return self.conv._conv_forward(
-            #     x,
-            #     self.conv.weight + (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling,
-            #     self.conv.bias
-            # )
-            weight = self.weight + (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling
-            bias = self.bias
-            return F.conv2d(x, weight, bias=bias, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
-        else:
-            return F.conv2d(x, self.weight, bias=self.bias, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
-class ConvTransposeLoRA(nn.ConvTranspose2d, LoRALayer):
-    def __init__(self, in_channels, out_channels, kernel_size, r=0, lora_alpha=1, lora_dropout=0., merge_weights=True, **kwargs):
-        #self.conv = conv_module(in_channels, out_channels, kernel_size, **kwargs)
-        nn.ConvTranspose2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
-        assert isinstance(kernel_size, int)
-        # Actual trainable parameters
-        if r > 0:
-            self.lora_A = nn.Parameter(
-                self.weight.new_zeros((r * kernel_size, in_channels * kernel_size))
-            )
-            self.lora_B = nn.Parameter(
-              self.weight.new_zeros((out_channels//self.groups*kernel_size, r*kernel_size))
-            )
-            self.scaling = self.lora_alpha / self.r
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False
-        self.reset_parameters()
-        self.merged = False
-    def reset_parameters(self):
-        #self.conv.reset_parameters()
-        if hasattr(self, 'lora_A'):
-            # initialize A the same way as the default for nn.Linear and B to zero
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-    # def train(self, mode=True):
-    #     super(ConvTransposeLoRA, self).train(mode)
-    #     if mode:
-    #         if self.merge_weights and self.merged:
-    #             if self.r > 0:
-    #                 # Make sure that the weights are not merged
-    #                 self.conv.weight.data -= (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
-    #             self.merged = False
-    #     else:
-    #         if self.merge_weights and not self.merged:
-    #             if self.r > 0:
-    #                 # Merge the weights and mark it
-    #                 self.conv.weight.data += (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
-    #             self.merged = True
-    def forward(self, x):
-        if self.r > 0 and not self.merged:
-            weight = self.weight + (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling
-            bias = self.bias
-            return F.conv_transpose2d(x, weight,
-                bias=bias, stride=self.stride, padding=self.padding, output_padding=self.output_padding,
-                groups=self.groups, dilation=self.dilation)
-        else:
-            return F.conv_transpose2d(x, self.weight,
-                bias=self.bias, stride=self.stride, padding=self.padding, output_padding=self.output_padding,
-                groups=self.groups, dilation=self.dilation)
-        #return self.conv(x)
-class Conv2dLoRA(ConvLoRA):
-    def __init__(self, *args, **kwargs):
-        super(Conv2dLoRA, self).__init__(*args, **kwargs)
-class ConvTranspose2dLoRA(ConvTransposeLoRA):
-    def __init__(self, *args, **kwargs):
-        super(ConvTranspose2dLoRA, self).__init__(*args, **kwargs)
-def compute_depth_expectation(prob, depth_values):
-    depth_values = depth_values.view(*depth_values.shape, 1, 1)
-    depth = torch.sum(prob * depth_values, 1)
-    return depth
-def interpolate_float32(x, size=None, scale_factor=None, mode='nearest', align_corners=None):
-    with torch.autocast(device_type='cuda', dtype=torch.bfloat16, enabled=False):
-        return F.interpolate(x.float(), size=size, scale_factor=scale_factor, mode=mode, align_corners=align_corners)
-# def upflow8(flow, mode='bilinear'):
-#     new_size = (8 * flow.shape[2], 8 * flow.shape[3])
-#     return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
-def upflow4(flow, mode='bilinear'):
-    new_size = (4 * flow.shape[2], 4 * flow.shape[3])
-    with torch.autocast(device_type='cuda', dtype=torch.bfloat16, enabled=False):
-        return  F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
-def coords_grid(batch, ht, wd):
-    # coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
-    coords = (torch.zeros((ht, wd)), torch.zeros((ht, wd)), torch.zeros((ht, wd)), torch.zeros((ht, wd)), torch.zeros((ht, wd)), torch.zeros((ht, wd)))
-    coords = torch.stack(coords[::-1], dim=0).float()
-    return coords[None].repeat(batch, 1, 1, 1)
-def norm_normalize(norm_out):
-    min_kappa = 0.01
-    norm_x, norm_y, norm_z, kappa = torch.split(norm_out, 1, dim=1)
-    norm = torch.sqrt(norm_x ** 2.0 + norm_y ** 2.0 + norm_z ** 2.0) + 1e-10
-    kappa = F.elu(kappa) + 1.0 + min_kappa
-    final_out = torch.cat([norm_x / norm, norm_y / norm, norm_z / norm, kappa], dim=1)
-    return final_out
-# uncertainty-guided sampling (only used during training)
-@torch.no_grad()
-def sample_points(init_normal, gt_norm_mask, sampling_ratio, beta):
-    device = init_normal.device
-    B, _, H, W = init_normal.shape
-    N = int(sampling_ratio * H * W)
-    beta = beta
-    # uncertainty map
-    uncertainty_map = -1 * init_normal[:, -1, :, :]  # B, H, W
-    # gt_invalid_mask (B, H, W)
-    if gt_norm_mask is not None:
-        gt_invalid_mask = F.interpolate(gt_norm_mask.float(), size=[H, W], mode='nearest')
-        gt_invalid_mask = gt_invalid_mask[:, 0, :, :] < 0.5
-        uncertainty_map[gt_invalid_mask] = -1e4
-    # (B, H*W)
-    _, idx = uncertainty_map.view(B, -1).sort(1, descending=True)
-    # importance sampling
-    if int(beta * N) > 0:
-        importance = idx[:, :int(beta * N)]    # B, beta*N
-        # remaining
-        remaining = idx[:, int(beta * N):]     # B, H*W - beta*N
-        # coverage
-        num_coverage = N - int(beta * N)
-        if num_coverage <= 0:
-            samples = importance
-        else:
-            coverage_list = []
-            for i in range(B):
-                idx_c = torch.randperm(remaining.size()[1])    # shuffles "H*W - beta*N"
-                coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))     # 1, N-beta*N
-            coverage = torch.cat(coverage_list, dim=0)                                      # B, N-beta*N
-            samples = torch.cat((importance, coverage), dim=1)                              # B, N
-    else:
-        # remaining
-        remaining = idx[:, :]  # B, H*W
-        # coverage
-        num_coverage = N
-        coverage_list = []
-        for i in range(B):
-            idx_c = torch.randperm(remaining.size()[1])  # shuffles "H*W - beta*N"
-            coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))  # 1, N-beta*N
-        coverage = torch.cat(coverage_list, dim=0)  # B, N-beta*N
-        samples = coverage
-    # point coordinates
-    rows_int = samples // W         # 0 for first row, H-1 for last row
-    rows_float = rows_int / float(H-1)         # 0 to 1.0
-    rows_float = (rows_float * 2.0) - 1.0       # -1.0 to 1.0
-    cols_int = samples % W          # 0 for first column, W-1 for last column
-    cols_float = cols_int / float(W-1)         # 0 to 1.0
-    cols_float = (cols_float * 2.0) - 1.0       # -1.0 to 1.0
-    point_coords = torch.zeros(B, 1, N, 2)
-    point_coords[:, 0, :, 0] = cols_float             # x coord
-    point_coords[:, 0, :, 1] = rows_float             # y coord
-    point_coords = point_coords.to(device)
-    return point_coords, rows_int, cols_int
-class FlowHead(nn.Module):
-    def __init__(self, input_dim=128, hidden_dim=256, output_dim_depth=2, output_dim_norm=4, tuning_mode=None):
-        super(FlowHead, self).__init__()
-        self.conv1d = Conv2dLoRA(input_dim, hidden_dim // 2, 3, padding=1, r = 8 if tuning_mode == 'lora' else 0)
-        self.conv2d = Conv2dLoRA(hidden_dim // 2, output_dim_depth, 3, padding=1, r = 8 if tuning_mode == 'lora' else 0)
-        self.conv1n = Conv2dLoRA(input_dim, hidden_dim // 2, 3, padding=1, r = 8 if tuning_mode == 'lora' else 0)
-        self.conv2n = Conv2dLoRA(hidden_dim // 2, output_dim_norm, 3, padding=1, r = 8 if tuning_mode == 'lora' else 0)
-        self.relu = nn.ReLU(inplace=True)
-    def forward(self, x):
-        depth = self.conv2d(self.relu(self.conv1d(x)))
-        normal = self.conv2n(self.relu(self.conv1n(x)))
-        return torch.cat((depth, normal), dim=1)
-class ConvGRU(nn.Module):
-    def __init__(self, hidden_dim, input_dim, kernel_size=3, tuning_mode=None):
-        super(ConvGRU, self).__init__()
-        self.convz = Conv2dLoRA(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2, r = 8 if tuning_mode == 'lora' else 0)
-        self.convr = Conv2dLoRA(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2, r = 8 if tuning_mode == 'lora' else 0)
-        self.convq = Conv2dLoRA(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2, r = 8 if tuning_mode == 'lora' else 0)
-    def forward(self, h, cz, cr, cq, *x_list):
-        x = torch.cat(x_list, dim=1)
-        hx = torch.cat([h, x], dim=1)
-        z = torch.sigmoid((self.convz(hx) + cz))
-        r = torch.sigmoid((self.convr(hx) + cr))
-        q = torch.tanh((self.convq(torch.cat([r*h, x], dim=1)) + cq))
-        # z = torch.sigmoid((self.convz(hx) + cz).float())
-        # r = torch.sigmoid((self.convr(hx) + cr).float())
-        # q = torch.tanh((self.convq(torch.cat([r*h, x], dim=1)) + cq).float())
-        h = (1-z) * h + z * q
-        return h
-def pool2x(x):
-    return F.avg_pool2d(x, 3, stride=2, padding=1)
-def pool4x(x):
-    return F.avg_pool2d(x, 5, stride=4, padding=1)
-def interp(x, dest):
-    interp_args = {'mode': 'bilinear', 'align_corners': True}
-    return interpolate_float32(x, dest.shape[2:], **interp_args)
-class BasicMultiUpdateBlock(nn.Module):
-    def __init__(self, args, hidden_dims=[], out_dims=2, tuning_mode=None):
-        super().__init__()
-        self.args = args
-        self.n_gru_layers = args.model.decode_head.n_gru_layers # 3
-        self.n_downsample = args.model.decode_head.n_downsample # 3, resolution of the disparity field (1/2^K)
-        # self.encoder = BasicMotionEncoder(args)
-        # encoder_output_dim = 128 # if there is corr volume
-        encoder_output_dim = 6 # no corr volume
-        self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (self.n_gru_layers > 1), tuning_mode=tuning_mode)
-        self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (self.n_gru_layers == 3) + hidden_dims[2], tuning_mode=tuning_mode)
-        self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1], tuning_mode=tuning_mode)
-        self.flow_head = FlowHead(hidden_dims[2], hidden_dim=2*hidden_dims[2], tuning_mode=tuning_mode)
-        factor = 2**self.n_downsample
-        self.mask = nn.Sequential(
-            Conv2dLoRA(hidden_dims[2], hidden_dims[2], 3, padding=1, r = 8 if tuning_mode == 'lora' else 0),
-            nn.ReLU(inplace=True),
-            Conv2dLoRA(hidden_dims[2], (factor**2)*9, 1, padding=0, r = 8 if tuning_mode == 'lora' else 0))
-    def forward(self, net, inp, corr=None, flow=None, iter08=True, iter16=True, iter32=True, update=True):
-        if iter32:
-            net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1]))
-        if iter16:
-            if self.n_gru_layers > 2:
-                net[1] = self.gru16(net[1], *(inp[1]), interp(pool2x(net[0]), net[1]), interp(net[2], net[1]))
-            else:
-                net[1] = self.gru16(net[1], *(inp[1]), interp(pool2x(net[0]), net[1]))
-        if iter08:
-            if corr is not None:
-                motion_features = self.encoder(flow, corr)
-            else:
-                motion_features = flow
-            if self.n_gru_layers > 1:
-                net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0]))
-            else:
-                net[0] = self.gru08(net[0], *(inp[0]), motion_features)
-        if not update:
-            return net
-        delta_flow = self.flow_head(net[0])
-        # scale mask to balence gradients
-        mask = .25 * self.mask(net[0])
-        return net, mask, delta_flow
-class LayerNorm2d(nn.LayerNorm):
-    def __init__(self, dim):
-        super(LayerNorm2d, self).__init__(dim)
-    def forward(self, x):
-        x = x.permute(0, 2, 3, 1).contiguous()
-        x = super(LayerNorm2d, self).forward(x)
-        x = x.permute(0, 3, 1, 2).contiguous()
-        return x
-class ResidualBlock(nn.Module):
-    def __init__(self, in_planes, planes, norm_fn='group', stride=1, tuning_mode=None):
-        super(ResidualBlock, self).__init__()
-        self.conv1 = Conv2dLoRA(in_planes, planes, kernel_size=3, padding=1, stride=stride, r = 8 if tuning_mode == 'lora' else 0)
-        self.conv2 = Conv2dLoRA(planes, planes, kernel_size=3, padding=1, r = 8 if tuning_mode == 'lora' else 0)
-        self.relu = nn.ReLU(inplace=True)
-        num_groups = planes // 8
-        if norm_fn == 'group':
-            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-            if not (stride == 1 and in_planes == planes):
-                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-        elif norm_fn == 'batch':
-            self.norm1 = nn.BatchNorm2d(planes)
-            self.norm2 = nn.BatchNorm2d(planes)
-            if not (stride == 1 and in_planes == planes):
-                self.norm3 = nn.BatchNorm2d(planes)
-        elif norm_fn == 'instance':
-            self.norm1 = nn.InstanceNorm2d(planes)
-            self.norm2 = nn.InstanceNorm2d(planes)
-            if not (stride == 1 and in_planes == planes):
-                self.norm3 = nn.InstanceNorm2d(planes)
-        elif norm_fn == 'layer':
-            self.norm1 = LayerNorm2d(planes)
-            self.norm2 = LayerNorm2d(planes)
-            if not (stride == 1 and in_planes == planes):
-                self.norm3 = LayerNorm2d(planes)
-        elif norm_fn == 'none':
-            self.norm1 = nn.Sequential()
-            self.norm2 = nn.Sequential()
-            if not (stride == 1 and in_planes == planes):
-                self.norm3 = nn.Sequential()
-        if stride == 1 and in_planes == planes:
-            self.downsample = None
-        else:
-            self.downsample = nn.Sequential(
-                Conv2dLoRA(in_planes, planes, kernel_size=1, stride=stride,  r = 8 if tuning_mode == 'lora' else 0), self.norm3)
-    def forward(self, x):
-        y = x
-        y = self.conv1(y)
-        y = self.norm1(y)
-        y = self.relu(y)
-        y = self.conv2(y)
-        y = self.norm2(y)
-        y = self.relu(y)
-        if self.downsample is not None:
-            x = self.downsample(x)
-        return self.relu(x+y)
-class ContextFeatureEncoder(nn.Module):
-    '''
-    Encoder features are used to:
-        1. initialize the hidden state of the update operator
-        2. and also injected into the GRU during each iteration of the update operator
-    '''
-    def __init__(self, in_dim, output_dim, tuning_mode=None):
-        '''
-        in_dim     = [x4, x8, x16, x32]
-        output_dim = [hindden_dims,   context_dims]
-                    [[x4,x8,x16,x32],[x4,x8,x16,x32]]
-        '''
-        super().__init__()
-        output_list = []
-        for dim in output_dim:
-            conv_out = nn.Sequential(
-                ResidualBlock(in_dim[0], dim[0], 'layer', stride=1, tuning_mode=tuning_mode),
-                Conv2dLoRA(dim[0], dim[0], 3, padding=1,  r = 8 if tuning_mode == 'lora' else 0))
-            output_list.append(conv_out)
-        self.outputs04 = nn.ModuleList(output_list)
-        output_list = []
-        for dim in output_dim:
-            conv_out = nn.Sequential(
-                ResidualBlock(in_dim[1], dim[1], 'layer', stride=1, tuning_mode=tuning_mode),
-                Conv2dLoRA(dim[1], dim[1], 3, padding=1, r = 8 if tuning_mode == 'lora' else 0))
-            output_list.append(conv_out)
-        self.outputs08 = nn.ModuleList(output_list)
-        output_list = []
-        for dim in output_dim:
-            conv_out = nn.Sequential(
-                ResidualBlock(in_dim[2], dim[2], 'layer', stride=1, tuning_mode=tuning_mode),
-                Conv2dLoRA(dim[2], dim[2], 3, padding=1,  r = 8 if tuning_mode == 'lora' else 0))
-            output_list.append(conv_out)
-        self.outputs16 = nn.ModuleList(output_list)
-        # output_list = []
-        # for dim in output_dim:
-        #     conv_out = Conv2dLoRA(in_dim[3], dim[3], 3, padding=1)
-        #     output_list.append(conv_out)
-        # self.outputs32 = nn.ModuleList(output_list)
-    def forward(self, encoder_features):
-        x_4, x_8, x_16, x_32 = encoder_features
-        outputs04 = [f(x_4) for f in self.outputs04]
-        outputs08 = [f(x_8) for f in self.outputs08]
-        outputs16 = [f(x_16)for f in self.outputs16]
-        # outputs32 = [f(x_32) for f in self.outputs32]
-        return (outputs04, outputs08, outputs16)
-class ConvBlock(nn.Module):
-    # reimplementation of DPT
-    def __init__(self, channels, tuning_mode=None):
-        super(ConvBlock, self).__init__()
-        self.act = nn.ReLU(inplace=True)
-        self.conv1 = Conv2dLoRA(
-            channels,
-            channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            r = 8 if tuning_mode == 'lora' else 0
-        )
-        self.conv2 = Conv2dLoRA(
-            channels,
-            channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            r = 8 if tuning_mode == 'lora' else 0
-        )
-    def forward(self, x):
-        out = self.act(x)
-        out = self.conv1(out)
-        out = self.act(out)
-        out = self.conv2(out)
-        return x + out
-class FuseBlock(nn.Module):
-    # reimplementation of DPT
-    def __init__(self, in_channels, out_channels, fuse=True, upsample=True, scale_factor=2, tuning_mode=None):
-        super(FuseBlock, self).__init__()
-        self.fuse = fuse
-        self.scale_factor = scale_factor
-        self.way_trunk = ConvBlock(in_channels, tuning_mode=tuning_mode)
-        if self.fuse:
-            self.way_branch = ConvBlock(in_channels, tuning_mode=tuning_mode)
-        self.out_conv = Conv2dLoRA(
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            r = 8 if tuning_mode == 'lora' else 0
-        )
-        self.upsample = upsample
-    def forward(self, x1, x2=None):
-        if x2 is not None:
-            x2 = self.way_branch(x2)
-            x1 = x1 + x2
-        out = self.way_trunk(x1)
-        if self.upsample:
-            out = interpolate_float32(
-                out, scale_factor=self.scale_factor, mode="bilinear", align_corners=True
-            )
-        out = self.out_conv(out)
-        return out
-class Readout(nn.Module):
-    # From DPT
-    def __init__(self, in_features, use_cls_token=True, num_register_tokens=0, tuning_mode=None):
-        super(Readout, self).__init__()
-        self.use_cls_token = use_cls_token
-        if self.use_cls_token == True:
-            self.project_patch = LoRALinear(in_features, in_features, r = 8 if tuning_mode == 'lora' else 0)
-            self.project_learn = LoRALinear((1 + num_register_tokens) * in_features, in_features, bias=False, r = 8 if tuning_mode == 'lora' else 0)
-            self.act = nn.GELU()
-        else:
-            self.project = nn.Identity()
-    def forward(self, x):
-        if self.use_cls_token == True:
-            x_patch = self.project_patch(x[0])
-            x_learn = self.project_learn(x[1])
-            x_learn = x_learn.expand_as(x_patch).contiguous()
-            features = x_patch + x_learn
-            return self.act(features)
-        else:
-            return self.project(x)
-class Token2Feature(nn.Module):
-    # From DPT
-    def __init__(self, vit_channel, feature_channel, scale_factor, use_cls_token=True, num_register_tokens=0, tuning_mode=None):
-        super(Token2Feature, self).__init__()
-        self.scale_factor = scale_factor
-        self.readoper = Readout(in_features=vit_channel, use_cls_token=use_cls_token, num_register_tokens=num_register_tokens,  tuning_mode=tuning_mode)
-        if scale_factor > 1 and isinstance(scale_factor, int):
-            self.sample = ConvTranspose2dLoRA(r = 8 if tuning_mode == 'lora' else 0,
-                in_channels=vit_channel,
-                out_channels=feature_channel,
-                kernel_size=scale_factor,
-                stride=scale_factor,
-                padding=0,
-            )
-        elif scale_factor > 1:
-            self.sample = nn.Sequential(
-                # Upsample2(upscale=scale_factor),
-                # nn.Upsample(scale_factor=scale_factor),
-                Conv2dLoRA(r = 8 if tuning_mode == 'lora' else 0,
-                    in_channels=vit_channel,
-                    out_channels=feature_channel,
-                    kernel_size=1,
-                    stride=1,
-                    padding=0,
-                ),
-            )
-        elif scale_factor < 1:
-            scale_factor = int(1.0 / scale_factor)
-            self.sample = Conv2dLoRA(r = 8 if tuning_mode == 'lora' else 0,
-                in_channels=vit_channel,
-                out_channels=feature_channel,
-                kernel_size=scale_factor+1,
-                stride=scale_factor,
-                padding=1,
-            )
-        else:
-            self.sample = nn.Identity()
-    def forward(self, x):
-        x = self.readoper(x)
-        #if use_cls_token == True:
-        x = x.permute(0, 3, 1, 2).contiguous()
-        if isinstance(self.scale_factor, float):
-            x = interpolate_float32(x.float(), scale_factor=self.scale_factor, mode='nearest')
-        x = self.sample(x)
-        return x
-class EncoderFeature(nn.Module):
-    def __init__(self, vit_channel, num_ch_dec=[256, 512, 1024, 1024], use_cls_token=True, num_register_tokens=0, tuning_mode=None):
-        super(EncoderFeature, self).__init__()
-        self.vit_channel = vit_channel
-        self.num_ch_dec = num_ch_dec
-        self.read_3 = Token2Feature(self.vit_channel, self.num_ch_dec[3], scale_factor=1, use_cls_token=use_cls_token, num_register_tokens=num_register_tokens, tuning_mode=tuning_mode)
-        self.read_2 = Token2Feature(self.vit_channel, self.num_ch_dec[2], scale_factor=1, use_cls_token=use_cls_token, num_register_tokens=num_register_tokens, tuning_mode=tuning_mode)
-        self.read_1 = Token2Feature(self.vit_channel, self.num_ch_dec[1], scale_factor=2, use_cls_token=use_cls_token, num_register_tokens=num_register_tokens, tuning_mode=tuning_mode)
-        self.read_0 = Token2Feature(self.vit_channel, self.num_ch_dec[0], scale_factor=7/2, use_cls_token=use_cls_token, num_register_tokens=num_register_tokens, tuning_mode=tuning_mode)
-    def forward(self, ref_feature):
-        x = self.read_3(ref_feature[3])  # 1/14
-        x2 = self.read_2(ref_feature[2]) # 1/14
-        x1 = self.read_1(ref_feature[1]) # 1/7
-        x0 = self.read_0(ref_feature[0]) # 1/4
-        return x, x2, x1, x0
-class DecoderFeature(nn.Module):
-    def __init__(self, vit_channel, num_ch_dec=[128, 256, 512, 1024, 1024], use_cls_token=True, tuning_mode=None):
-        super(DecoderFeature, self).__init__()
-        self.vit_channel = vit_channel
-        self.num_ch_dec = num_ch_dec
-        self.upconv_3 = FuseBlock(
-            self.num_ch_dec[4],
-            self.num_ch_dec[3],
-        fuse=False, upsample=False, tuning_mode=tuning_mode)
-        self.upconv_2 = FuseBlock(
-            self.num_ch_dec[3],
-            self.num_ch_dec[2],
-        tuning_mode=tuning_mode)
-        self.upconv_1 = FuseBlock(
-            self.num_ch_dec[2],
-            self.num_ch_dec[1] + 2,
-            scale_factor=7/4,
-        tuning_mode=tuning_mode)
-        # self.upconv_0 = FuseBlock(
-        #     self.num_ch_dec[1],
-        #     self.num_ch_dec[0] + 1,
-        # )
-    def forward(self, ref_feature):
-        x, x2, x1, x0 = ref_feature # 1/14 1/14 1/7 1/4
-        x = self.upconv_3(x)     # 1/14
-        x = self.upconv_2(x, x2) # 1/7
-        x = self.upconv_1(x, x1) # 1/4
-        # x = self.upconv_0(x, x0) # 4/7
-        return x
-class RAFTDepthNormalDPT5(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.in_channels = cfg.model.decode_head.in_channels # [1024, 1024, 1024, 1024]
-        self.feature_channels = cfg.model.decode_head.feature_channels # [256, 512, 1024, 1024] [2/7, 1/7, 1/14, 1/14]
-        self.decoder_channels = cfg.model.decode_head.decoder_channels # [128, 256, 512, 1024, 1024] [-, 1/4, 1/7, 1/14, 1/14]
-        self.use_cls_token = cfg.model.decode_head.use_cls_token
-        self.up_scale = cfg.model.decode_head.up_scale
-        self.num_register_tokens = cfg.model.decode_head.num_register_tokens
-        self.min_val = cfg.data_basic.depth_normalize[0]
-        self.max_val = cfg.data_basic.depth_normalize[1]
-        self.regress_scale = 100.0\
-        try:
-            tuning_mode = cfg.model.decode_head.tuning_mode
-        except:
-            tuning_mode = None
-        self.tuning_mode = tuning_mode
-        self.hidden_dims = self.context_dims = cfg.model.decode_head.hidden_channels # [128, 128, 128, 128]
-        self.n_gru_layers = cfg.model.decode_head.n_gru_layers # 3
-        self.n_downsample = cfg.model.decode_head.n_downsample # 3, resolution of the disparity field (1/2^K)
-        self.iters = cfg.model.decode_head.iters # 22
-        self.slow_fast_gru = cfg.model.decode_head.slow_fast_gru # True
-        self.num_depth_regressor_anchor = 256 # 512
-        self.used_res_channel = self.decoder_channels[1] # now, use 2/7 res
-        self.token2feature = EncoderFeature(self.in_channels[0], self.feature_channels, self.use_cls_token, self.num_register_tokens, tuning_mode=tuning_mode)
-        self.decoder_mono = DecoderFeature(self.in_channels, self.decoder_channels, tuning_mode=tuning_mode)
-        self.depth_regressor = nn.Sequential(
-            Conv2dLoRA(self.used_res_channel,
-                      self.num_depth_regressor_anchor,
-                      kernel_size=3,
-                      padding=1, r = 8 if tuning_mode == 'lora' else 0),
-            # nn.BatchNorm2d(self.num_depth_regressor_anchor),
-            nn.ReLU(inplace=True),
-            Conv2dLoRA(self.num_depth_regressor_anchor,
-                      self.num_depth_regressor_anchor,
-                      kernel_size=1, r = 8 if tuning_mode == 'lora' else 0),
-        )
-        self.normal_predictor = nn.Sequential(
-            Conv2dLoRA(self.used_res_channel,
-                      128,
-                      kernel_size=3,
-                      padding=1, r = 8 if tuning_mode == 'lora' else 0,),
-            # nn.BatchNorm2d(128),
-            nn.ReLU(inplace=True),
-            Conv2dLoRA(128, 128, kernel_size=1, r = 8 if tuning_mode == 'lora' else 0), nn.ReLU(inplace=True),
-            Conv2dLoRA(128, 128, kernel_size=1, r = 8 if tuning_mode == 'lora' else 0), nn.ReLU(inplace=True),
-            Conv2dLoRA(128, 3, kernel_size=1, r = 8 if tuning_mode == 'lora' else 0),
-        )
-        self.context_feature_encoder = ContextFeatureEncoder(self.feature_channels, [self.hidden_dims, self.context_dims], tuning_mode=tuning_mode)
-        self.context_zqr_convs = nn.ModuleList([Conv2dLoRA(self.context_dims[i], self.hidden_dims[i]*3, 3, padding=3//2, r = 8 if tuning_mode == 'lora' else 0) for i in range(self.n_gru_layers)])
-        self.update_block = BasicMultiUpdateBlock(cfg, hidden_dims=self.hidden_dims, out_dims=6, tuning_mode=tuning_mode)
-        self.relu = nn.ReLU(inplace=True)
-    def get_bins(self, bins_num):
-        depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cuda")
-        depth_bins_vec = torch.exp(depth_bins_vec)
-        return depth_bins_vec
-    def register_depth_expectation_anchor(self, bins_num, B):
-        depth_bins_vec = self.get_bins(bins_num)
-        depth_bins_vec = depth_bins_vec.unsqueeze(0).repeat(B, 1)
-        self.register_buffer('depth_expectation_anchor', depth_bins_vec, persistent=False)
-    def clamp(self, x):
-        y = self.relu(x - self.min_val) + self.min_val
-        y = self.max_val - self.relu(self.max_val - y)
-        return y
-    def regress_depth(self, feature_map_d):
-        prob_feature = self.depth_regressor(feature_map_d)
-        prob = prob_feature.softmax(dim=1)
-        #prob = prob_feature.float().softmax(dim=1)
-        ## Error logging
-        if torch.isnan(prob).any():
-            print('prob_feat_nan!!!')
-        if torch.isinf(prob).any():
-            print('prob_feat_inf!!!')
-        # h = prob[0,:,0,0].cpu().numpy().reshape(-1)
-        # import matplotlib.pyplot as plt
-        # plt.bar(range(len(h)), h)
-        B = prob.shape[0]
-        if "depth_expectation_anchor" not in self._buffers:
-            self.register_depth_expectation_anchor(self.num_depth_regressor_anchor, B)
-        d = compute_depth_expectation(
-            prob,
-            self.depth_expectation_anchor[:B, ...]).unsqueeze(1)
-        ## Error logging
-        if torch.isnan(d ).any():
-            print('d_nan!!!')
-        if torch.isinf(d ).any():
-            print('d_inf!!!')
-        return (self.clamp(d) - self.max_val)/ self.regress_scale, prob_feature
-    def pred_normal(self, feature_map, confidence):
-        normal_out = self.normal_predictor(feature_map)
-        ## Error logging
-        if torch.isnan(normal_out).any():
-            print('norm_nan!!!')
-        if torch.isinf(normal_out).any():
-            print('norm_feat_inf!!!')
-        return norm_normalize(torch.cat([normal_out, confidence], dim=1))
-        #return norm_normalize(torch.cat([normal_out, confidence], dim=1).float())
-    def create_mesh_grid(self, height, width, batch, device="cuda", set_buffer=True):
-        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=device),
-                               torch.arange(0, width, dtype=torch.float32, device=device)], indexing='ij')
-        meshgrid = torch.stack((x, y))
-        meshgrid = meshgrid.unsqueeze(0).repeat(batch, 1, 1, 1)
-        #self.register_buffer('meshgrid', meshgrid, persistent=False)
-        return meshgrid
-    def upsample_flow(self, flow, mask):
-        """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
-        N, D, H, W = flow.shape
-        factor = 2 ** self.n_downsample
-        mask = mask.view(N, 1, 9, factor, factor, H, W)
-        mask = torch.softmax(mask, dim=2)
-        #mask = torch.softmax(mask.float(), dim=2)
-        #up_flow = F.unfold(factor * flow, [3,3], padding=1)
-        up_flow = F.unfold(flow, [3,3], padding=1)
-        up_flow = up_flow.view(N, D, 9, 1, 1, H, W)
-        up_flow = torch.sum(mask * up_flow, dim=2)
-        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
-        return up_flow.reshape(N, D, factor*H, factor*W)
-    def initialize_flow(self, img):
-        """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
-        N, _, H, W = img.shape
-        coords0 = coords_grid(N, H, W).to(img.device)
-        coords1 = coords_grid(N, H, W).to(img.device)
-        return coords0, coords1
-    def upsample(self, x, scale_factor=2):
-        """Upsample input tensor by a factor of 2
-        """
-        return interpolate_float32(x, scale_factor=scale_factor*self.up_scale/8, mode="nearest")
-    def forward(self, vit_features, **kwargs):
-        ## read vit token to multi-scale features
-        B, H, W, _, _, num_register_tokens = vit_features[1]
-        vit_features = vit_features[0]
-        ## Error logging
-        if torch.isnan(vit_features[0]).any():
-            print('vit_feature_nan!!!')
-        if torch.isinf(vit_features[0]).any():
-            print('vit_feature_inf!!!')
-        if self.use_cls_token == True:
-            vit_features = [[ft[:, 1+num_register_tokens:, :].view(B, H, W, self.in_channels[0]), \
-                ft[:, 0:1+num_register_tokens, :].view(B, 1, 1, self.in_channels[0] * (1+num_register_tokens))] for ft in vit_features]
-        else:
-            vit_features = [ft.view(B, H, W, self.in_channels[0]) for ft in vit_features]
-        encoder_features = self.token2feature(vit_features) # 1/14, 1/14, 1/7, 1/4
-        ## Error logging
-        for en_ft in encoder_features:
-            if torch.isnan(en_ft).any():
-                print('decoder_feature_nan!!!')
-                print(en_ft.shape)
-            if torch.isinf(en_ft).any():
-                print('decoder_feature_inf!!!')
-                print(en_ft.shape)
-        ## decode features to init-depth (and confidence)
-        ref_feat= self.decoder_mono(encoder_features) # now, 1/4 for depth
-        ## Error logging
-        if torch.isnan(ref_feat).any():
-            print('ref_feat_nan!!!')
-        if torch.isinf(ref_feat).any():
-            print('ref_feat_inf!!!')
-        feature_map = ref_feat[:, :-2, :, :] # feature map share of depth and normal prediction
-        depth_confidence_map = ref_feat[:, -2:-1, :, :]
-        normal_confidence_map = ref_feat[:, -1:, :, :]
-        depth_pred, binmap = self.regress_depth(feature_map) # regress bin for depth
-        normal_pred = self.pred_normal(feature_map, normal_confidence_map) # mlp for normal
-        depth_init = torch.cat((depth_pred, depth_confidence_map, normal_pred), dim=1) # (N, 1+1+4, H, W)
-        ## encoder features to context-feature for init-hidden-state and contex-features
-        cnet_list = self.context_feature_encoder(encoder_features[::-1])
-        net_list = [torch.tanh(x[0]) for x in cnet_list] # x_4, x_8, x_16 of hidden state
-        inp_list = [torch.relu(x[1]) for x in cnet_list] # x_4, x_8, x_16 context features
-        # Rather than running the GRU's conv layers on the context features multiple times, we do it once at the beginning
-        inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i,conv in zip(inp_list, self.context_zqr_convs)]
-        coords0, coords1 = self.initialize_flow(net_list[0])
-        if depth_init is not None:
-            coords1 = coords1 + depth_init
-        if self.training:
-            low_resolution_init = [self.clamp(depth_init[:,:1] * self.regress_scale + self.max_val), depth_init[:,1:2], norm_normalize(depth_init[:,2:].clone())]
-            init_depth = upflow4(depth_init)
-            flow_predictions = [self.clamp(init_depth[:,:1] * self.regress_scale + self.max_val)]
-            conf_predictions = [init_depth[:,1:2]]
-            normal_outs = [norm_normalize(init_depth[:,2:].clone())]
-        else:
-            flow_predictions = []
-            conf_predictions = []
-            samples_pred_list = []
-            coord_list = []
-            normal_outs = []
-            low_resolution_init = []
-        for itr in range(self.iters):
-            # coords1 = coords1.detach()
-            flow = coords1 - coords0
-            if self.n_gru_layers == 3 and self.slow_fast_gru: # Update low-res GRU
-                net_list = self.update_block(net_list, inp_list, iter32=True, iter16=False, iter08=False, update=False)
-            if self.n_gru_layers >= 2 and self.slow_fast_gru:# Update low-res GRU and mid-res GRU
-                net_list = self.update_block(net_list, inp_list, iter32=self.n_gru_layers==3, iter16=True, iter08=False, update=False)
-            net_list, up_mask, delta_flow = self.update_block(net_list, inp_list, None, flow, iter32=self.n_gru_layers==3, iter16=self.n_gru_layers>=2)
-            # F(t+1) = F(t) + \Delta(t)
-            coords1 = coords1 + delta_flow
-            # We do not need to upsample or output intermediate results in test_mode
-            #if (not self.training) and itr < self.iters-1:
-                #continue
-            # upsample predictions
-            if up_mask is None:
-                flow_up = self.upsample(coords1-coords0, 4)
-            else:
-                flow_up = self.upsample_flow(coords1 - coords0, up_mask)
-                # flow_up = self.upsample(coords1-coords0, 4)
-            flow_predictions.append(self.clamp(flow_up[:,:1] * self.regress_scale + self.max_val))
-            conf_predictions.append(flow_up[:,1:2])
-            normal_outs.append(norm_normalize(flow_up[:,2:].clone()))
-        outputs=dict(
-            prediction=flow_predictions[-1],
-            predictions_list=flow_predictions,
-            confidence=conf_predictions[-1],
-            confidence_list=conf_predictions,
-            pred_logit=None,
-            # samples_pred_list=samples_pred_list,
-            # coord_list=coord_list,
-            prediction_normal=normal_outs[-1],
-            normal_out_list=normal_outs,
-            low_resolution_init=low_resolution_init,
-        )
-        return outputs
-if __name__ == "__main__":
-    try:
-        from mmcv.utils import Config
-    except:
-        from mmengine import Config
-    cfg = Config.fromfile('/cpfs01/shared/public/users/mu.hu/monodepth/mono/configs/RAFTDecoder/vit.raft.full2t.py')
-    cfg.model.decode_head.in_channels = [384, 384, 384, 384]
-    cfg.model.decode_head.feature_channels = [96, 192, 384, 768]
-    cfg.model.decode_head.decoder_channels = [48, 96, 192, 384, 384]
-    cfg.model.decode_head.hidden_channels = [48, 48, 48, 48, 48]
-    cfg.model.decode_head.up_scale = 7
-    # cfg.model.decode_head.use_cls_token = True
-    # vit_feature = [[torch.rand((2, 20, 60, 384)).cuda(), torch.rand(2, 384).cuda()], \
-    #         [torch.rand((2, 20, 60, 384)).cuda(), torch.rand(2, 384).cuda()], \
-    #         [torch.rand((2, 20, 60, 384)).cuda(), torch.rand(2, 384).cuda()], \
-    #         [torch.rand((2, 20, 60, 384)).cuda(), torch.rand(2, 384).cuda()]]
-    cfg.model.decode_head.use_cls_token = True
-    cfg.model.decode_head.num_register_tokens = 4
-    vit_feature = [[torch.rand((2, (74 * 74) + 5, 384)).cuda(),\
-                    torch.rand((2, (74 * 74) + 5, 384)).cuda(), \
-                    torch.rand((2, (74 * 74) + 5, 384)).cuda(), \
-                    torch.rand((2, (74 * 74) + 5, 384)).cuda()], (2, 74, 74, 1036, 1036, 4)]
-    decoder = RAFTDepthNormalDPT5(cfg).cuda()
-    output = decoder(vit_feature)
-    temp = 1

mono/model/decode_heads/__init__.py DELETED Viewed

@@ -1,4 +0,0 @@
-from .HourGlassDecoder import HourglassDecoder
-from .RAFTDepthNormalDPTDecoder5 import RAFTDepthNormalDPT5
-__all__=['HourglassDecoder', 'RAFTDepthNormalDPT5']

mono/model/decode_heads/__pycache__/HourGlassDecoder.cpython-39.pyc DELETED Viewed

Binary file (8.65 kB)

mono/model/decode_heads/__pycache__/__init__.cpython-39.pyc DELETED Viewed

Binary file (252 Bytes)

mono/model/model_pipelines/__base_model__.py DELETED Viewed

@@ -1,20 +0,0 @@
-import torch
-import torch.nn as nn
-from mono.utils.comm import get_func
-class BaseDepthModel(nn.Module):
-    def __init__(self, cfg, **kwargs) -> None:
-        super(BaseDepthModel, self).__init__()
-        model_type = cfg.model.type
-        self.depth_model = get_func('mono.model.model_pipelines.' + model_type)(cfg)
-    def forward(self, data):
-        output = self.depth_model(**data)
-        return output['prediction'], output['confidence'], output
-    def inference(self, data):
-        with torch.no_grad():
-            pred_depth, confidence, _ = self.forward(data)
-        return pred_depth, confidence

mono/model/model_pipelines/__init__.py DELETED Viewed

@@ -1,6 +0,0 @@
-from .dense_pipeline import DensePredModel
-from .__base_model__ import BaseDepthModel
-__all__ = [
-    'DensePredModel', 'BaseDepthModel',
-]

mono/model/model_pipelines/__pycache__/__base_model__.cpython-39.pyc DELETED Viewed

Binary file (1.19 kB)

mono/model/model_pipelines/__pycache__/__init__.cpython-39.pyc DELETED Viewed

Binary file (313 Bytes)

mono/model/model_pipelines/__pycache__/dense_pipeline.cpython-39.pyc DELETED Viewed

Binary file (1.01 kB)

mono/model/model_pipelines/dense_pipeline.py DELETED Viewed

@@ -1,16 +0,0 @@
-import torch
-import torch.nn as nn
-from mono.utils.comm import get_func
-class DensePredModel(nn.Module):
-    def __init__(self, cfg) -> None:
-        super(DensePredModel, self).__init__()
-        self.encoder = get_func('mono.model.' + cfg.model.backbone.prefix + cfg.model.backbone.type)(**cfg.model.backbone)
-        self.decoder = get_func('mono.model.' + cfg.model.decode_head.prefix + cfg.model.decode_head.type)(cfg)
-    def forward(self, input, **kwargs):
-        # [f_32, f_16, f_8, f_4]
-        features = self.encoder(input)
-        out = self.decoder(features, **kwargs)
-        return out

mono/model/monodepth_model.py DELETED Viewed

@@ -1,37 +0,0 @@
-import torch
-import torch.nn as nn
-from .model_pipelines.__base_model__ import BaseDepthModel
-class DepthModel(BaseDepthModel):
-    def __init__(self, cfg, **kwards):
-        super(DepthModel, self).__init__(cfg)
-        model_type = cfg.model.type
-    def inference(self, data):
-        with torch.no_grad():
-            pred_depth, confidence, output_dict = self.forward(data)
-        return pred_depth, confidence, output_dict
-def get_monodepth_model(
-    cfg : dict,
-    **kwargs
-    ) -> nn.Module:
-    # config depth  model
-    model = DepthModel(cfg, **kwargs)
-    #model.init_weights(load_imagenet_model, imagenet_ckpt_fpath)
-    assert isinstance(model, nn.Module)
-    return model
-def get_configured_monodepth_model(
-    cfg: dict,
-    ) -> nn.Module:
-    """
-        Args:
-        @ configs: configures for the network.
-        @ load_imagenet_model: whether to initialize from ImageNet-pretrained model.
-        @ imagenet_ckpt_fpath: string representing path to file with weights to initialize model with.
-        Returns:
-        # model: depth model.
-    """
-    model = get_monodepth_model(cfg)
-    return model

mono/tools/test_scale_cano.py DELETED Viewed

@@ -1,158 +0,0 @@
-import os
-import os.path as osp
-import cv2
-import time
-import sys
-CODE_SPACE=os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-sys.path.append(CODE_SPACE)
-import argparse
-import mmcv
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-try:
-    from mmcv.utils import Config, DictAction
-except:
-    from mmengine import Config, DictAction
-from datetime import timedelta
-import random
-import numpy as np
-from mono.utils.logger import setup_logger
-import glob
-from mono.utils.comm import init_env
-from mono.model.monodepth_model import get_configured_monodepth_model
-from mono.utils.running import load_ckpt
-from mono.utils.do_test import do_scalecano_test_with_custom_data
-from mono.utils.mldb import load_data_info, reset_ckpt_path
-from mono.utils.custom_data import load_from_annos, load_data
-def parse_args():
-    parser = argparse.ArgumentParser(description='Train a segmentor')
-    parser.add_argument('config', help='train config file path')
-    parser.add_argument('--show-dir', help='the dir to save logs and visualization results')
-    parser.add_argument('--load-from', help='the checkpoint file to load weights from')
-    parser.add_argument('--node_rank', type=int, default=0)
-    parser.add_argument('--nnodes', type=int, default=1, help='number of nodes')
-    parser.add_argument('--options', nargs='+', action=DictAction, help='custom options')
-    parser.add_argument('--launcher', choices=['None', 'pytorch', 'slurm', 'mpi', 'ror'], default='slurm', help='job launcher')
-    parser.add_argument('--test_data_path', default='None', type=str, help='the path of test data')
-    args = parser.parse_args()
-    return args
-def main(args):
-    os.chdir(CODE_SPACE)
-    cfg = Config.fromfile(args.config)
-    if args.options is not None:
-        cfg.merge_from_dict(args.options)
-    # show_dir is determined in this priority: CLI > segment in file > filename
-    if args.show_dir is not None:
-        # update configs according to CLI args if args.show_dir is not None
-        cfg.show_dir = args.show_dir
-    else:
-        # use condig filename + timestamp as default show_dir if args.show_dir is None
-        cfg.show_dir = osp.join('./show_dirs',
-                                osp.splitext(osp.basename(args.config))[0],
-                                args.timestamp)
-    # ckpt path
-    if args.load_from is None:
-        raise RuntimeError('Please set model path!')
-    cfg.load_from = args.load_from
-    # load data info
-    data_info = {}
-    load_data_info('data_info', data_info=data_info)
-    cfg.mldb_info = data_info
-    # update check point info
-    reset_ckpt_path(cfg.model, data_info)
-    # create show dir
-    os.makedirs(osp.abspath(cfg.show_dir), exist_ok=True)
-    # init the logger before other steps
-    cfg.log_file = osp.join(cfg.show_dir, f'{args.timestamp}.log')
-    logger = setup_logger(cfg.log_file)
-    # log some basic info
-    logger.info(f'Config:\n{cfg.pretty_text}')
-    # init distributed env dirst, since logger depends on the dist info
-    if args.launcher == 'None':
-        cfg.distributed = False
-    else:
-        cfg.distributed = True
-        init_env(args.launcher, cfg)
-    logger.info(f'Distributed training: {cfg.distributed}')
-    # dump config
-    cfg.dump(osp.join(cfg.show_dir, osp.basename(args.config)))
-    test_data_path = args.test_data_path
-    if not os.path.isabs(test_data_path):
-        test_data_path = osp.join(CODE_SPACE, test_data_path)
-    if 'json' in test_data_path:
-        test_data = load_from_annos(test_data_path)
-    else:
-        test_data = load_data(args.test_data_path)
-    if not cfg.distributed:
-        main_worker(0, cfg, args.launcher, test_data)
-    else:
-        # distributed training
-        if args.launcher == 'ror':
-            local_rank = cfg.dist_params.local_rank
-            main_worker(local_rank, cfg, args.launcher, test_data)
-        else:
-            mp.spawn(main_worker, nprocs=cfg.dist_params.num_gpus_per_node, args=(cfg, args.launcher, test_data))
-def main_worker(local_rank: int, cfg: dict, launcher: str, test_data: list):
-    if cfg.distributed:
-        cfg.dist_params.global_rank = cfg.dist_params.node_rank * cfg.dist_params.num_gpus_per_node + local_rank
-        cfg.dist_params.local_rank = local_rank
-        if launcher == 'ror':
-            init_torch_process_group(use_hvd=False)
-        else:
-            torch.cuda.set_device(local_rank)
-            default_timeout = timedelta(minutes=30)
-            dist.init_process_group(
-                backend=cfg.dist_params.backend,
-                init_method=cfg.dist_params.dist_url,
-                world_size=cfg.dist_params.world_size,
-                rank=cfg.dist_params.global_rank,
-                timeout=default_timeout)
-    logger = setup_logger(cfg.log_file)
-    # build model
-    model = get_configured_monodepth_model(cfg, )
-    # config distributed training
-    if cfg.distributed:
-        model = torch.nn.parallel.DistributedDataParallel(model.cuda(),
-                                                          device_ids=[local_rank],
-                                                          output_device=local_rank,
-                                                          find_unused_parameters=True)
-    else:
-        model = torch.nn.DataParallel(model).cuda()
-    # load ckpt
-    model, _,  _, _ = load_ckpt(cfg.load_from, model, strict_match=False)
-    model.eval()
-    do_scalecano_test_with_custom_data(
-        model,
-        cfg,
-        test_data,
-        logger,
-        cfg.distributed,
-        local_rank
-    )
-if __name__ == '__main__':
-    args = parse_args()
-    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
-    args.timestamp = timestamp
-    main(args)

mono/utils/__init__.py DELETED Viewed

	@@ -1 +0,0 @@
1	-

mono/utils/__pycache__/__init__.cpython-39.pyc DELETED Viewed

Binary file (160 Bytes)

mono/utils/__pycache__/avg_meter.cpython-39.pyc DELETED Viewed

Binary file (10.1 kB)

mono/utils/__pycache__/comm.cpython-39.pyc DELETED Viewed

Binary file (9.72 kB)

mono/utils/__pycache__/custom_data.cpython-39.pyc DELETED Viewed

Binary file (1.21 kB)

mono/utils/__pycache__/do_test.cpython-39.pyc DELETED Viewed

Binary file (8.71 kB)

mono/utils/__pycache__/logger.cpython-39.pyc DELETED Viewed

Binary file (3.17 kB)

mono/utils/__pycache__/mldb.cpython-39.pyc DELETED Viewed

Binary file (1.18 kB)

mono/utils/__pycache__/running.cpython-39.pyc DELETED Viewed

Binary file (2.09 kB)

mono/utils/__pycache__/transform.cpython-39.pyc DELETED Viewed

Binary file (11.5 kB)

mono/utils/__pycache__/unproj_pcd.cpython-39.pyc DELETED Viewed

Binary file (2.61 kB)