homa 0.0.19__tar.gz → 0.1.94__tar.gz

homa-0.1.94/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.4
+Name: homa
+Version: 0.1.94
+Summary: A curated list of machine learning and deep learning helpers.
+Author-email: Taha Shieenavaz <tahashieenavaz@gmail.com>
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: fire
+
+# Core
+
+### Device Management
+
+```py
+from homa import cpu, mps, cuda, device
+
+torch.tensor([1, 2, 3, 4, 5]).to(cpu())
+torch.tensor([1, 2, 3, 4, 5]).to(cuda())
+torch.tensor([1, 2, 3, 4, 5]).to(mps())
+torch.tensor([1, 2, 3, 4, 5]).to(device())
+```
+
+# Vision
+
+## Resnet
+
+This is the standard ResNet50 module.
+
+You can train the model with a `DataLoader` object.
+
+```py
+from homa.vision import Resnet
+
+model = Resnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    model.train(train_dataloader)
+```
+
+Similarly you can manually take care of decomposition of data from the `DataLoader`.
+
+```py
+from homa.vision import Resnet
+
+model = Resnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    for x, y in train_dataloader:
+        model.train(x, y)
+```
+
+## StochasticResnet
+
+This is a ResNet module whose activation functions are replaced from a pool of different activation functions randomly. Read more on the [(paper)](https://www.mdpi.com/1424-8220/22/16/6129).
+
+You can train the model with a `DataLoader` object.
+
+```py
+from homa.vision import StochasticResnet
+
+model = StochasticResnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    model.train(train_dataloader)
+```
+
+Similarly you can manually take care of decomposition of data from the `DataLoader`.
+
+```py
+from homa.vision import StochasticResnet
+
+model = StochasticResnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    for x, y in train_dataloader:
+        model.train(x, y)
+```

homa-0.1.94/README.md

@@ -0,0 +1,64 @@
+# Core
+
+### Device Management
+
+```py
+from homa import cpu, mps, cuda, device
+
+torch.tensor([1, 2, 3, 4, 5]).to(cpu())
+torch.tensor([1, 2, 3, 4, 5]).to(cuda())
+torch.tensor([1, 2, 3, 4, 5]).to(mps())
+torch.tensor([1, 2, 3, 4, 5]).to(device())
+```
+
+# Vision
+
+## Resnet
+
+This is the standard ResNet50 module.
+
+You can train the model with a `DataLoader` object.
+
+```py
+from homa.vision import Resnet
+
+model = Resnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    model.train(train_dataloader)
+```
+
+Similarly you can manually take care of decomposition of data from the `DataLoader`.
+
+```py
+from homa.vision import Resnet
+
+model = Resnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    for x, y in train_dataloader:
+        model.train(x, y)
+```
+
+## StochasticResnet
+
+This is a ResNet module whose activation functions are replaced from a pool of different activation functions randomly. Read more on the [(paper)](https://www.mdpi.com/1424-8220/22/16/6129).
+
+You can train the model with a `DataLoader` object.
+
+```py
+from homa.vision import StochasticResnet
+
+model = StochasticResnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    model.train(train_dataloader)
+```
+
+Similarly you can manually take care of decomposition of data from the `DataLoader`.
+
+```py
+from homa.vision import StochasticResnet
+
+model = StochasticResnet(num_classes=10, lr=0.001)
+for epoch in range(10):
+    for x, y in train_dataloader:
+        model.train(x, y)
+```

{homa-0.0.19 → homa-0.1.94}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "homa"
-version = "0.0.19"
+version = "0.1.94"
 description = "A curated list of machine learning and deep learning helpers."
 authors = [
   { name="Taha Shieenavaz", email="tahashieenavaz@gmail.com" },

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/APLU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class APLU(torch.nn.Module):
@@ -12,6 +13,7 @@ class APLU(torch.nn.Module):
         self.psi = None
         self.mu = None
         self._num_channels = None
+        self.device = get_device()
 
     def _initialize_parameters(self, x):
         if x.ndim < 2:
@@ -25,18 +27,23 @@ class APLU(torch.nn.Module):
         param_shape = [1] * x.ndim
         param_shape[1] = num_channels
 
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(param_shape))
+        self.alpha = torch.nn.Parameter(torch.zeros(param_shape), device=device())
+        self.beta = torch.nn.Parameter(torch.zeros(param_shape), device=device())
+        self.gamma = torch.nn.Parameter(torch.zeros(param_shape), device=device())
 
-        self.xi = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
-        self.psi = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
-        self.mu = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
+        self.xi = torch.nn.Parameter(self.max_input * torch.rand(param_shape)).to(
+            self.device
+        )
+        self.psi = torch.nn.Parameter(self.max_input * torch.rand(param_shape)).to(
+            self.device
+        )
+        self.mu = torch.nn.Parameter(self.max_input * torch.rand(param_shape)).to(
+            self.device
+        )
 
     def forward(self, x):
         if self.alpha is None:
             self._initialize_parameters(x)
-
         a = torch.relu(x)
 
         # following are called hinges
@@ -44,5 +51,4 @@ class APLU(torch.nn.Module):
         c = self.beta * torch.relu(-x + self.psi)
         d = self.gamma * torch.relu(-x + self.mu)
         z = a + b + c + d
-
         return z

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/GALU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class GALU(torch.nn.Module):
@@ -12,6 +13,7 @@ class GALU(torch.nn.Module):
         self.gamma = None
         self.delta = None
         self._num_channels = None
+        self.device = get_device()
 
     def _initialize_parameters(self, x):
         if x.ndim < 2:
@@ -23,10 +25,10 @@ class GALU(torch.nn.Module):
         self._num_channels = num_channels
         param_shape = [1] * x.ndim
         param_shape[1] = num_channels
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(param_shape))
-        self.delta = torch.nn.Parameter(torch.zeros(param_shape))
+        self.alpha = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
+        self.beta = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
+        self.gamma = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
+        self.delta = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
 
     def forward(self, x):
         if self.alpha is None:

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/MELU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class MELU(torch.nn.Module):
@@ -12,6 +13,7 @@ class MELU(torch.nn.Module):
         self.xi = None
         self.psi = None
         self._initialized = False
+        self.device = get_device()
 
     def _initialize_parameters(self, X: torch.Tensor):
         if X.dim() != 4:
@@ -20,12 +22,12 @@ class MELU(torch.nn.Module):
             )
         num_channels = X.shape[1]
         shape = (1, num_channels, 1, 1)
