parallel-metalearn 0.1.0__tar.gz

parallel_metalearn-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Ali Razani
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

parallel_metalearn-0.1.0/PKG-INFO

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: parallel-metalearn
+Version: 0.1.0
+Summary: A high-performance, vmap-accelerated functional meta-learning framework built natively on PyTorch.
+Author: Ali Razani
+Project-URL: Homepage, https://github.com/Razani-Ali/parallel-metalearn
+Project-URL: Bug Tracker, https://github.com/Razani-Ali/parallel-metalearn/issues
+Keywords: meta-learning,maml,anil,few-shot-learning,pytorch,vmap,functional-api,fault-diagnosis,deep-learning
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.5.0
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: tqdm>=4.60.0
+Dynamic: license-file
+
+
+
+# 🚀 Parallel-MetaLearn: Blazing-Fast, VMAP-Powered Functional Meta-Learning for PyTorch
+
+**Stop writing slow `for` loops over your meta-batches. Stop rewriting your PyTorch models into awkward functional syntax.**
+
+MetaLearn is a next-generation, high-performance meta-learning framework built natively on top of PyTorch 2.0+ `torch.func`. Designed for researchers and production engineers, it delivers **massive speedups** by vectorizing the outer-loop task processing while keeping your code clean, modular, and purely object-oriented.
+
+Whether you are doing Few-Shot Classification, Domain Adaptation, Semantic Segmentation, or Regression, MetaLearn adapts to your task—not the other way around.
+
+---
+
+## 🔥 Why Choose Parallel-MetaLearn? (The Game Changers)
+
+Existing libraries (like `learn2learn` or `higher`) force you into difficult compromises: they either use sequential `for` loops that bottleneck your GPU, or they require you to completely rewrite your model's forward pass to accept explicit parameters (e.g., `torch.functional.conv1d(x, weight=params['w'])`).
+
+**MetaLearn solves all of this:**
+
+* ⚡ **True Parallelism via `vmap`:** We eliminated the task `for` loop. By leveraging PyTorch's `vmap`, MetaLearn processes the entire meta-batch simultaneously. Expect speedups directly proportional to your task batch size (e.g., up to **Q-times faster** where Q is the number of tasks).
+* 🧠 **Zero-Friction Model Definitions:** Write your `nn.Module` exactly as you normally would. No need to pass parameter dictionaries into your `forward()` method. We handle the stateless functional calls completely under the hood.
+* 🎭 **Dynamic Task Imbalance & Masking:** `vmap` usually crashes if tasks have different batch sizes. We engineered a robust **Masking & Padding engine** under the hood. You can now train on highly imbalanced tasks (`support_shot=(min_shot,max_shot)`) without breaking vectorization!
+* 🎯 **Class-Agnostic & Class-Specific Modes:** Seamlessly switch between Class-Agnostic encoding (perfect for Out-Of-Distribution (OOD) generalization to unseen classes) and standard Class-Specific targets.
+* 🧩 **Task-Agnostic Architecture:** MetaLearn doesn't care if you are doing Classification, Regression, or Segmentation. Just swap out the Dataset and Loss classes. The core MAML remain 100% untouched.
+* ⏱️ **Step-Aware Inner Loop:** Your inner models and optimizers can be fully aware of the current gradient step, allowing for per-step learning rates and independent buffer management (crucial for MAML++).
+
+---
+
+## 🛠️ Supported Algorithms
+
+Currently, the library natively supports the most powerful gradient-based meta-learning algorithms out of the box:
+
+* ✅ **MAML** (Model-Agnostic Meta-Learning)
+* ✅ **FOMAML** (First Order MAML)
+* ✅ **ANIL** (Almost No Inner Loop)
+* ✅ **Meta-SGD** (Learnable inner learning rates)
+* ✅ **MAML++** (Multi-Step Loss Optimization & Per-step parameters)
+
+---
+
+## 📦 Core Features at a Glance
+
+* **Customizable Data Pipelines:** Use our highly flexible `MetaTaskDataset` to randomly or deterministically sample N-way K-shot tasks, or easily subclass it for your own custom data logic.
+* **Plug-and-Play Optimizers:** Build your own custom Inner-Optimizers effortlessly, and use any standard PyTorch optimizer (Adam, SGD, etc.) for the Outer-Loop.
+* **Automated Pipeline:** Say goodbye to boilerplate code. Our `MetaTrain` engine automatically handles the meta-training loop, validation intervals, metric logging, early stopping, and checkpoint saving.
+
+---
+
+## 🚀 Quick Start
+
+The complete pipeline works out of the box. Check out `main.py` for a fully working example on the CWRU Fault Diagnosis dataset. Here is how simple it is to initialize and train:
+
+```python
+import torch
+from metalearn.model_wrappers import MAML_Model
+from metalearn.loss import LabelEncoder, CrossEntropy, CategoricalAccuracy
+from metalearn.inner_optimizers import InnerSGD
+from metalearn.algorithms import MAML
+from metalearn.train import MetaTrain
+
+# 1. Define your standard PyTorch models (No functional rewrites needed!)
+backbone = MyCNNBackbone() 
+head = MyLinearHead()
+model = MAML_Model(backbone=backbone, head=head, drop_rate=0.5)
+
+# 2. Setup Class-Agnostic Encoding & Loss
+label_encoder = LabelEncoder(num_classes=10, max_n_way=3, shuffle=True)
+loss_fn = CrossEntropy(metric_fn=CategoricalAccuracy())
+
+# 3. Define Optimizers
+inner_optimizer = InnerSGD(initial_fast_weights=model.get_fast_weights(), inner_lr=0.01)
+outer_optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+# 4. Initialize Algorithm (MAML, ANIL, MAML++, etc.)
+algorithm = MAML(
+    model=model,
+    optimizer=outer_optimizer,
+    inner_optimizer=inner_optimizer,
+    support_loss_fn=loss_fn,
+    encoder=label_encoder,
+)
+
+# 5. Train with Automated Logging & Checkpointing!
+trainer = MetaTrain(
+    TrainLoader=train_loader, 
+    ValLoader=val_loader, 
+    algorithm=algorithm
+)
+
+history, best_metric, best_loss = trainer.train(
+    epochs=1500, 
+    check_idx=10, 
+    log_checkpoint_path="logs"
+)
+
+```
+
+---
+
+## 🗺️ Roadmap (Upcoming Features)
+
+We are constantly pushing the boundaries of what is possible in functional meta-learning. In our upcoming releases, look forward to:
+
+* **Recurrent Network Support:** Native, `vmap`-safe support for `LSTM` and `GRU` layers.
+* **Advanced Noise Management:** Robust meta-learning under input perturbations.
+* **New Meta-Algorithms:** Integration of cutting-edge algorithms (e.g., ProtoMAML, Reptile).
+* **New Inner Optimizers:** Second-order approximation optimizers and adaptive inner-loop schedulers.
+
+---
+
+## 🛠️ Unmatched Extensibility for Researchers (Developer Guide)
+
+MetaLearn is architected around **strict separation of concerns**. The core MAML execution engine operates purely on standardized output dictionaries (`out_dict`) and target dictionaries (`targets`). This means you can extend MetaLearn to cutting-edge research paradigms **without ever touching the core MAML execution loop or `vmap` logic**:
+
+### 1. 🔀 Multi-Task Learning (MTL)
+Need joint classification and auxiliary regression/reconstruction?
+* **Data:** Return auxiliary targets alongside labels in `MetaTaskDataset` (e.g., `y_dict = {"labels": y, "reg_targets": reg_y}`).
+* **Loss:** Subclass `BaseLoss` to compute composite loss (`cls_loss + lambda * reg_loss`).
+* *MAML engine automatically propagates gradients across all tasks!*
+
+### 2. 🌐 Meta-Domain Adaptation (MDA)
+Want to align feature distributions across shifting domains?
+* **Data:** Pass domain indicators inside your dataset targets (e.g., `y_dict = {"labels": y, "domain_id": d}`).
+* **Loss:** Extract features from `out_dict["features"]` and compute domain alignment loss (e.g., MMD, Wasserstein Distance, or Adversarial Loss) inside your custom Loss class.
+
+### 3. 🌐 Federated Meta-Learning (FedMeta)
+Want to simulate decentralized client adaptation or privacy-preserving meta-learning?
+* **Data & Algorithm:** Keep the same functional `MAML` step, but customize the task assignment logic to simulate client-side local updates before global aggregation.
+
+### 4. 🌐 Federated Learning (FedAvg, FedGrad, FedProx) & FedMeta
+Because MetaLearn processes inner-loop updates in a stateless, functional manner, you can effortlessly simulate **Pure Federated Learning algorithms** (e.g., FedAvg, FedGrad) alongside **Federated Meta-Learning (FedMeta)**:
+* **Parallel Client Simulation via `vmap`:** Instead of sequentially looping through individual clients, MetaLearn simulates dozens of local client updates *simultaneously* on the GPU using `vmap`.
+* **Zero-Overhead Aggregation:** Extract adapted local parameters $\theta_i'$ from each client task, perform global server aggregation (e.g., weighted averaging via `torch.stack(client_weights).mean(dim=0)`), and seamlessly set the new global start state for the next communication round.
+
+---
+
+## 🤝 Contributing & Citation
+
+If you use MetaLearn in your research or production pipelines, we'd love to hear about it! Contributions, issues, and feature requests are always welcome.
+
+> *Fully functional example available in `main.py`.* Just run `python main.py` and watch the `vmap` magic happen!

