Upload directory

models/vit_kprpe/RPE/__init__.py

@@ -0,0 +1,48 @@
+from .KPRPE import kprpe_shared
+import torch
+import warnings
+import subprocess
+import sys
+import os
+
+try:
+ from .rpe_ops.rpe_index import RPEIndexFunction
+except ImportError:
+ try:
+ # Attempt to install the module from the setup.py script
+ dirname = os.path.dirname(os.path.abspath(__file__))
+ cwd = os.getcwd()
+ os.chdir(os.path.join(dirname, 'rpe_ops'))
+ subprocess.check_call([sys.executable, 'setup.py', 'install', '--user'])
+ GREEN_STR = "\033[92m{}\033[00m"
+ print(GREEN_STR.format("\n[INFO] Successfully installed `rpe_ops`. Restart Application"),)
+ sys.exit()
+ except subprocess.CalledProcessError as install_error:
+ RED_STR = "\033[91m{}\033[00m"
+ warnings.warn(RED_STR.format("\n[WARNING] Failed to install `rpe_ops`. "
+ "Please check the installation script."),)
+ except ImportError as import_error:
+ RED_STR = "\033[91m{}\033[00m"
+ warnings.warn(RED_STR.format("\n[WARNING] The module `rpe_ops` is not built. "
+ "For better training performance, please build `rpe_ops`."),)
+
+
+def build_rpe(rpe_config, head_dim, num_heads):
+ if rpe_config is None:
+ return None
+ else:
+ name = rpe_config.name
+ if name == 'KPRPE_shared':
+ rpe_config = kprpe_shared.get_rpe_config(
+ ratio=rpe_config.ratio,
+ method=rpe_config.method,
+ mode=rpe_config.mode,
+ shared_head=rpe_config.shared_head,
+ skip=0,
+ rpe_on=rpe_config.rpe_on,
+ )
+ return kprpe_shared.build_rpe(rpe_config, head_dim=head_dim, num_heads=num_heads)
+
+ else:
+ raise NotImplementedError(f"Unknow RPE: {name}")
+