-        self.alpha = torch.nn.Parameter(torch.zeros(shape))
-        self.beta = torch.nn.Parameter(torch.zeros(shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(shape))
-        self.delta = torch.nn.Parameter(torch.zeros(shape))
-        self.xi = torch.nn.Parameter(torch.zeros(shape))
-        self.psi = torch.nn.Parameter(torch.zeros(shape))
+        self.alpha = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.beta = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.gamma = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.delta = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.xi = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.psi = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
         self._initialized = True
 
     def forward(self, X: torch.Tensor) -> torch.Tensor:

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/PDELU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class PDELU(torch.nn.Module):
@@ -12,6 +13,7 @@ class PDELU(torch.nn.Module):
         self._power_val = 1.0 / (1.0 - self.theta)
         self.alpha = torch.nn.UninitializedParameter()
         self._num_channels = None
+        self.device = get_device()
 
     def _initialize_parameters(self, x: torch.Tensor):
         if x.ndim < 2:
@@ -23,14 +25,14 @@ class PDELU(torch.nn.Module):
         self._num_channels = num_channels
         param_shape = [1] * x.ndim
         param_shape[1] = num_channels
-        init_tensor = torch.zeros(param_shape) + 0.1
-        self.alpha = torch.nn.Parameter(init_tensor)
+        init_tensor = torch.zeros(param_shape, device=self.device) + 0.1
+        self.alpha = torch.nn.Parameter(init_tensor).to(self.device)
 
     def forward(self, x: torch.Tensor):
         if self.alpha is None:
             self._initialize_parameters(x)
 
-        zero = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+        zero = torch.tensor(0.0, device=self.device, dtype=x.dtype)
         positive_part = torch.relu(x)
         inner_term = torch.relu(1.0 + (1.0 - self.theta) * x)
         powered_term = torch.pow(inner_term, self._power_val)

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/SReLU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class SReLU(torch.nn.Module):
@@ -18,6 +19,7 @@ class SReLU(torch.nn.Module):
         self.beta = torch.nn.UninitializedParameter()
         self.gamma = torch.nn.UninitializedParameter()
         self.delta = torch.nn.UninitializedParameter()
+        self.device = get_device()
 
     def _initialize_parameters(self, x: torch.Tensor):
         if isinstance(self.alpha, torch.nn.UninitializedParameter):
@@ -31,14 +33,16 @@ class SReLU(torch.nn.Module):
             param_shape[1] = num_channels
             self.alpha = torch.nn.Parameter(
                 torch.full(param_shape, self.alpha_init_val)
-            )
-            self.beta = torch.nn.Parameter(torch.full(param_shape, self.beta_init_val))
+            ).to(self.device)
+            self.beta = torch.nn.Parameter(
+                torch.full(param_shape, self.beta_init_val)
+            ).to(self.device)
             self.gamma = torch.nn.Parameter(
                 torch.full(param_shape, self.gamma_init_val)
-            )
+            ).to(self.device)
             self.delta = torch.nn.Parameter(
                 torch.full(param_shape, self.delta_init_val)
-            )
+            ).to(self.device)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         self._initialize_parameters(x)

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/SmallGALU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class SmallGALU(torch.nn.Module):
@@ -10,6 +11,7 @@ class SmallGALU(torch.nn.Module):
         self.alpha = None
         self.beta = None
         self._num_channels = None
