torchEBM 0.0.0__tar.gz

torchebm-0.0.0/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: torchEBM
+Version: 0.0.0
+Summary: Components and algorithms for energy-based models
+Author-email: Soran Ghaderi <soran.gdr.cs@gmail.com>
+Project-URL: homepage, https://github.com/soran-ghaderi/torchebm
+Description-Content-Type: text/markdown

torchebm-0.0.0/pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torchEBM"  # Replace with your package name
+#version = "0.1.1"
+dynamic = ["version"]
+description = "Components and algorithms for energy-based models"
+readme = "README.md"
+license = { file = "LICENSE" }
+authors = [{ name="Soran Ghaderi", email="soran.gdr.cs@gmail.com" }]
+dependencies = [
+    # List your dependencies here, e.g., "numpy", "torch"
+]
+[project.urls]
+homepage = "https://github.com/soran-ghaderi/torchebm"
+
+

torchebm-0.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torchebm-0.0.0/torchEBM.egg-info/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: torchEBM
+Version: 0.0.0
+Summary: Components and algorithms for energy-based models
+Author-email: Soran Ghaderi <soran.gdr.cs@gmail.com>
+Project-URL: homepage, https://github.com/soran-ghaderi/torchebm
+Description-Content-Type: text/markdown

torchebm-0.0.0/torchEBM.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+pyproject.toml
+torchEBM.egg-info/PKG-INFO
+torchEBM.egg-info/SOURCES.txt
+torchEBM.egg-info/dependency_links.txt
+torchEBM.egg-info/top_level.txt
+torchebm/__init__.py
+torchebm/core/__init__.py
+torchebm/core/energy_function.py
+torchebm/core/optimizer.py
+torchebm/core/sampler.py
+torchebm/core/trainer.py
+torchebm/cuda/__init__.py
+torchebm/examples/__init__.py
+torchebm/models/__init__.py
+torchebm/models/base_model.py
+torchebm/samplers/__init__.py
+torchebm/utils/__init__.py

torchebm-0.0.0/torchEBM.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

torchebm-0.0.0/torchEBM.egg-info/top_level.txt

@@ -0,0 +1 @@
+torchebm

torchebm-0.0.0/torchebm/core/energy_function.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+import torch
+
+class EnergyFunction(ABC):
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def gradient(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    def cuda_forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward(x)  # Default to PyTorch implementation
+
+    def cuda_gradient(self, x: torch.Tensor) -> torch.Tensor:
+        return self.gradient(x)

torchebm-0.0.0/torchebm/core/sampler.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+import torch
+
+from torchebm.core.energy_function import EnergyFunction
+
+
+class Sampler(ABC):
+    @abstractmethod
+    def sample(self, energy_function: EnergyFunction, initial_state: torch.Tensor, num_steps: int) -> torch.Tensor:
+        pass
+
+    def cuda_sample(self, energy_function: EnergyFunction, initial_state: torch.Tensor, num_steps: int) -> torch.Tensor:
+        return self.sample(energy_function, initial_state, num_steps)  # Default to PyTorch implementation

torchebm-0.0.0/torchebm/core/trainer.py

@@ -0,0 +1,29 @@
+import torch
+from .energy_function import EnergyFunction
+from .sampler import Sampler
+
+class ContrastiveDivergenceTrainer:
+    def __init__(self, energy_function: EnergyFunction, sampler: Sampler, learning_rate: float = 0.01):
+        self.energy_function = energy_function
+        self.sampler = sampler
+        self.optimizer = torch.optim.Adam(self.energy_function.parameters(), lr=learning_rate)
+
+    def train_step(self, real_data: torch.Tensor) -> dict:
+        self.optimizer.zero_grad()
+
+        # Positive phase
+        positive_energy = self.energy_function(real_data)
+
+        # Negative phase
+        initial_samples = torch.randn_like(real_data)
+        negative_samples = self.sampler.sample(self.energy_function, initial_samples, num_steps=10)
+        negative_energy = self.energy_function(negative_samples)
+
+        # Compute loss
+        loss = positive_energy.mean() - negative_energy.mean()
+
+        # Backpropagation
+        loss.backward()
+        self.optimizer.step()
+
+        return {"loss": loss.item(), "positive_energy": positive_energy.mean().item(), "negative_energy": negative_energy.mean().item()}

torchebm-0.0.0/torchebm/models/base_model.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+import torch
+from ..core import EnergyFunction, Sampler
+
+class BaseModel(ABC):
+    def __init__(self, energy_function: EnergyFunction, sampler: Sampler):
+        self.energy_function = energy_function
+        self.sampler = sampler
+
+    @abstractmethod
+    def energy(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def sample(self, num_samples: int) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def train_step(self, real_data: torch.Tensor) -> dict:
+        pass
+