flaxdiff 0.1.35.6__py3-none-any.whl → 0.1.36.1__py3-none-any.whl

flaxdiff/trainer/simple_trainer.py

@@ -1,442 +0,0 @@
-import orbax.checkpoint
-import tqdm
-from flax import linen as nn
-import jax
-from typing import Callable
-from dataclasses import field
-import jax.numpy as jnp
-import numpy as np
-from functools import partial
-from clu import metrics
-from flax.training import train_state  # Useful dataclass to keep train state
-import optax
-from flax import struct                # Flax dataclasses
-import flax
-import time
-import os
-import orbax
-from flax.training import orbax_utils
-from jax.sharding import Mesh, PartitionSpec as P
-from jax.experimental import mesh_utils
-from jax.experimental.shard_map import shard_map
-from orbax.checkpoint.utils import fully_replicated_host_local_array_to_global_array
-from termcolor import colored
-from typing import Dict, Callable, Sequence, Any, Union, Tuple
-from flax.training.dynamic_scale import DynamicScale
-from flaxdiff.utils import RandomMarkovState
-
-PROCESS_COLOR_MAP = {
-    0: "green",
-    1: "yellow",
-    2: "magenta",
-    3: "cyan", 
-    4: "white",
-    5: "light_blue",
-    6: "light_red",
-    7: "light_cyan"
-}
-
-def _build_global_shape_and_sharding(
-    local_shape: tuple[int, ...], global_mesh: Mesh
-) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]:
-  sharding = jax.sharding.NamedSharding(global_mesh, P(global_mesh.axis_names))
-  global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:]
-  return global_shape, sharding
-
-
-def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array:
-  """Put local sharded array into local devices"""
-  global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh)
-  try:
-    local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0)
-  except ValueError as array_split_error:
-    raise ValueError(
-        f"Unable to put to devices shape {array.shape} with "
-        f"local device count {len(global_mesh.local_devices)} "
-    ) from array_split_error
-  local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices)
-  return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers)
-
-def convert_to_global_tree(global_mesh, pytree):
-    return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree)
-
-@struct.dataclass
-class Metrics(metrics.Collection):
-    accuracy: metrics.Accuracy
-    loss: metrics.Average.from_output('loss')
-
-# Define the TrainState
-class SimpleTrainState(train_state.TrainState):
-    metrics: Metrics
-    dynamic_scale: DynamicScale
-
-class SimpleTrainer:
-    state: SimpleTrainState
-    best_state: SimpleTrainState
-    best_loss: float
-    model: nn.Module
-    ema_decay: float = 0.999
-
-    def __init__(self,
-                 model: nn.Module,
-                 input_shapes: Dict[str, Tuple[int]],
-                 optimizer: optax.GradientTransformation,
-                 rngs: jax.random.PRNGKey,
-                 train_state: SimpleTrainState = None,
-                 name: str = "Simple",
-                 load_from_checkpoint: str = None,
-                 checkpoint_suffix: str = "",
-                 loss_fn=optax.l2_loss,
-                 param_transforms: Callable = None,
-                 wandb_config: Dict[str, Any] = None,
-                 distributed_training: bool = None,
-                 checkpoint_base_path: str = "./checkpoints",
-                 checkpoint_step: int = None,
-                 use_dynamic_scale: bool = False,
-                 ):
-        if distributed_training is None or distributed_training is True:
-            # Auto-detect if we are running on multiple devices
-            distributed_training = jax.device_count() > 1
-            self.mesh = jax.sharding.Mesh(jax.devices(), 'data')
-        else:
-            self.mesh = None
-
-        self.distributed_training = distributed_training
-        self.model = model
-        self.name = name
-        self.loss_fn = loss_fn
-        self.input_shapes = input_shapes
-        self.checkpoint_base_path = checkpoint_base_path
-        
-        
-        if wandb_config is not None and jax.process_index() == 0:
-            run = wandb.init(**wandb_config)
-            self.wandb = run
-            
-            # define our custom x axis metric
-            self.wandb.define_metric("train/step")
-            self.wandb.define_metric("train/epoch")
-            
-            self.wandb.define_metric("train/loss", step_metric="train/step")
-            
-            self.wandb.define_metric("train/epoch_time", step_metric="train/epoch")
-            self.wandb.define_metric("train/avg_time_per_step", step_metric="train/epoch")
-            self.wandb.define_metric("train/avg_loss", step_metric="train/epoch")
-            self.wandb.define_metric("train/best_loss", step_metric="train/epoch")
-            
-        # checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-        async_checkpointer = orbax.checkpoint.AsyncCheckpointer(orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=60)
-
-        options = orbax.checkpoint.CheckpointManagerOptions(
-            max_to_keep=4, create=True)
-        self.checkpointer = orbax.checkpoint.CheckpointManager(
-            self.checkpoint_path() + checkpoint_suffix, async_checkpointer, options)
-
-        if load_from_checkpoint is not None:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step)
-        else:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = 0, 0, None, None, None
-
-        self.latest_step = latest_step
-        
-        if rngstate:
-            self.rngstate = RandomMarkovState(**rngstate)
-        else:
-            self.rngstate = RandomMarkovState(rngs)
-            
-        self.rngstate, subkey = self.rngstate.get_random_key()
-
-        if train_state == None:
-            state, best_state = self.generate_states(
-                optimizer, subkey, old_state, old_best_state, model, param_transforms, use_dynamic_scale
-            )
-            self.init_state(state, best_state)
-        else:
-            self.state = train_state
-            self.best_state = train_state
-            self.best_loss = 1e9
-
-    def get_input_ones(self):
-        return {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()}
-
-    def generate_states(
-        self,
-        optimizer: optax.GradientTransformation,
-        rngs: jax.random.PRNGKey,
-        existing_state: dict = None,
-        existing_best_state: dict = None,
-        model: nn.Module = None,
-        param_transforms: Callable = None,
-        use_dynamic_scale: bool = False
-    ) -> Tuple[SimpleTrainState, SimpleTrainState]:
-        print("Generating states for SimpleTrainer")
-        rngs, subkey = jax.random.split(rngs)
-
-        if existing_state == None:
-            input_vars = self.get_input_ones()
-            params = model.init(subkey, **input_vars)
-        else:
-            params = existing_state['params']
-            
-        if param_transforms is not None:
-            params = param_transforms(params)
-
-        state = SimpleTrainState.create(
-            apply_fn=model.apply,
-            params=params,
-            tx=optimizer,
-            metrics=Metrics.empty(),
-            dynamic_scale = DynamicScale() if use_dynamic_scale else None
-        )
-        if existing_best_state is not None:
-            best_state = state.replace(
-                params=existing_best_state['params'])
-        else:
-            best_state = state
-            
-        return state, best_state
-
-    def init_state(
-        self,
-        state: SimpleTrainState,
-        best_state: SimpleTrainState,
-    ):
-        self.best_loss = 1e9
-
-        self.state = state
-        self.best_state = best_state
-
-    def get_state(self):
-        # return fully_replicated_host_local_array_to_global_array()
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.state)
-
-    def get_best_state(self):
-        # return convert_to_global_tree(self.mesh, flax.jax_utils.replicate(self.best_state, jax.local_devices()))
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.best_state)
-        
-    def get_rngstate(self):
-        # return convert_to_global_tree(self.mesh, flax.jax_utils.replicate(self.rngstate, jax.local_devices()))
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.rngstate)
-
-    def checkpoint_path(self):
-        path = os.path.join(self.checkpoint_base_path, self.name.replace(' ', '_').lower())
-        if not os.path.exists(path):
-            os.makedirs(path)
-        return path
-
-    def tensorboard_path(self):
-        experiment_name = self.name
-        path = os.path.join(os.path.abspath('./tensorboard'), experiment_name)
-        if not os.path.exists(path):
-            os.makedirs(path)
-        return path
-
-    def load(self, checkpoint_path=None, checkpoint_step=None):
-        if checkpoint_path is None:
-            checkpointer = self.checkpointer
-        else:
-            checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-            options = orbax.checkpoint.CheckpointManagerOptions(
-                max_to_keep=4, create=False)
-            checkpointer = orbax.checkpoint.CheckpointManager(
-                checkpoint_path, checkpointer, options)    
-        
-        if checkpoint_step is None:
-            step = checkpointer.latest_step()
-        else:
-            step = checkpoint_step
-        
-        print("Loading model from checkpoint at step ", step)
-        ckpt = checkpointer.restore(step)
-        state = ckpt['state']
-        best_state = ckpt['best_state']
-        rngstate = ckpt['rngs']
-        # Convert the state to a TrainState
-        self.best_loss = ckpt['best_loss']
-        current_epoch = ckpt.get('epoch', step) # Must be a checkpoint from an older version which used epochs instead of steps
-        print(
-            f"Loaded model from checkpoint at epoch {current_epoch} step {step}", ckpt['best_loss'])
-        return current_epoch, step, state, best_state, rngstate
-
-    def save(self, epoch=0, step=0):
-        print(f"Saving model at epoch {epoch} step {step}")
-        ckpt = {
-            # 'model': self.model,
-            'rngs': self.get_rngstate(),
-            'state': self.get_state(),
-            'best_state': self.get_best_state(),
-            'best_loss': np.array(self.best_loss),
-            'epoch': epoch,
-        }
-        try:
-            save_args = orbax_utils.save_args_from_target(ckpt)
-            self.checkpointer.save(step, ckpt, save_kwargs={
-                                   'save_args': save_args}, force=True)
-            self.checkpointer.wait_until_finished()
-            pass
-        except Exception as e:
-            print("Error saving checkpoint", e)
-
-    def _define_train_step(self, **kwargs):
-        model = self.model
-        loss_fn = self.loss_fn
-        distributed_training = self.distributed_training
-
-        def train_step(train_state: SimpleTrainState, rng_state: RandomMarkovState, batch, local_device_indexes):
-            """Train for a single step."""
-            images = batch['image']
-            labels = batch['label']
-
-            def model_loss(params):
-                preds = model.apply(params, images)
-                expected_output = labels
-                nloss = loss_fn(preds, expected_output)
-                loss = jnp.mean(nloss)
-                return loss
-            loss, grads = jax.value_and_grad(model_loss)(train_state.params)
-            if distributed_training:
-                grads = jax.lax.pmean(grads, "data")
-            train_state = train_state.apply_gradients(grads=grads)
-            return train_state, loss, rng_state
-        
-        if distributed_training:
-            train_step = shard_map(train_step, mesh=self.mesh, in_specs=(P(), P(), P('data'), P('data')), out_specs=(P(), P('data'), P()))
-            train_step = jax.pmap(train_step)
-        return train_step
-
-    def _define_compute_metrics(self):
-        model = self.model
-        loss_fn = self.loss_fn
-
-        @jax.jit
-        def compute_metrics(state: SimpleTrainState, batch):
-            preds = model.apply(state.params, batch['image'])
-            expected_output = batch['label']
-            loss = jnp.mean(loss_fn(preds, expected_output))
-            metric_updates = state.metrics.single_from_model_output(
-                loss=loss, logits=preds, labels=expected_output)
-            metrics = state.metrics.merge(metric_updates)
-            state = state.replace(metrics=metrics)
-            return state
-        return compute_metrics
-
-    def summary(self):
-        input_vars = self.get_input_ones()
-        print(self.model.tabulate(jax.random.key(0), **input_vars,
-              console_kwargs={"width": 200, "force_jupyter": True, }))
-
-    def config(self):
-        return {
-            "model": self.model,
-            "state": self.state,
-            "name": self.name,
-            "input_shapes": self.input_shapes
-        }
-
-    def init_tensorboard(self, batch_size, steps_per_epoch, epochs):
-        from flax.metrics import tensorboard
-        summary_writer = tensorboard.SummaryWriter(self.tensorboard_path())
-        summary_writer.hparams({
-            **self.config(),
-            "steps_per_epoch": steps_per_epoch,
-            "epochs": epochs,
-            "batch_size": batch_size
-        })
-        return summary_writer
-
-    def fit(self, data, steps_per_epoch, epochs, train_step_args={}):
-        train_ds = iter(data['train']())
-        if 'test' in data:
-            test_ds = data['test']
-        else:
-            test_ds = None
-        train_step = self._define_train_step(**train_step_args)
-        compute_metrics = self._define_compute_metrics()
-        train_state = self.state
-        rng_state = self.rngstate
-        global_device_count = jax.device_count()
-        local_device_count = jax.local_device_count()
-        process_index = jax.process_index()
-        if self.distributed_training:
-            global_device_indexes = jnp.arange(global_device_count)
-        else:
-            global_device_indexes = 0
-
-        def train_loop(current_step, pbar: tqdm.tqdm, train_state, rng_state):
-            epoch_loss = 0
-            current_epoch = current_step // steps_per_epoch
-            last_save_time = time.time()
-            for i in range(steps_per_epoch):
-                batch = next(train_ds)
-                if self.distributed_training and global_device_count > 1:
-                    # Convert the local device batches to a unified global jax.Array 
-                    batch = convert_to_global_tree(self.mesh, batch)
-                train_state, loss, rng_state = train_step(train_state, rng_state, batch, global_device_indexes)
-
-                if self.distributed_training:
-                    loss = jax.experimental.multihost_utils.process_allgather(loss)
-                    loss = jnp.mean(loss) # Just to make sure its a scaler value
-                    
-                if loss <= 1e-6:
-                    # If the loss is too low, we can assume the model has diverged
-                    print(colored(f"Loss too low at step {current_step} => {loss}", 'red'))
-                    # Exit the training loop
-                    exit(1)
-                            
-                epoch_loss += loss
-                current_step += 1
-                if i % 100 == 0:
-                    if pbar is not None:
-                        pbar.set_postfix(loss=f'{loss:.4f}')
-                        pbar.update(100)
-                        if self.wandb is not None:
-                            self.wandb.log({
-                                "train/step" : current_step,
-                                "train/loss": loss,
-                            }, step=current_step)
-                    # Save the model every 40 minutes
-                    if time.time() - last_save_time > 40 * 60:
-                        print(f"Saving model after 40 minutes at step {current_step}")
-                        self.save(current_epoch, current_step)
-                        last_save_time = time.time()
-            print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/steps_per_epoch}", 'green'))
-            return epoch_loss, current_step, train_state, rng_state
-
-        while self.latest_step < epochs * steps_per_epoch:
-            current_epoch = self.latest_step // steps_per_epoch
-            print(f"\nEpoch {current_epoch}/{epochs}")
-            start_time = time.time()
-            epoch_loss = 0
-
-            if process_index == 0:
-                with tqdm.tqdm(total=steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step') as pbar:
-                    epoch_loss, current_step, train_state, rng_state = train_loop(self.latest_step, pbar, train_state, rng_state)
-            else:
-                epoch_loss, current_step, train_state, rng_state = train_loop(self.latest_step, None, train_state, rng_state)
-                print(colored(f"Epoch done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
-            
-            self.latest_step = current_step
-            end_time = time.time()
-            self.state = train_state
-            self.rngstate = rng_state
-            total_time = end_time - start_time
-            avg_time_per_step = total_time / steps_per_epoch
-            avg_loss = epoch_loss / steps_per_epoch
-            if avg_loss < self.best_loss:
-                self.best_loss = avg_loss
-                self.best_state = train_state
-                self.save(current_epoch, current_step)
-                
-            if process_index == 0:
-                if self.wandb is not None:
-                    self.wandb.log({
-                        "train/epoch_time": total_time,
-                        "train/avg_time_per_step": avg_time_per_step,
-                        "train/avg_loss": avg_loss,
-                        "train/best_loss": self.best_loss,
-                        "train/epoch": current_epoch,
-                    }, step=current_step)
-                print(colored(f"\n\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}", 'green'))
-        print("Training done")
-        self.save(epochs)
-        return self.state

flaxdiff-0.1.35.6.dist-info/RECORD

@@ -1,40 +0,0 @@
-flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
-flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
-flaxdiff/data/online_loader.py,sha256=DoHrMZCi5gMd9tmkCpZIUU9lGxvfYtuaz58943_lCRc,11315
-flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
-flaxdiff/models/attention.py,sha256=JvrP7-09MV6IfRLRBhqjPmNUU-lkEMk9TOnJSBKcar8,13289
-flaxdiff/models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
-flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
-flaxdiff/models/simple_unet.py,sha256=L5m2j5580QP7pJ5VIme7U5xYA22PZiGP7qdvcKUnB38,11463
-flaxdiff/models/simple_vit.py,sha256=UCDDr0XVnpf6tbJWKFtEt3_nAqMqOoakXf5amyVWZNo,7929
-flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
-flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
-flaxdiff/models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
-flaxdiff/models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
-flaxdiff/predictors/__init__.py,sha256=SKkYYRF9Wfgk2zhtZw4vCXOdOeRlrm2Mk6cvuaEvAzc,4403
-flaxdiff/samplers/__init__.py,sha256=_S-9TwDeshrI0VmapV-J2hqjTByOa0-oOeUs_IdovjU,285
-flaxdiff/samplers/common.py,sha256=_an5h5Niz9Joz_-ppridLrGHpu8X0VVvhNGknPu6AUY,5272
-flaxdiff/samplers/ddim.py,sha256=XHMBX06S5hMTnKMaGh6fmq189pQcaGkA6fnX6YPbHP0,511
-flaxdiff/samplers/ddpm.py,sha256=d_58hfVShJHsRPQf5h1-4YKrD43-HGjWz9Vd8hltZBg,2627
-flaxdiff/samplers/euler.py,sha256=Epf7LBKUky7B8b-1ZyIlLWdRMgjmP08BQraGSKmr_3I,2726
-flaxdiff/samplers/heun_sampler.py,sha256=hhWnSM26OfOIFAcsuWYa1z-2QPjASuoYTop2byLWqzE,1388
-flaxdiff/samplers/multistep_dpm.py,sha256=ocmEq2sCvsULy6oTFaD5BhTU4c8VHsge4bdg6tfxW80,2724
-flaxdiff/samplers/rk4_sampler.py,sha256=BF-dMV1KauO-SYShqrCfm3U3V-1n4clqQXBeoG8RWQo,1728
-flaxdiff/schedulers/__init__.py,sha256=3id390WEfdf-MN-oLSPAhlRFIXrFWr6ioAHPAwURJyE,375
-flaxdiff/schedulers/common.py,sha256=b-W4iI-aqScpVE8VZbBpiYvAVI6rqDkUP-C_hEVBwCI,4151
-flaxdiff/schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
-flaxdiff/schedulers/cosine.py,sha256=9ban0dFHLMm35wQvaBT4nCQwPGmzNsXwQ1xI0oppmJI,2005
-flaxdiff/schedulers/discrete.py,sha256=O54wH2HVu3olJA71NxgAXFW9cr6B6Gl-DR_uZeytpds,3319
-flaxdiff/schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
-flaxdiff/schedulers/karras.py,sha256=4GN120kGwdxxU-h2mVdhBVy9IORkUMm_vvz3XjthBcI,3355
-flaxdiff/schedulers/linear.py,sha256=6003F5ISq1Wc0h6UAzY95MJgsDIKGMhBzbiVALpea0k,581
-flaxdiff/schedulers/sqrt.py,sha256=1F84ZgQPuoNMhe6yxGTR2G0h7dPOZtm4UDQOakbSsEU,445
-flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,128
-flaxdiff/trainer/autoencoder_trainer.py,sha256=al7AsZ7yeDMEiDD-gbcXf0ADq_xfk1VMxvg24GfA-XQ,7008
-flaxdiff/trainer/diffusion_trainer.py,sha256=wKkg63DWZjx2MoM3VQNCDIr40rWN8fUGxH9jWWxfZao,9373
-flaxdiff/trainer/simple_trainer.py,sha256=cawm6fZNQoLLATMneAU2gQ9j7kefqHnBPHuaIj3i_a4,18237
-flaxdiff-0.1.35.6.dist-info/METADATA,sha256=NVCk5V7Zc3iq-nrWTivzO17dQa1fIjYgjJb800ZrZhQ,22085
-flaxdiff-0.1.35.6.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-flaxdiff-0.1.35.6.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.35.6.dist-info/RECORD,,