parallel_metalearn-0.1.0/README.md

@@ -0,0 +1,140 @@
+
+
+# 🚀 Parallel-MetaLearn: Blazing-Fast, VMAP-Powered Functional Meta-Learning for PyTorch
+
+**Stop writing slow `for` loops over your meta-batches. Stop rewriting your PyTorch models into awkward functional syntax.**
+
+MetaLearn is a next-generation, high-performance meta-learning framework built natively on top of PyTorch 2.0+ `torch.func`. Designed for researchers and production engineers, it delivers **massive speedups** by vectorizing the outer-loop task processing while keeping your code clean, modular, and purely object-oriented.
+
+Whether you are doing Few-Shot Classification, Domain Adaptation, Semantic Segmentation, or Regression, MetaLearn adapts to your task—not the other way around.
+
+---
+
+## 🔥 Why Choose Parallel-MetaLearn? (The Game Changers)
+
+Existing libraries (like `learn2learn` or `higher`) force you into difficult compromises: they either use sequential `for` loops that bottleneck your GPU, or they require you to completely rewrite your model's forward pass to accept explicit parameters (e.g., `torch.functional.conv1d(x, weight=params['w'])`).
+
+**MetaLearn solves all of this:**
+
+* ⚡ **True Parallelism via `vmap`:** We eliminated the task `for` loop. By leveraging PyTorch's `vmap`, MetaLearn processes the entire meta-batch simultaneously. Expect speedups directly proportional to your task batch size (e.g., up to **Q-times faster** where Q is the number of tasks).
+* 🧠 **Zero-Friction Model Definitions:** Write your `nn.Module` exactly as you normally would. No need to pass parameter dictionaries into your `forward()` method. We handle the stateless functional calls completely under the hood.
+* 🎭 **Dynamic Task Imbalance & Masking:** `vmap` usually crashes if tasks have different batch sizes. We engineered a robust **Masking & Padding engine** under the hood. You can now train on highly imbalanced tasks (`support_shot=(min_shot,max_shot)`) without breaking vectorization!
+* 🎯 **Class-Agnostic & Class-Specific Modes:** Seamlessly switch between Class-Agnostic encoding (perfect for Out-Of-Distribution (OOD) generalization to unseen classes) and standard Class-Specific targets.
+* 🧩 **Task-Agnostic Architecture:** MetaLearn doesn't care if you are doing Classification, Regression, or Segmentation. Just swap out the Dataset and Loss classes. The core MAML remain 100% untouched.
+* ⏱️ **Step-Aware Inner Loop:** Your inner models and optimizers can be fully aware of the current gradient step, allowing for per-step learning rates and independent buffer management (crucial for MAML++).
+
+---
+
+## 🛠️ Supported Algorithms
+
+Currently, the library natively supports the most powerful gradient-based meta-learning algorithms out of the box:
+
+* ✅ **MAML** (Model-Agnostic Meta-Learning)
+* ✅ **FOMAML** (First Order MAML)
+* ✅ **ANIL** (Almost No Inner Loop)
+* ✅ **Meta-SGD** (Learnable inner learning rates)
+* ✅ **MAML++** (Multi-Step Loss Optimization & Per-step parameters)
+
+---
+
+## 📦 Core Features at a Glance
+
+* **Customizable Data Pipelines:** Use our highly flexible `MetaTaskDataset` to randomly or deterministically sample N-way K-shot tasks, or easily subclass it for your own custom data logic.
+* **Plug-and-Play Optimizers:** Build your own custom Inner-Optimizers effortlessly, and use any standard PyTorch optimizer (Adam, SGD, etc.) for the Outer-Loop.
+* **Automated Pipeline:** Say goodbye to boilerplate code. Our `MetaTrain` engine automatically handles the meta-training loop, validation intervals, metric logging, early stopping, and checkpoint saving.
+
+---
+
+## 🚀 Quick Start
+
+The complete pipeline works out of the box. Check out `main.py` for a fully working example on the CWRU Fault Diagnosis dataset. Here is how simple it is to initialize and train:
+
+```python
+import torch
+from metalearn.model_wrappers import MAML_Model
+from metalearn.loss import LabelEncoder, CrossEntropy, CategoricalAccuracy
+from metalearn.inner_optimizers import InnerSGD
+from metalearn.algorithms import MAML
+from metalearn.train import MetaTrain
+
+# 1. Define your standard PyTorch models (No functional rewrites needed!)
+backbone = MyCNNBackbone() 
+head = MyLinearHead()
+model = MAML_Model(backbone=backbone, head=head, drop_rate=0.5)
+
+# 2. Setup Class-Agnostic Encoding & Loss
+label_encoder = LabelEncoder(num_classes=10, max_n_way=3, shuffle=True)
+loss_fn = CrossEntropy(metric_fn=CategoricalAccuracy())
+
+# 3. Define Optimizers
+inner_optimizer = InnerSGD(initial_fast_weights=model.get_fast_weights(), inner_lr=0.01)
+outer_optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+# 4. Initialize Algorithm (MAML, ANIL, MAML++, etc.)
+algorithm = MAML(
+    model=model,
+    optimizer=outer_optimizer,
+    inner_optimizer=inner_optimizer,
+    support_loss_fn=loss_fn,
+    encoder=label_encoder,
+)
+
+# 5. Train with Automated Logging & Checkpointing!
+trainer = MetaTrain(
+    TrainLoader=train_loader, 
+    ValLoader=val_loader, 
+    algorithm=algorithm
+)
+
+history, best_metric, best_loss = trainer.train(
+    epochs=1500, 
+    check_idx=10, 
+    log_checkpoint_path="logs"
+)
+
+```
+
+---
+
+## 🗺️ Roadmap (Upcoming Features)
+
+We are constantly pushing the boundaries of what is possible in functional meta-learning. In our upcoming releases, look forward to:
+
+* **Recurrent Network Support:** Native, `vmap`-safe support for `LSTM` and `GRU` layers.
+* **Advanced Noise Management:** Robust meta-learning under input perturbations.
+* **New Meta-Algorithms:** Integration of cutting-edge algorithms (e.g., ProtoMAML, Reptile).
+* **New Inner Optimizers:** Second-order approximation optimizers and adaptive inner-loop schedulers.
+
+---
+
+## 🛠️ Unmatched Extensibility for Researchers (Developer Guide)
+
+MetaLearn is architected around **strict separation of concerns**. The core MAML execution engine operates purely on standardized output dictionaries (`out_dict`) and target dictionaries (`targets`). This means you can extend MetaLearn to cutting-edge research paradigms **without ever touching the core MAML execution loop or `vmap` logic**:
+
+### 1. 🔀 Multi-Task Learning (MTL)
+Need joint classification and auxiliary regression/reconstruction?
+* **Data:** Return auxiliary targets alongside labels in `MetaTaskDataset` (e.g., `y_dict = {"labels": y, "reg_targets": reg_y}`).
+* **Loss:** Subclass `BaseLoss` to compute composite loss (`cls_loss + lambda * reg_loss`).
+* *MAML engine automatically propagates gradients across all tasks!*
+
+### 2. 🌐 Meta-Domain Adaptation (MDA)
+Want to align feature distributions across shifting domains?
+* **Data:** Pass domain indicators inside your dataset targets (e.g., `y_dict = {"labels": y, "domain_id": d}`).
+* **Loss:** Extract features from `out_dict["features"]` and compute domain alignment loss (e.g., MMD, Wasserstein Distance, or Adversarial Loss) inside your custom Loss class.
+
+### 3. 🌐 Federated Meta-Learning (FedMeta)
+Want to simulate decentralized client adaptation or privacy-preserving meta-learning?
+* **Data & Algorithm:** Keep the same functional `MAML` step, but customize the task assignment logic to simulate client-side local updates before global aggregation.
+
+### 4. 🌐 Federated Learning (FedAvg, FedGrad, FedProx) & FedMeta
+Because MetaLearn processes inner-loop updates in a stateless, functional manner, you can effortlessly simulate **Pure Federated Learning algorithms** (e.g., FedAvg, FedGrad) alongside **Federated Meta-Learning (FedMeta)**:
+* **Parallel Client Simulation via `vmap`:** Instead of sequentially looping through individual clients, MetaLearn simulates dozens of local client updates *simultaneously* on the GPU using `vmap`.
+* **Zero-Overhead Aggregation:** Extract adapted local parameters $\theta_i'$ from each client task, perform global server aggregation (e.g., weighted averaging via `torch.stack(client_weights).mean(dim=0)`), and seamlessly set the new global start state for the next communication round.
+
+---
+
+## 🤝 Contributing & Citation
+
+If you use MetaLearn in your research or production pipelines, we'd love to hear about it! Contributions, issues, and feature requests are always welcome.
+
+> *Fully functional example available in `main.py`.* Just run `python main.py` and watch the `vmap` magic happen!