+        self.device = get_device()
 
     def _initialize_parameters(self, x):
         if x.ndim < 2:
@@ -21,8 +23,8 @@ class SmallGALU(torch.nn.Module):
         self._num_channels = num_channels
         param_shape = [1] * x.ndim
         param_shape[1] = num_channels
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
+        self.alpha = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
+        self.beta = torch.nn.Parameter(torch.zeros(param_shape)).to(self.device)
 
     def forward(self, x):
         if self.alpha is None:

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/StochasticActivation.py

@@ -13,7 +13,21 @@ from .SReLU import SReLU
 class StochasticActivation(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.gate = random.choice([APLU, GALU, SmallGALU, MELU, WideMELU, PDELU, SReLU])
+        self.gate = random.choice(
+            [
+                APLU,
+                GALU,
+                SmallGALU,
+                MELU,
+                WideMELU,
+                PDELU,
+                SReLU,
+                torch.nn.ReLU,
+                torch.nn.PReLU,
+                torch.nn.LeakyReLU,
+                torch.nn.ELU,
+            ]
+        )
         self.gate = self.gate()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:

{homa-0.0.19 → homa-0.1.94}/src/homa/activations/classes/WideMELU.py

@@ -1,4 +1,5 @@
 import torch
+from ...device import get_device
 
 
 class WideMELU(torch.nn.Module):
@@ -14,6 +15,7 @@ class WideMELU(torch.nn.Module):
         self.theta = None
         self.lam = None
         self._initialized = False
+        self.device = get_device()
 
     def _initialize_parameters(self, X: torch.Tensor):
         if X.dim() != 4:
@@ -24,14 +26,14 @@ class WideMELU(torch.nn.Module):
         num_channels = X.shape[1]
         shape = (1, num_channels, 1, 1)
 
-        self.alpha = torch.nn.Parameter(torch.zeros(shape))
-        self.beta = torch.nn.Parameter(torch.zeros(shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(shape))
-        self.delta = torch.nn.Parameter(torch.zeros(shape))
-        self.xi = torch.nn.Parameter(torch.zeros(shape))
-        self.psi = torch.nn.Parameter(torch.zeros(shape))
-        self.theta = torch.nn.Parameter(torch.zeros(shape))
-        self.lam = torch.nn.Parameter(torch.zeros(shape))
+        self.alpha = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.beta = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.gamma = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.delta = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.xi = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.psi = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.theta = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
+        self.lam = torch.nn.Parameter(torch.zeros(shape)).to(self.device)
         self._initialized = True
 
     def forward(self, X: torch.Tensor) -> torch.Tensor:

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/Ensemble.py

@@ -1,8 +1,7 @@
 from .concerns import (
     ReportsSize,
-    RecordsStateDictionaries,
+    StoresModels,
     ReportsClassificationMetrics,
-    HasNetwork,
     PredictsProbabilities,
 )
 
@@ -10,9 +9,8 @@ from .concerns import (
 class Ensemble(
     ReportsSize,
     ReportsClassificationMetrics,
-    RecordsStateDictionaries,
     PredictsProbabilities,
-    HasNetwork,
+    StoresModels,
 ):
     def __init__(self):
         super().__init__()

homa-0.1.94/src/homa/ensemble/concerns/CalculatesMetricNecessities.py

@@ -0,0 +1,24 @@
+import torch
+from ...device import get_device
+
+
+class CalculatesMetricNecessities:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @torch.no_grad()
+    def metric_necessities(self, dataloader):
+        predictions, labels = [], []
+        device = get_device()
+        for x, y in dataloader:
+            x, y = x.to(device), y.to(device)
+            sum_logits = None
+            for model in self.models:
+                model.to(device)
+                model.eval()
+                logits = model(x)
+                sum_logits = logits if sum_logits is None else sum_logits + logits
+            batch_predictions = sum_logits.argmax(dim=1)
+            predictions.extend(batch_predictions.cpu().numpy())
+            labels.extend(y.cpu().numpy())
+        return predictions, labels

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/concerns/PredictsProbabilities.py

@@ -9,3 +9,7 @@ class PredictsProbabilities(ReportsLogits):
     def predict(self, x: torch.Tensor) -> torch.Tensor:
         logits = self.logits(x)
         return torch.nn.functional.softmax(logits, dim=1)
+
+    @torch.no_grad()
+    def predict_(self, x: torch.Tensor) -> torch.Tensor:
+        return self.predict(x)

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/concerns/ReportsEnsembleAccuracy.py

@@ -1,10 +1,11 @@
 from sklearn.metrics import accuracy_score as accuracy
+from torch.utils.data import DataLoader
 
 
 class ReportsEnsembleAccuracy:
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def accuracy(self) -> float:
-        predictions, labels = self.metric_necessities()
+    def accuracy(self, dataloader: DataLoader) -> float:
+        predictions, labels = self.metric_necessities(dataloader)
         return accuracy(labels, predictions)

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/concerns/ReportsLogits.py

@@ -11,3 +11,7 @@ class ReportsLogits:
         for model in self.models:
             logits += model(x)
         return logits
+
+    @torch.no_grad()
+    def logits_(self, *args, **kwargs):
+        return self.logits(*args, **kwargs)

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/concerns/ReportsSize.py

@@ -4,8 +4,8 @@ class ReportsSize:
 
     @property
     def size(self):
-        return len(self.state_dicts)
+        return len(self.models)
 
     @property
     def length(self):
-        return len(self.state_dicts)
+        return len(self.models)

homa-0.1.94/src/homa/ensemble/concerns/StoresModels.py

@@ -0,0 +1,29 @@
+import torch
+from copy import deepcopy
+from typing import List
+from ...vision import Model
+
+
+class StoresModels:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.models: List[torch.nn.Module] = []
+
+    def record(self, model: Model | torch.nn.Module):
+        model_: torch.nn.Module | None = None
+        if isinstance(model, Model):
+            model_ = deepcopy(model.network)
+        elif isinstance(model, torch.nn.Module):
+            model_ = deepcopy(model)
+        else:
+            raise TypeError("Wrong input to ensemble record")
+        self.models.append(model_)
+
+    def push(self, *args, **kwargs):
+        self.record(*args, **kwargs)
+
+    def append(self, *args, **kwargs):
+        self.record(*args, **kwargs)
+
+    def add(self, *args, **kwargs):
+        self.record(*args, **kwargs)

{homa-0.0.19 → homa-0.1.94}/src/homa/ensemble/concerns/__init__.py

@@ -1,10 +1,9 @@
 from .CalculatesMetricNecessities import CalculatesMetricNecessities
-from .HasNetwork import HasNetwork
 from .PredictsProbabilities import PredictsProbabilities
-from .RecordsStateDictionaries import RecordsStateDictionaries
 from .ReportsClassificationMetrics import ReportsClassificationMetrics
 from .ReportsEnsembleAccuracy import ReportsEnsembleAccuracy
 from .ReportsEnsembleF1 import ReportsEnsembleF1
 from .ReportsEnsembleKappa import ReportsEnsembleKappa
 from .ReportsLogits import ReportsLogits
 from .ReportsSize import ReportsSize
+from .StoresModels import StoresModels

homa-0.1.94/src/homa/loss/LogitNormLoss.py

@@ -0,0 +1,12 @@
+import torch
+from .Loss import Loss
+
+
+class LogitNormLoss(Loss):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, logits, target):
+        norms = torch.norm(logits, p=2, dim=-1, keepdim=True) + 1e-7
+        normalized_logits = torch.div(logits, norms)
+        return torch.nn.functional.cross_entropy(normalized_logits, target)

homa-0.1.94/src/homa/loss/Loss.py

@@ -0,0 +1,2 @@
+class Loss:
+    pass

homa-0.1.94/src/homa/loss/__init__.py

@@ -0,0 +1,2 @@
+from .LogitNormLoss import LogitNormLoss
+from .Loss import Loss

homa-0.1.94/src/homa/torch/__init__.py

@@ -0,0 +1 @@
+from .helpers import *

homa-0.1.94/src/homa/vision/ClassificationModel.py

@@ -0,0 +1,5 @@
+from .Model import Model
+
+
+class ClassificationModel(Model):
+    pass

homa-0.1.94/src/homa/vision/Resnet.py

@@ -0,0 +1,13 @@
+import torch
+from .modules import ResnetModule
+from .ClassificationModel import ClassificationModel
+from .concerns import Trainable, ReportsMetrics
+from ..device import get_device
+
+
+class Resnet(ClassificationModel, Trainable, ReportsMetrics):
+    def __init__(self, num_classes: int, lr: float = 0.001):
+        super().__init__()
+        self.network = ResnetModule(num_classes).to(get_device())
+        self.criterion = torch.nn.CrossEntropyLoss()
+        self.optimizer = torch.optim.SGD(self.network.parameters(), lr=lr, momentum=0.9)

homa-0.1.94/src/homa/vision/StochasticResnet.py

@@ -0,0 +1,9 @@
+from .Resnet import Resnet
+from .utils import replace_relu
+from ..activations import StochasticActivation
+
+
+class StochasticResnet(Resnet):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        replace_relu(self.network, StochasticActivation)

homa-0.1.94/src/homa/vision/StochasticSwin.py

@@ -0,0 +1,9 @@
+from .Swin import Swin
+from .utils import replace_gelu
+from ..activations import StochasticActivation
+
+
+class StochasticSwin(Swin):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        replace_gelu(self.network.encoder, StochasticActivation)

homa-0.1.94/src/homa/vision/Swin.py

@@ -0,0 +1,12 @@
+import torch
+from .ClassificationModel import ClassificationModel
+from .concerns import Trainable, ReportsMetrics
+from .modules import SwinModule
+
+
+class Swin(ClassificationModel, Trainable, ReportsMetrics):
+    def __init__(self, num_classes: int, lr: float = 0.0001):
+        super().__init__()
+        self.network = SwinModule(num_classes=num_classes)
+        self.optimizer = torch.optim.AdamW(self.network.parameters(), lr=lr)
+        self.criterion = torch.nn.CrossEntropyLoss()

{homa-0.0.19 → homa-0.1.94}/src/homa/vision/__init__.py

@@ -1,3 +1,4 @@
 from .Model import Model
 from .Resnet import Resnet
 from .StochasticResnet import StochasticResnet
+from .Swin import Swin

homa-0.1.94/src/homa/vision/concerns/HasLabels.py

@@ -0,0 +1,13 @@
+import torch
+
+
+class HasLabels:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def predict(self, x: torch.Tensor):
+        return torch.argmax(self.logits(x), dim=1)
+
+    @torch.no_grad()
+    def predict_(self, x: torch.Tensor):
+        return torch.argmax(self.logits(x), dim=1)

homa-0.1.94/src/homa/vision/concerns/HasLogits.py

@@ -0,0 +1,12 @@
+import torch
+
+
+class HasLogits:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def logits(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+    def logits_(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)