Allow str device

diffusion/__init__.py

@@ -24,12 +24,12 @@ class Model(Generic[D]):
     data: Data
     schedule: Schedule
     noise: Noise[D]
-    loss: Loss[D]
     net: Net
+    loss: Loss[D]
     time: Time = field(default_factory=Discrete)
     guidance: Optional[Guidance] = None  # TODO remove hardcoding
     optimizer: Optional[Optimizer | Callable[..., Optimizer]] = None
-    device: torch.device = torch.device("cpu")
+    device: str | torch.device = torch.device("cpu")
     compile: bool = True
 
     @torch.no_grad()
@@ -43,28 +43,11 @@ def __post_init__(self):
         self.net = self.net.to(self.device)
         for name, value in vars(self.data).items():
             if isinstance(value, nn.Module):
-                setattr(self.data, name, value.to(self.device)) 
+                setattr(self.data, name, value.to(self.device))
         if self.compile and sys.version_info < (3, 11):
             self.net = torch.compile(self.net)  # type: ignore[union-attr]
-
-    @torch.no_grad()
-    def load(self, path: Path | str):
-        state = torch.load(path)
-        self.net.load_state_dict(state["net"])
-        for name, dict in state["data"].items():
-            getattr(self.data, name).load_state_dict(dict)
-
-    @torch.no_grad()
-    def save(self, path: Path | str):
-        state = {
-            "net": self.net.state_dict(),
-            "data": {
-                name: value.state_dict()
-                for name, value in vars(self.data).items()
-                if isinstance(value, nn.Module)
-            }
-        }
-        torch.save(state, path)
+        if isinstance(self.device, str):
+            self.device = torch.device(self.device)
 
     @torch.enable_grad()
     def train(self, epochs: int = 1, progress: bool = True) -> Iterator[float]:

examples/conditional-diffusion.py

@@ -34,7 +34,7 @@
     net=UNet(channels=(1, 64, 128, 256), labels=10),
     guidance=ClassifierFree(dropout=0.1, strength=2),
     loss=Simple(parameter="epsilon"),
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 if (output / "model.pt").exists():

examples/embedding-diffusion.py

@@ -43,7 +43,7 @@
     noise=Gaussian(parameter="x", variance="fixed"),
     loss=Simple(parameter="x"),
     net=Transformer(input=32, width=1024, depth=16, heads=16),
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 if (output / "model.pt").exists():

examples/text-diffusion.py

@@ -45,7 +45,7 @@
     noise=Absorbing(len(v)),
     loss=VLB() + 1e-2 * Lambda[Cat](lambda batch: Cat(batch.hat[0]).nll(batch.x).sum()),
     net=Transformer(input=len(v), width=1024, depth=16, heads=16) | Softmax(3),
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 if (output / "model.pt").exists():

examples/transformer-diffusion.py

@@ -51,7 +51,7 @@
     noise=Gaussian(parameter="epsilon", variance="fixed"),
     net=Transformer(input=x.shape[2], width=768, depth=12, heads=12),
     loss=Simple(parameter="epsilon"),
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 if (output / "model.pt").exists():

examples/unconditional-diffusion.py

@@ -32,7 +32,7 @@
     noise=Gaussian(parameter="epsilon", variance="fixed"),
     net=UNet(channels=(1, 64, 128, 256)),
     loss=Simple(parameter="epsilon"),
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 if (output / "model.pt").exists():