Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,816 Bytes
a7dedf9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
from torch import nn
from typing import Optional
from functools import partial
from .utils import _init_weights, interpolate_pos_embed
from .blocks import DepthSeparableConv2d, conv1x1, conv3x3, Conv2dLayerNorm
from .refine import ConvRefine, LightConvRefine, LighterConvRefine
from .downsample import ConvDownsample, LightConvDownsample, LighterConvDownsample
from .upsample import ConvUpsample, LightConvUpsample, LighterConvUpsample
from .multi_scale import MultiScale
from .blocks import ConvAdapter, ViTAdapter
def _get_norm_layer(model: nn.Module) -> Optional[nn.Module]:
for module in model.modules():
if isinstance(module, nn.BatchNorm2d):
return nn.BatchNorm2d
elif isinstance(module, nn.GroupNorm):
num_groups = module.num_groups
return partial(nn.GroupNorm, num_groups=num_groups)
elif isinstance(module, (nn.LayerNorm, Conv2dLayerNorm)):
return Conv2dLayerNorm
return None
def _get_activation(model: nn.Module) -> Optional[nn.Module]:
for module in model.modules():
if isinstance(module, nn.BatchNorm2d):
return nn.ReLU(inplace=True)
elif isinstance(module, nn.GroupNorm):
return nn.ReLU(inplace=True)
elif isinstance(module, (nn.LayerNorm, Conv2dLayerNorm)):
return nn.GELU()
return nn.GELU()
__all__ = [
"_init_weights", "_check_norm_layer", "_check_activation",
"conv1x1",
"conv3x3",
"Conv2dLayerNorm",
"interpolate_pos_embed",
"DepthSeparableConv2d",
"ConvRefine",
"LightConvRefine",
"LighterConvRefine",
"ConvDownsample",
"LightConvDownsample",
"LighterConvDownsample",
"ConvUpsample",
"LightConvUpsample",
"LighterConvUpsample",
"MultiScale",
"ConvAdapter", "ViTAdapter",
]
|