158 lines
5.1 KiB
Python
158 lines
5.1 KiB
Python
from ._fx import (
|
|
create_feature_extractor,
|
|
get_graph_node_names,
|
|
register_notrace_function,
|
|
register_notrace_module,
|
|
is_notrace_module,
|
|
is_notrace_function,
|
|
get_notrace_modules,
|
|
get_notrace_functions,
|
|
)
|
|
from .activations import *
|
|
from .adaptive_avgmax_pool import (
|
|
adaptive_avgmax_pool2d,
|
|
select_adaptive_pool2d,
|
|
AdaptiveAvgMaxPool2d,
|
|
SelectAdaptivePool2d,
|
|
)
|
|
from .attention import Attention, AttentionRope, maybe_add_mask
|
|
from .attention2d import MultiQueryAttention2d, Attention2d, MultiQueryAttentionV2
|
|
from .attention_pool import AttentionPoolLatent
|
|
from .attention_pool2d import AttentionPool2d, RotAttentionPool2d, RotaryEmbedding
|
|
from .blur_pool import BlurPool2d, create_aa
|
|
from .classifier import create_classifier, ClassifierHead, NormMlpClassifierHead, ClNormMlpClassifierHead
|
|
from .cond_conv2d import CondConv2d, get_condconv_initializer
|
|
from .config import (
|
|
is_exportable,
|
|
is_scriptable,
|
|
is_no_jit,
|
|
use_fused_attn,
|
|
set_exportable,
|
|
set_scriptable,
|
|
set_no_jit,
|
|
set_layer_config,
|
|
set_fused_attn,
|
|
set_reentrant_ckpt,
|
|
use_reentrant_ckpt,
|
|
)
|
|
from .conv2d_same import Conv2dSame, conv2d_same
|
|
from .conv_bn_act import ConvNormAct, ConvNormActAa, ConvBnAct
|
|
from .create_act import create_act_layer, get_act_layer, get_act_fn
|
|
from .create_attn import get_attn, create_attn
|
|
from .create_conv2d import create_conv2d
|
|
from .create_norm import get_norm_layer, create_norm_layer
|
|
from .create_norm_act import get_norm_act_layer, create_norm_act_layer, get_norm_act_layer
|
|
from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path, calculate_drop_path_rates
|
|
from .eca import EcaModule, CecaModule, EfficientChannelAttn, CircularEfficientChannelAttn
|
|
from .evo_norm import (
|
|
EvoNorm2dB0,
|
|
EvoNorm2dB1,
|
|
EvoNorm2dB2,
|
|
EvoNorm2dS0,
|
|
EvoNorm2dS0a,
|
|
EvoNorm2dS1,
|
|
EvoNorm2dS1a,
|
|
EvoNorm2dS2,
|
|
EvoNorm2dS2a,
|
|
)
|
|
from .fast_norm import is_fast_norm, set_fast_norm, fast_group_norm, fast_layer_norm
|
|
from .filter_response_norm import FilterResponseNormTlu2d, FilterResponseNormAct2d
|
|
from .format import Format, get_channel_dim, get_spatial_dim, nchw_to, nhwc_to
|
|
from .gather_excite import GatherExcite
|
|
from .global_context import GlobalContext
|
|
from .grid import ndgrid, meshgrid
|
|
from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple, make_divisible, extend_tuple
|
|
from .hybrid_embed import HybridEmbed, HybridEmbedWithSize
|
|
from .inplace_abn import InplaceAbn
|
|
from .layer_scale import LayerScale, LayerScale2d
|
|
from .linear import Linear
|
|
from .mixed_conv2d import MixedConv2d
|
|
from .mlp import Mlp, GluMlp, GatedMlp, SwiGLU, SwiGLUPacked, ConvMlp, GlobalResponseNormMlp
|
|
from .non_local_attn import NonLocalAttn, BatNonLocalAttn
|
|
from .norm import (
|
|
GroupNorm,
|
|
GroupNorm1,
|
|
LayerNorm,
|
|
LayerNorm2d,
|
|
LayerNormFp32,
|
|
LayerNorm2dFp32,
|
|
RmsNorm,
|
|
RmsNorm2d,
|
|
RmsNormFp32,
|
|
RmsNorm2dFp32,
|
|
SimpleNorm,
|
|
SimpleNorm2d,
|
|
SimpleNormFp32,
|
|
SimpleNorm2dFp32,
|
|
)
|
|
from .norm_act import (
|
|
BatchNormAct2d,
|
|
GroupNormAct,
|
|
GroupNorm1Act,
|
|
LayerNormAct,
|
|
LayerNormAct2d,
|
|
LayerNormActFp32,
|
|
LayerNormAct2dFp32,
|
|
RmsNormAct,
|
|
RmsNormAct2d,
|
|
RmsNormActFp32,
|
|
RmsNormAct2dFp32,
|
|
SyncBatchNormAct,
|
|
convert_sync_batchnorm,
|
|
FrozenBatchNormAct2d,
|
|
freeze_batch_norm_2d,
|
|
unfreeze_batch_norm_2d,
|
|
)
|
|
from .padding import get_padding, get_same_padding, pad_same
|
|
from .patch_dropout import PatchDropout, PatchDropoutWithIndices, patch_dropout_forward
|
|
from .patch_embed import PatchEmbed, PatchEmbedWithSize, PatchEmbedInterpolator, resample_patch_embed
|
|
from .pool1d import global_pool_nlc
|
|
from .pool2d_same import AvgPool2dSame, create_pool2d
|
|
from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
|
|
from .pos_embed_rel import (
|
|
RelPosMlp,
|
|
RelPosBias,
|
|
RelPosBiasTf,
|
|
gen_relative_position_index,
|
|
gen_relative_log_coords,
|
|
resize_rel_pos_bias_table,
|
|
resize_rel_pos_bias_table_simple,
|
|
resize_rel_pos_bias_table_levit,
|
|
)
|
|
from .pos_embed_sincos import (
|
|
pixel_freq_bands,
|
|
freq_bands,
|
|
build_sincos2d_pos_embed,
|
|
build_fourier_pos_embed,
|
|
build_rotary_pos_embed,
|
|
apply_rot_embed,
|
|
apply_rot_embed_cat,
|
|
apply_rot_embed_list,
|
|
apply_keep_indices_nlc,
|
|
FourierEmbed,
|
|
RotaryEmbedding,
|
|
RotaryEmbeddingCat,
|
|
RotaryEmbeddingMixed,
|
|
RotaryEmbeddingDinoV3,
|
|
get_mixed_freqs,
|
|
create_rope_embed,
|
|
)
|
|
from .squeeze_excite import SEModule, SqueezeExcite, EffectiveSEModule, EffectiveSqueezeExcite
|
|
from .selective_kernel import SelectiveKernel
|
|
from .separable_conv import SeparableConv2d, SeparableConvNormAct
|
|
from .space_to_depth import SpaceToDepth, DepthToSpace
|
|
from .split_attn import SplitAttn
|
|
from .split_batchnorm import SplitBatchNorm2d, convert_splitbn_model
|
|
from .std_conv import StdConv2d, StdConv2dSame, ScaledStdConv2d, ScaledStdConv2dSame
|
|
from .test_time_pool import TestTimePoolHead, apply_test_time_pool
|
|
from .trace_utils import _assert, _float_to_int
|
|
from .typing import LayerType, PadType, disable_compiler
|
|
from .weight_init import (
|
|
trunc_normal_,
|
|
trunc_normal_tf_,
|
|
variance_scaling_,
|
|
lecun_normal_,
|
|
init_weight_jax,
|
|
init_weight_vit,
|
|
)
|