torchEBM 0.1.0__py3-none-any.whl

torchEBM-0.1.0.dist-info/METADATA

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

torchEBM-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+torchebm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/core/energy_function.py,sha256=sozA5nIV6Li3mzbBSsWgzqSu-Pf1aqijhPhDdn8YqGE,481
+torchebm/core/optimizer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/core/sampler.py,sha256=K_qpvbX5nInIhoV5oAfjE-xVXZf48hsz-BYrYHd2X3g,448
+torchebm/core/trainer.py,sha256=xNnGrD2njbT3vv4U-2cX79hGXiCSIRdetdHn3geaYdc,1129
+torchebm/cuda/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/models/base_model.py,sha256=pQ5uvnIjsyywaZe675spSBrsot8lnX46UtCdX1EIgOM,541
+torchebm/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchebm/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchEBM-0.1.0.dist-info/METADATA,sha256=ADUxPZ_uauBOP3X7QDc93NTwHGqWPdJtYFsHnvnQE_U,205
+torchEBM-0.1.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+torchEBM-0.1.0.dist-info/top_level.txt,sha256=FTDHo7yycQ1rtcscAqb-fXXaoOVdtugPVNEUnm_FYHs,9
+torchEBM-0.1.0.dist-info/RECORD,,

torchEBM-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torchEBM-0.1.0.dist-info/top_level.txt

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

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/core/sampler.py

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+import torch
+
+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/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/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
+