__init__.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. from ._fx import (
  2. create_feature_extractor,
  3. get_graph_node_names,
  4. register_notrace_function,
  5. register_notrace_module,
  6. is_notrace_module,
  7. is_notrace_function,
  8. get_notrace_modules,
  9. get_notrace_functions,
  10. )
  11. from .activations import *
  12. from .adaptive_avgmax_pool import (
  13. adaptive_avgmax_pool2d,
  14. select_adaptive_pool2d,
  15. AdaptiveAvgMaxPool2d,
  16. SelectAdaptivePool2d,
  17. )
  18. from .attention import Attention, AttentionRope, maybe_add_mask
  19. from .attention2d import MultiQueryAttention2d, Attention2d, MultiQueryAttentionV2
  20. from .attention_pool import AttentionPoolLatent
  21. from .attention_pool2d import AttentionPool2d, RotAttentionPool2d, RotaryEmbedding
  22. from .blur_pool import BlurPool2d, create_aa
  23. from .classifier import create_classifier, ClassifierHead, NormMlpClassifierHead, ClNormMlpClassifierHead
  24. from .cond_conv2d import CondConv2d, get_condconv_initializer
  25. from .config import (
  26. is_exportable,
  27. is_scriptable,
  28. is_no_jit,
  29. use_fused_attn,
  30. set_exportable,
  31. set_scriptable,
  32. set_no_jit,
  33. set_layer_config,
  34. set_fused_attn,
  35. set_reentrant_ckpt,
  36. use_reentrant_ckpt,
  37. )
  38. from .conv2d_same import Conv2dSame, conv2d_same
  39. from .conv_bn_act import ConvNormAct, ConvNormActAa, ConvBnAct
  40. from .create_act import create_act_layer, get_act_layer, get_act_fn
  41. from .create_attn import get_attn, create_attn
  42. from .create_conv2d import create_conv2d
  43. from .create_norm import get_norm_layer, create_norm_layer
  44. from .create_norm_act import get_norm_act_layer, create_norm_act_layer, get_norm_act_layer
  45. from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path, calculate_drop_path_rates
  46. from .eca import EcaModule, CecaModule, EfficientChannelAttn, CircularEfficientChannelAttn
  47. from .evo_norm import (
  48. EvoNorm2dB0,
  49. EvoNorm2dB1,
  50. EvoNorm2dB2,
  51. EvoNorm2dS0,
  52. EvoNorm2dS0a,
  53. EvoNorm2dS1,
  54. EvoNorm2dS1a,
  55. EvoNorm2dS2,
  56. EvoNorm2dS2a,
  57. )
  58. from .fast_norm import is_fast_norm, set_fast_norm, fast_group_norm, fast_layer_norm
  59. from .filter_response_norm import FilterResponseNormTlu2d, FilterResponseNormAct2d
  60. from .format import Format, get_channel_dim, get_spatial_dim, nchw_to, nhwc_to
  61. from .gather_excite import GatherExcite
  62. from .global_context import GlobalContext
  63. from .grid import ndgrid, meshgrid
  64. from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple, make_divisible, extend_tuple
  65. from .hybrid_embed import HybridEmbed, HybridEmbedWithSize
  66. from .inplace_abn import InplaceAbn
  67. from .layer_scale import LayerScale, LayerScale2d
  68. from .linear import Linear
  69. from .mixed_conv2d import MixedConv2d
  70. from .mlp import Mlp, GluMlp, GatedMlp, SwiGLU, SwiGLUPacked, ConvMlp, GlobalResponseNormMlp
  71. from .non_local_attn import NonLocalAttn, BatNonLocalAttn
  72. from .norm import (
  73. GroupNorm,
  74. GroupNorm1,
  75. LayerNorm,
  76. LayerNorm2d,
  77. LayerNormFp32,
  78. LayerNorm2dFp32,
  79. RmsNorm,
  80. RmsNorm2d,
  81. RmsNormFp32,
  82. RmsNorm2dFp32,
  83. SimpleNorm,
  84. SimpleNorm2d,
  85. SimpleNormFp32,
  86. SimpleNorm2dFp32,
  87. )
  88. from .norm_act import (
  89. BatchNormAct2d,
  90. GroupNormAct,
  91. GroupNorm1Act,
  92. LayerNormAct,
  93. LayerNormAct2d,
  94. LayerNormActFp32,
  95. LayerNormAct2dFp32,
  96. RmsNormAct,
  97. RmsNormAct2d,
  98. RmsNormActFp32,
  99. RmsNormAct2dFp32,
  100. SyncBatchNormAct,
  101. convert_sync_batchnorm,
  102. FrozenBatchNormAct2d,
  103. freeze_batch_norm_2d,
  104. unfreeze_batch_norm_2d,
  105. )
  106. from .padding import get_padding, get_same_padding, pad_same
  107. from .patch_dropout import PatchDropout, PatchDropoutWithIndices, patch_dropout_forward
  108. from .patch_embed import PatchEmbed, PatchEmbedWithSize, PatchEmbedInterpolator, resample_patch_embed
  109. from .pool1d import global_pool_nlc
  110. from .pool2d_same import AvgPool2dSame, create_pool2d
  111. from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
  112. from .pos_embed_rel import (
  113. RelPosMlp,
  114. RelPosBias,
  115. RelPosBiasTf,
  116. gen_relative_position_index,
  117. gen_relative_log_coords,
  118. resize_rel_pos_bias_table,
  119. resize_rel_pos_bias_table_simple,
  120. resize_rel_pos_bias_table_levit,
  121. )
  122. from .pos_embed_sincos import (
  123. pixel_freq_bands,
  124. freq_bands,
  125. build_sincos2d_pos_embed,
  126. build_fourier_pos_embed,
  127. build_rotary_pos_embed,
  128. apply_rot_embed,
  129. apply_rot_embed_cat,
  130. apply_rot_embed_list,
  131. apply_keep_indices_nlc,
  132. FourierEmbed,
  133. RotaryEmbedding,
  134. RotaryEmbeddingCat,
  135. RotaryEmbeddingMixed,
  136. RotaryEmbeddingDinoV3,
  137. get_mixed_freqs,
  138. create_rope_embed,
  139. )
  140. from .squeeze_excite import SEModule, SqueezeExcite, EffectiveSEModule, EffectiveSqueezeExcite
  141. from .selective_kernel import SelectiveKernel
  142. from .separable_conv import SeparableConv2d, SeparableConvNormAct
  143. from .space_to_depth import SpaceToDepth, DepthToSpace
  144. from .split_attn import SplitAttn
  145. from .split_batchnorm import SplitBatchNorm2d, convert_splitbn_model
  146. from .std_conv import StdConv2d, StdConv2dSame, ScaledStdConv2d, ScaledStdConv2dSame
  147. from .test_time_pool import TestTimePoolHead, apply_test_time_pool
  148. from .trace_utils import _assert, _float_to_int
  149. from .typing import LayerType, PadType, disable_compiler
  150. from .weight_init import (
  151. trunc_normal_,
  152. trunc_normal_tf_,
  153. variance_scaling_,
  154. lecun_normal_,
  155. init_weight_jax,
  156. init_weight_vit,
  157. )