derived value for image preprocessing belongs to the data config class

README.md

@@ -1078,6 +1078,7 @@ This library would not have gotten to this working state without the help of
 - [x] use an experimental tracker agnostic setup, as done <a href="https://github.com/lucidrains/tf-bind-transformer#simple-trainer-class-for-fine-tuning">here</a>
 - [x] use pydantic for config drive training
 - [x] for both diffusion prior and decoder, all exponential moving averaged models needs to be saved and restored as well (as well as the step number)
+- [x] offer save / load methods on the trainer classes to automatically take care of state dicts for scalers / optimizers / saving versions and checking for breaking changes
 - [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet (test out unet² in ddpm repo) - consider https://github.com/lucidrains/uformer-pytorch attention-based unet
 - [ ] transcribe code to Jax, which lowers the activation energy for distributed training, given access to TPUs
 - [ ] train on a toy task, offer in colab
@@ -1087,11 +1088,9 @@ This library would not have gotten to this working state without the help of
 - [ ] test out grid attention in cascading ddpm locally, decide whether to keep or remove
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
 - [ ] make sure FILIP works with DALL-E2 from x-clip https://arxiv.org/abs/2111.07783
-- [ ] offer save / load methods on the trainer classes to automatically take care of state dicts for scalers / optimizers / saving versions and checking for breaking changes
 - [ ] bring in skip-layer excitatons (from lightweight gan paper) to see if it helps for either decoder of unet or vqgan-vae training
 - [ ] decoder needs one day worth of refactor for tech debt
 - [ ] allow for unet to be able to condition non-cross attention style as well
-- [ ] for all model classes with hyperparameters that changes the network architecture, make it requirement that they must expose a config property, and write a simple function that asserts that it restores the object correctly
 - [ ] read the paper, figure it out, and build it https://github.com/lucidrains/DALLE2-pytorch/issues/89
 
 ## Citations

dalle2_pytorch/train_configs.py

@@ -64,6 +64,22 @@ class DecoderDataConfig(BaseModel):
     resample_train: bool = False
     preprocessing: Dict[str, Any] = {'ToTensor': True}
 
+    @property
+    def img_preproc(self):
+        def _get_transformation(transformation_name, **kwargs):
+            if transformation_name == "RandomResizedCrop":
+                return T.RandomResizedCrop(**kwargs)
+            elif transformation_name == "RandomHorizontalFlip":
+                return T.RandomHorizontalFlip()
+            elif transformation_name == "ToTensor":
+                return T.ToTensor()
+
+        transforms = []
+        for transform_name, transform_kwargs_or_bool in self.preprocessing.items():
+            transform_kwargs = {} if not isinstance(transform_kwargs_or_bool, dict) else transform_kwargs_or_bool
+            transforms.append(_get_transformation(transform_name, **transform_kwargs))
+        return T.Compose(transforms)
+
 class DecoderTrainConfig(BaseModel):
     epochs: int = 20
     lr: float = 1e-4
@@ -117,19 +133,3 @@ class TrainDecoderConfig(BaseModel):
         with open(json_path) as f:
             config = json.load(f)
         return cls(**config)
-
-    @property
-    def img_preproc(self):
-        def _get_transformation(transformation_name, **kwargs):
-            if transformation_name == "RandomResizedCrop":
-                return T.RandomResizedCrop(**kwargs)
-            elif transformation_name == "RandomHorizontalFlip":
-                return T.RandomHorizontalFlip()
-            elif transformation_name == "ToTensor":
-                return T.ToTensor()
-        
-        transforms = []
-        for transform_name, transform_kwargs_or_bool in self.data.preprocessing.items():
-            transform_kwargs = {} if not isinstance(transform_kwargs_or_bool, dict) else transform_kwargs_or_bool
-            transforms.append(_get_transformation(transform_name, **transform_kwargs))
-        return T.Compose(transforms)

dalle2_pytorch/trainer.py

@@ -133,12 +133,6 @@ def split_args_and_kwargs(*args, split_size = None, **kwargs):
         chunk_size_frac = chunk_size / batch_size
         yield chunk_size_frac, (chunked_args, chunked_kwargs)
 
-# print helpers
-
-def print_ribbon(s, symbol = '=', repeat = 40):
-    flank = symbol * repeat
-    return f'{flank} {s} {flank}'
-
 # saving and loading functions
 
 # for diffusion prior

dalle2_pytorch/utils.py

@@ -1,5 +1,7 @@
 import time
 
+# time helpers
+
 class Timer:
     def __init__(self):
         self.reset()
@@ -9,3 +11,9 @@ class Timer:
 
     def elapsed(self):
         return time.time() - self.last_time
+
+# print helpers
+
+def print_ribbon(s, symbol = '=', repeat = 40):
+    flank = symbol * repeat
+    return f'{flank} {s} {flank}'

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.4.5',
+  version = '0.4.7',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',

train_decoder.py

@@ -1,9 +1,9 @@
 from dalle2_pytorch import Unet, Decoder
-from dalle2_pytorch.trainer import DecoderTrainer, print_ribbon
+from dalle2_pytorch.trainer import DecoderTrainer
 from dalle2_pytorch.dataloaders import create_image_embedding_dataloader
 from dalle2_pytorch.trackers import WandbTracker, ConsoleTracker
 from dalle2_pytorch.train_configs import TrainDecoderConfig
-from dalle2_pytorch.utils import Timer
+from dalle2_pytorch.utils import Timer, print_ribbon
 
 import torchvision
 import torch
@@ -420,7 +420,7 @@ def initialize_training(config):
 
     dataloaders = create_dataloaders (
         available_shards=all_shards,
-        img_preproc = config.img_preproc,
+        img_preproc = config.data.img_preproc,
         train_prop = config.data.splits.train,
         val_prop = config.data.splits.val,
         test_prop = config.data.splits.test,

train_diffusion_prior.py

@@ -9,10 +9,10 @@ from torch import nn
 
 from dalle2_pytorch.dataloaders import make_splits
 from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork, OpenAIClipAdapter
-from dalle2_pytorch.trainer import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model, print_ribbon
+from dalle2_pytorch.trainer import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model
 
 from dalle2_pytorch.trackers import ConsoleTracker, WandbTracker
-from dalle2_pytorch.utils import Timer
+from dalle2_pytorch.utils import Timer, print_ribbon
 
 from embedding_reader import EmbeddingReader