From 96d8e18ee0c2d09b1212bfa99d34161c1b3b3701 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:25:38 +0100 Subject: [PATCH 1/7] updated software version --- pyproject.toml | 2 +- setup.cfg | 2 +- src/aiia/__init__.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 129bf1a..e4d1f2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ include = '\.pyi?$' [project] name = "aiia" -version = "0.2.0" +version = "0.2.1" description = "AIIA Deep Learning Model Implementation" readme = "README.md" authors = [ diff --git a/setup.cfg b/setup.cfg index 198e211..eeb961d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = aiia -version = 0.2.0 +version = 0.2.1 author = Falko Habel author_email = falko.habel@gmx.de description = AIIA deep learning model implementation diff --git a/src/aiia/__init__.py b/src/aiia/__init__.py index cd5188b..6d716f8 100644 --- a/src/aiia/__init__.py +++ b/src/aiia/__init__.py @@ -1,7 +1,7 @@ -from .model.Model import AIIABase, AIIABaseShared, AIIAchunked, AIIAExpert, AIIAmoe, AIIA, AIIArecursive +from .model.Model import AIIABase, AIIABaseShared, AIIAchunked, AIIAExpert, AIIAmoe, AIIArecursive, AIIASparseMoe from .model.config import AIIAConfig from .data.DataLoader import DataLoader from .pretrain.pretrainer import Pretrainer, ProjectionHead -__version__ = "0.2.0" +__version__ = "0.2.1" -- 2.34.1 From 1fcb31b0443b0bc6651edad3f20d7ef70c0d7186 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:25:52 +0100 Subject: [PATCH 2/7] removed aiia from being able to get called --- src/aiia/model/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aiia/model/__init__.py b/src/aiia/model/__init__.py index f68a42a..c0ad6e7 100644 --- a/src/aiia/model/__init__.py +++ b/src/aiia/model/__init__.py @@ -1,21 +1,21 @@ from .Model import ( - AIIA, AIIABase, AIIABaseShared, AIIAchunked, AIIAExpert, AIIAmoe, - AIIArecursive + AIIArecursive, + AIIASparseMoe ) from .config import AIIAConfig __all__ = [ - "AIIA", "AIIABase", "AIIABaseShared", "AIIAchunked", "AIIAExpert", "AIIAmoe", "AIIArecursive", - "AIIAConfig" + "AIIAConfig", + "AIIASparseMoe" ] \ No newline at end of file -- 2.34.1 From 10967ea880563ac7ae4a2ca99cb65e87015b042e Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:26:00 +0100 Subject: [PATCH 3/7] added sparse moe --- src/aiia/model/Model.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src/aiia/model/Model.py b/src/aiia/model/Model.py index 06e2924..f21a0db 100644 --- a/src/aiia/model/Model.py +++ b/src/aiia/model/Model.py @@ -260,6 +260,40 @@ class AIIAmoe(AIIA): return merged_output +class AIIASparseMoe(AIIAmoe): + def __init__(self, config: AIIAConfig, num_experts: int = 3, top_k: int = 2, base_class=AIIABase, **kwargs): + super().__init__(config=config, num_experts=num_experts, base_class=base_class, **kwargs) + self.top_k = top_k + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Compute the gate_weights similar to standard moe. + expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1) + spatial_avg = torch.mean(expert_outputs, dim=(3, 4)) + gate_input = torch.mean(spatial_avg, dim=1) + gate_weights = self.gate(gate_input) + + # Select the top-k experts for each input based on gating weights. + _, top_k_indices = gate_weights.topk(self.top_k, dim=-1) + + # Initialize a list to store outputs from selected experts. + merged_outputs = [] + + # Iterate over batch dimension to apply top-k selection per instance. + for i in range(x.size(0)): + # Get the indices of top-k experts for current instance. + instance_top_k_indices = top_k_indices[i] + + # Select outputs from top-k experts. + selected_expert_outputs = expert_outputs[i][instance_top_k_indices] + + # Average over the selected experts to get a single output per instance. + averaged_output = torch.mean(selected_expert_outputs, dim=0) + merged_outputs.append(averaged_output.unsqueeze(0)) + + # Stack outputs from all instances back into a batch tensor. + return torch.cat(merged_outputs, dim=0) + + class AIIAchunked(AIIA): def __init__(self, config: AIIAConfig, patch_size: int = 16, base_class=AIIABase, **kwargs): super().__init__(config=config, **kwargs) -- 2.34.1 From c0e36cd5799b69484f401a22ed94dc515efd9880 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:26:14 +0100 Subject: [PATCH 4/7] added basic tests for sparse moe --- tests/model/test_aiia.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/model/test_aiia.py b/tests/model/test_aiia.py index b8904a2..340537b 100644 --- a/tests/model/test_aiia.py +++ b/tests/model/test_aiia.py @@ -1,6 +1,6 @@ import os import torch -from aiia import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig +from aiia import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig, AIIASparseMoe def test_aiiabase_creation(): config = AIIAConfig() @@ -106,6 +106,32 @@ def test_aiiamoe_save_load(): os.remove(os.path.join(save_path, "config.json")) os.rmdir(save_path) +def test_aiiasparsemoe_creation(): + config = AIIAConfig() + model = AIIASparseMoe(config, num_experts=5, top_k=2) + assert isinstance(model, AIIASparseMoe) + +def test_aiiasparsemoe_save_load(): + config = AIIAConfig() + model = AIIASparseMoe(config, num_experts=3, top_k=1) + save_path = "test_aiiasparsemoe_save_load" + + # Save the model + model.save(save_path) + assert os.path.exists(os.path.join(save_path, "model.pth")) + assert os.path.exists(os.path.join(save_path, "config.json")) + + # Load the model + loaded_model = AIIASparseMoe.load(save_path) + + # Check if the loaded model is an instance of AIIASparseMoe + assert isinstance(loaded_model, AIIASparseMoe) + + # Clean up + os.remove(os.path.join(save_path, "model.pth")) + os.remove(os.path.join(save_path, "config.json")) + os.rmdir(save_path) + def test_aiiachunked_creation(): config = AIIAConfig() model = AIIAchunked(config) -- 2.34.1 From 458d78a23e84563fdabdcd97647c17fc963567f5 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:26:30 +0100 Subject: [PATCH 5/7] removed full vectorizing --- .gitea/workflows/embed.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitea/workflows/embed.yaml b/.gitea/workflows/embed.yaml index 6804e2f..8036bc7 100644 --- a/.gitea/workflows/embed.yaml +++ b/.gitea/workflows/embed.yaml @@ -34,4 +34,4 @@ jobs: VECTORDB_TOKEN: ${{ secrets.VECTORDB_TOKEN }} run: | cd VectorLoader - python -m src.run --full + python -m src.run -- 2.34.1 From 5987a130f63945bf45f64208501ef8607425602b Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 26 Mar 2025 21:57:17 +0100 Subject: [PATCH 6/7] updated init to remove expeort model and move the smoe higher --- src/aiia/__init__.py | 2 +- src/aiia/model/__init__.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/aiia/__init__.py b/src/aiia/__init__.py index 6d716f8..838c765 100644 --- a/src/aiia/__init__.py +++ b/src/aiia/__init__.py @@ -1,4 +1,4 @@ -from .model.Model import AIIABase, AIIABaseShared, AIIAchunked, AIIAExpert, AIIAmoe, AIIArecursive, AIIASparseMoe +from .model.Model import AIIABase, AIIABaseShared, AIIAchunked, AIIAmoe, AIIASparseMoe, AIIArecursive from .model.config import AIIAConfig from .data.DataLoader import DataLoader from .pretrain.pretrainer import Pretrainer, ProjectionHead diff --git a/src/aiia/model/__init__.py b/src/aiia/model/__init__.py index c0ad6e7..a45512a 100644 --- a/src/aiia/model/__init__.py +++ b/src/aiia/model/__init__.py @@ -2,10 +2,9 @@ from .Model import ( AIIABase, AIIABaseShared, AIIAchunked, - AIIAExpert, AIIAmoe, - AIIArecursive, - AIIASparseMoe + AIIASparseMoe, + AIIArecursive ) from .config import AIIAConfig @@ -13,9 +12,9 @@ __all__ = [ "AIIABase", "AIIABaseShared", "AIIAchunked", - "AIIAExpert", "AIIAmoe", + "AIIASparseMoe", "AIIArecursive", "AIIAConfig", - "AIIASparseMoe" + ] \ No newline at end of file -- 2.34.1 From fb1b25507cbebc25f6563927d0e74a11d017fcb5 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Fri, 28 Mar 2025 16:10:22 +0100 Subject: [PATCH 7/7] added Expert to model to fix test --- tests/model/test_aiia.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/model/test_aiia.py b/tests/model/test_aiia.py index 340537b..bffa616 100644 --- a/tests/model/test_aiia.py +++ b/tests/model/test_aiia.py @@ -1,6 +1,6 @@ import os import torch -from aiia import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig, AIIASparseMoe +from aiia.model.Model import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig, AIIASparseMoe def test_aiiabase_creation(): config = AIIAConfig() -- 2.34.1