let the pydantic config base model take care of loading configuration from json path

configs/train_decoder_config.example.json

@@ -0,0 +1,99 @@
+{
+    "unets": [
+        {
+            "dim": 128,
+            "image_embed_dim": 768,
+            "cond_dim": 64,
+            "channels": 3,
+            "dim_mults": [1, 2, 4, 8],
+            "attn_dim_head": 32,
+            "attn_heads": 16
+        }
+    ],
+    "decoder": {
+        "image_sizes": [64],
+        "channels": 3,
+        "timesteps": 1000,
+        "loss_type": "l2",
+        "beta_schedule": "cosine",
+        "learned_variance": true
+    },
+    "data": {
+        "webdataset_base_url": "pipe:s3cmd get s3://bucket/path/{}.tar -",
+        "embeddings_url": "s3://bucket/embeddings/path/",
+        "num_workers": 4,
+        "batch_size": 64,
+        "start_shard": 0,
+        "end_shard": 9999999,
+        "shard_width": 6,
+        "index_width": 4,
+        "splits": {
+            "train": 0.75,
+            "val": 0.15,
+            "test": 0.1
+        },
+        "shuffle_train": true,
+        "resample_train": false,
+        "preprocessing": {
+            "RandomResizedCrop": {
+                "size": [128, 128],
+                "scale": [0.75, 1.0],
+                "ratio": [1.0, 1.0]
+            },
+            "ToTensor": true
+        }
+    },
+    "train": {
+        "epochs": 20,
+        "lr": 1e-4,
+        "wd": 0.01,
+        "max_grad_norm": 0.5,
+        "save_every_n_samples": 100000,
+        "n_sample_images": 6,
+        "device": "cuda:0",
+        "epoch_samples": null,
+        "validation_samples": null,
+        "use_ema": true,
+        "ema_beta": 0.99,
+        "amp": false,
+        "save_all": false,
+        "save_latest": true,
+        "save_best": true,
+        "unet_training_mask": [true]
+    },
+    "evaluate": {
+        "n_evaluation_samples": 1000,
+        "FID": {
+            "feature": 64
+        },
+        "IS": {
+            "feature": 64,
+            "splits": 10
+        },
+        "KID": {
+            "feature": 64,
+            "subset_size": 10
+        },
+        "LPIPS": {
+            "net_type": "vgg",
+            "reduction": "mean"
+        }
+    },
+    "tracker": {
+        "tracker_type": "console",
+        "data_path": "./models",
+
+        "wandb_entity": "",
+        "wandb_project": "",
+
+        "verbose": false
+    },
+    "load": {
+        "source": null,
+
+        "run_path": "",
+        "file_path": "",
+
+        "resume": false
+    }
+}

dalle2_pytorch/train_configs.py

@@ -1,5 +1,6 @@
+import json
 from torchvision import transforms as T
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, validator, root_validator
 from typing import List, Iterable, Optional, Union, Tuple, Dict, Any
 
 def exists(val):
@@ -38,6 +39,17 @@ class DecoderConfig(BaseModel):
     class Config:
         extra = "allow"
 
+class TrainSplitConfig(BaseModel):
+    train: float = 0.75
+    val: float = 0.15
+    test: float = 0.1
+
+    @root_validator
+    def validate_all(cls, fields):
+        if sum([*fields.values()]) != 1.:
+            raise ValueError(f'{fields.keys()} must sum to 1.0')
+        return fields
+
 class DecoderDataConfig(BaseModel):
     webdataset_base_url: str     # path to a webdataset with jpg images
     embeddings_url: str          # path to .npy files with embeddings
@@ -47,11 +59,7 @@ class DecoderDataConfig(BaseModel):
     end_shard: int = 9999999
     shard_width: int = 6
     index_width: int = 4
-    splits: Dict[str, float] = {
+    splits: TrainSplitConfig
-        'train': 0.75,
-        'val': 0.15,
-        'test': 0.1
-    }
     shuffle_train: bool = True
     resample_train: bool = False
     preprocessing: Dict[str, Any] = {'ToTensor': True}
@@ -104,6 +112,12 @@ class TrainDecoderConfig(BaseModel):
     tracker: TrackerConfig
     load: DecoderLoadConfig
 
+    @classmethod
+    def from_json_path(cls, json_path):
+        with open(json_path) as f:
+            config = json.load(f)
+        return cls(**config)
+
     @property
     def img_preproc(self):
         def _get_transformation(transformation_name, **kwargs):

setup.py

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

train_decoder.py

@@ -5,7 +5,6 @@ from dalle2_pytorch.trackers import WandbTracker, ConsoleTracker
 from dalle2_pytorch.train_configs import TrainDecoderConfig
 from dalle2_pytorch.utils import Timer
 
-import json
 import torchvision
 import torch
 from torchmetrics.image.fid import FrechetInceptionDistance
@@ -422,9 +421,9 @@ def initialize_training(config):
     dataloaders = create_dataloaders (
         available_shards=all_shards,
         img_preproc = config.img_preproc,
-        train_prop = config.data["splits"]["train"],
+        train_prop = config.data.splits.train,
-        val_prop = config.data["splits"]["val"],
+        val_prop = config.data.splits.val,
-        test_prop = config.data["splits"]["test"],
+        test_prop = config.data.splits.test,
         n_sample_images=config.train.n_sample_images,
         **config.data.dict()
     )
@@ -449,9 +448,7 @@ def initialize_training(config):
 @click.option("--config_file", default="./train_decoder_config.json", help="Path to config file")
 def main(config_file):
     print("Recalling config from {}".format(config_file))
-    with open(config_file) as f:
+    config = TrainDecoderConfig.from_json_path(config_file)
-        config = json.load(f)
-    config = TrainDecoderConfig(**config)
     initialize_training(config)