parallel_metalearn-0.1.0/metalearn/algorithms/BaseLearner.py

@@ -0,0 +1,57 @@
+from abc import ABC, abstractmethod
+from typing import Tuple, Any, Dict, Optional
+import torch
+import torch.nn as nn
+
+
+class MetaOptimizer(nn.Module, ABC):
+    """
+    Abstract Base Class for all Meta-Learning Algorithms (e.g., MAML, ProtoNet).
+    
+    Establishes a unified interface for meta-training/evaluation steps 
+    and test-time adaptation/deployment.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    @abstractmethod
+    def step(
+        self, 
+        task_loader: Any, 
+        training: bool = True, 
+        **kwargs: Any
+    ) -> Tuple[float, float]:
+        """
+        Executes a single meta-training or meta-validation step over a batch of tasks.
+
+        Args:
+            task_loader: Dataloader/Iterable yielding batches of tasks 
+                         (typically yielding Xs, Ys, Xq, Yq).
+            training (bool): If True, computes meta-gradients and updates outer model parameters.
+            **kwargs: Additional contextual arguments (e.g., epoch, inner_steps override).
+
+        Returns:
+            Tuple[float, float]: A tuple containing:
+                - mean_meta_loss (float): Average meta-loss value across the batch of tasks.
+                - mean_metric (float): Average evaluation metric (e.g., accuracy) across tasks.
+        """
+        pass
+
+    @abstractmethod
+    def adapt_and_update(
+        self, 
+        Xsupport: torch.Tensor, 
+        Ysupport: torch.Tensor, 
+        **kwargs: Any
+    ) -> None:
+        """
+        Adapts parameters on the provided support set and permanently updates the internal model.
+        Intended exclusively for test-time adaptation, fine-tuning, or deployment.
+
+        Args:
+            Xsupport (torch.Tensor): Support set inputs of shape (num_tasks, K_support, ...).
+            Ysupport (torch.Tensor): Support set targets of shape (num_tasks, K_support, ...).
+            **kwargs: Additional context arguments (e.g., custom inner adaptation steps).
+        """
+        pass