From c0ce19ad846ef0c1084437ed361ce1cf009cde6a Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 15:09:35 +0100
Subject: [PATCH 01/12] most of first draft, no testing yet. Still some
redundancy btwn trainer and model. Need to look at load/save usages also
---
.../utils_dataclass/real_data_gen.py | 29 ++--
model/trainer.py | 37 ++---
run/parameter.py | 127 +++++++++++-------
run/train_model.py | 2 +-
4 files changed, 98 insertions(+), 97 deletions(-)
diff --git a/dataclass_creator/utils_dataclass/real_data_gen.py b/dataclass_creator/utils_dataclass/real_data_gen.py
index b362c84..77216d0 100644
--- a/dataclass_creator/utils_dataclass/real_data_gen.py
+++ b/dataclass_creator/utils_dataclass/real_data_gen.py
@@ -11,13 +11,11 @@ from os.path import join
from pathlib import Path
from typing import Any, Iterable
-#from optree import tree_map
-
from opt_tools import SysVarSet, Traj, Batch, tree_map_relaxed
from ros_datahandler import get_bags_iter
from ros_datahandler import JointPosData, ForceData, JointState, WrenchStamped
from dataclass_creator.franka_dataclass import FrankaData
-
+from run.parameter import DatasetParams
def min_max_normalize_data(data: jnp.ndarray):
for i in range(data.shape[-1]):
@@ -105,7 +103,7 @@ def bag_info(root='data/without_search_strategy/', element: int = 0):
print('Topics of one Rosbag:', topics)
-def load_dataset(root='experiments/example_data/train', name: str = 'SUCCESS') -> Iterable[Traj]:
+def load_dataset(dp: DatasetParams) -> Iterable[Traj]:
"""collects rosbags in specified dir and return a formatted dataclass containg ragged list*3
Args:
@@ -116,26 +114,15 @@ def load_dataset(root='experiments/example_data/train', name: str = 'SUCCESS') -
FormattedDataclass: data containing the modalities corresponding to the attributes initialized for RosDatahandler
"""
- bags = glob.glob(join(root, f'*{name}*.bag'))
- assert len(bags) > 0, f"I expected bags! No bags on {root}"
- bags = [Path(bag) for bag in bags]
+ bags = glob.glob(join(dp.root, dp.glob_filter))
+ assert len(bags) > 0, f"I expected bags! No bags on {dp.root} with filter {dp.glob_filter}"
- print(f'Containing # {len(bags)} rosbags, named {name}')
-
- init_frankadata = FrankaData(jt = JointPosData('/joint_states', JointState),
- force = ForceData('/franka_state_controller/F_ext', WrenchStamped))
-
- # @kevin 2.7.24: data is now a list of Trajs; leaves are np arrays for each bags
- data = get_bags_iter(bags, init_frankadata, sync_stream='force')
+ bags = [Path(bag) for bag in bags]
+ logging.info(f'Containing # {len(bags)} rosbags for filter {dp.glob_filter}')
+ data = get_bags_iter(bags, dp.datatype, sync_stream=dp.sync_stream)
data = [d.cast_traj() for d in data]
-
- assert len(data) == len(bags)
- assert isinstance(data[0], FrankaData), 'data is somehow not a FrankaData? huh?'
- assert isinstance(data[0], SysVarSet), 'data is somehow not a sysvarset? huh?'
- assert isinstance(data[0], Traj), 'data should be a traj'
- assert not isinstance(data[0], Batch), 'data should not be a batch yet'
-
+
return data
diff --git a/model/trainer.py b/model/trainer.py
index 621c18f..216b7c6 100644
--- a/model/trainer.py
+++ b/model/trainer.py
@@ -14,7 +14,7 @@ from flax.training.train_state import TrainState
from tqdm.auto import tqdm
from .autoencoder import AutoEncoder
-from run.parameter import LoggerParams, OptimizerHParams
+from run.parameter import LoggerParams, TrainerConfig
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from utils.loss_calculation import LossCalculator, ReconstructionLoss, PredictionLoss
from opt_tools import SysVarSet, Batch
@@ -28,17 +28,6 @@ def train_step(loss_grad_fn: LossCalculator.loss,
opt_state = opt_state.apply_gradients(grads=grads)
return opt_state, batch_loss
-@dataclass
-class TrainerConfig:
- purpose: str
- optimizer_hparams: OptimizerHParams
- model: AutoEncoder
- logger_params: LoggerParams
- seed: int = 0
- n_epochs: int = 50
- check_val_every_n_epoch: int = 1
-
-
class Trainer:
"""
Train the model to the dataset, using optimizer and logger settings from config.
@@ -58,14 +47,13 @@ class Trainer:
prep_batches: if true, compile batches before starting
"""
- self.n_epochs = config.n_epochs
+ self.epochs = config.epochs
self.check_val_every_n_epoch = config.check_val_every_n_epoch
self.batches = train_dataloader.prep_batches() if prep_batches else train_dataloader
if val_dataloader:
raise NotImplementedError("val_dataloader not implemented")
-
self.state = self.init_state(config, train_dataloader.get_init_vec(), len(train_dataloader))
#self.logger = self.init_logger(config.logger_params)
self.train_step = self.init_train_step(config.purpose, config.model)
@@ -83,7 +71,7 @@ class Trainer:
init_vec (SysVarsSet): Initialization input vector of model
seed (int): start for PRNGKey
"""
- optimizer = config.optimizer_hparams.create(config.n_epochs,
+ optimizer = config.optimizer_hparams.create(config.epochs,
n_steps_per_epoch)
params = config.model.init(jax.random.PRNGKey(config.seed), init_vec)
@@ -103,19 +91,14 @@ class Trainer:
return tf.summary.create_file_writer(log_dir)
@staticmethod
- def init_train_step(purpose: str,
- model: AutoEncoder
+ def init_train_step(config: TrainerConfig,
+ model: AutoEncoder,
+ loss: Callable[[Model], float]
) -> Callable[[optax.OptState, Batch], float]:
"""
Initializes the loss_fn matching the purpose.
"""
- if purpose.lower() == 'reconstruction':
- loss = ReconstructionLoss(model).batch_loss
- elif purpose.lower() == 'prediction':
- loss = PredictionLoss(model).batch_loss
- else:
- raise TypeError(f"Unknown purpose {purpose}")
- loss_grad_fn = jax.value_and_grad(loss)
+ loss_grad_fn = jax.value_and_grad(config.loss)
# We attach the loss_grad_fn into train_step so we dont have to keep track of it
return jit(partial(train_step, loss_grad_fn))
@@ -147,8 +130,8 @@ class Trainer:
epoch_loss_array = epoch_loss_array.at[i].set(jnp.mean(batch_loss_array))
return state, epoch_loss_array
- loss_array = jnp.empty(self.n_epochs)
- for i in tqdm(range(1, self.n_epochs+1)):
+ loss_array = jnp.empty(self.epochs)
+ for i in tqdm(range(1, self.epochs+1)):
self.state, loss_array = epoch(i-1, (self.state, loss_array))
#with self.logger.as_default():
@@ -170,7 +153,7 @@ class Trainer:
"""
loss_array = []
- for i in tqdm(range(1, self.n_epochs + 1)):
+ for i in tqdm(range(1, self.epochs + 1)):
batch_loss_array = []
for batch in self.batches:
self.state, batch_loss = self.train_step(self.state, batch)
diff --git a/run/parameter.py b/run/parameter.py
index 51964c4..899064c 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -2,10 +2,75 @@ import optax
import os
from pathlib import Path
import logging
+from typing import Iterable
from dataclasses import dataclass
from datetime import datetime
+from opt_tools.jax_tools import SysVarSet, Traj
+from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
+from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+
+logging.addLevelName(logging.INFO,
+ "{}{}{}".format('\033[92m',
+ logging.getLevelName(logging.INFO),
+ '\033[0m'))
+@dataclass
+class DatasetParams:
+ # Dataset
+ datatype: SysVarSet
+ sync_stream: str
+ root: Path = 'experiments/example_data/train'
+ glob_filter: str = "*SUCCESS*.bag"
+
+ # Dataloader
+ batch_size: int
+ # Dataloader, also shared with model
+ window_size: int = None
+ pred_offset: int = 0
+
+ def __post_init__(self):
+ assert hasattr(self.datatype, self.sync_stream), \
+ f"The sync_stream should be an attribute of datatype, \
+ got sync_stream {self.sync_stream} and datatype {self.datatype}"
+
+ def get_dataloader(self) -> DynSliceDataloader:
+ data = load_dataset(self)
+ return DynSliceDataloader(data,
+ window_size=self.window_size,
+ batch_size=self.batch_size,
+ pred_offset=self.pred_offset)
+
+@dataclass
+class ModelParams:
+ root: Path # previously path_to_models
+ latent_dim: int
+ window_size: int
+ pred_offset: int = 0 # only relevant for prediction models
+
+ @property
+ def path_to_model(self):
+ return os.path.join(self.root, Path(f"{self}"))
+
+ def save(self, model, optimized_params):
+ if not os.path.exists(self.path_to_model):
+ os.makedirs(self.path_to_model)
+ time_stamp = datetime.now().strftime("%d_%m_%Y-%H_%M_%S")
+ logging.info(f"Path to logging directory: {self.path_to_model} \n Time: {time_stamp}")
+
+ to_dill((model, optimized_params, self), path_to_model, "model.dill")
+
+ def load(self):
+ assert os.path.exists(os.path.join(self.path_to_model, "model.dill")), \
+ f"No model.dill found in {self.path_to_model}, did you train?"
+ return from_dill(
+ def __str__(self):
+ return f"w{self.window_size}_l{self.latent_dim}"
+
+@dataclass
+class AEParams(ModelParams):
+ c_hid: int = 50 # parameters in hidden layer
+
@dataclass
class OptimizerHParams:
@@ -13,8 +78,8 @@ class OptimizerHParams:
lr: float = 1e-5
schedule: bool = False
warmup: float = 0.0
-
- def create(self, n_epochs: int, n_steps_per_epoch: int):
+
+ def create(self, epochs: int = None, n_steps_per_epoch: int = None):
"""
Initializes the optimizer with learning rate scheduler/constant learning rate.
@@ -25,62 +90,28 @@ class OptimizerHParams:
optimizer_class = self.optimizer
lr = self.lr
if self.schedule:
+ assert epochs is not None and n_steps_per_epoch is not None, "Args required when using scheduler"
lr = optax.warmup_cosine_decay_schedule(
initial_value=0.0,
peak_value=lr,
warmup_steps=self.warmup,
- decay_steps=int(n_epochs * n_steps_per_epoch),
+ decay_steps=int(epochs * n_steps_per_epoch),
end_value=0.01 * lr
)
return optimizer_class(lr)
-class TerminalColour:
- """
- Terminal colour formatting codes, essential not to be misled by red logging messages, as everything is smooth
- """
- MAGENTA = '\033[95m'
- BLUE = '\033[94m'
- GREEN = '\033[92m'
- YELLOW = '\033[93m'
- RED = '\033[91m'
- GREY = '\033[0m' # normal
- WHITE = '\033[1m' # bright white
- UNDERLINE = '\033[4m'
-
-class LoggerParams:
-
- def __init__(self, window_size, batch_size, epochs, latent_dim, path_to_models):
- self.window_size = window_size
- self.batch_size = batch_size
- self.epochs = epochs
- self.latent_dim = latent_dim
- self.path_to_models = path_to_models
- self.path_to_model = os.path.join(self.path_to_models, Path(f"test_w{window_size}_b{batch_size}_e{epochs}_l{latent_dim}"))
- self.time_stamp = datetime.now().strftime("%d_%m_%Y-%H_%M_%S")
-
- logging.basicConfig(level=logging.INFO)
- logging.addLevelName(logging.INFO,
- "{}{}{}".format(TerminalColour.GREEN, logging.getLevelName(logging.INFO), TerminalColour.GREY))
- if not os.path.exists(self.path_to_model):
- os.makedirs(self.path_to_model)
- print(f"Path to logging directory: {self.path_to_model} \n Time: {self.time_stamp}")
-
-
+
@dataclass
-class AEParams:
- c_hid: int = 50 # parameters in hidden layer
- bottleneck_size: int = 20 # latent_dim
-
- #@kevin 10.8.24: a major advantage of dataclass is the automatic __init__,
- # by defining __init__ by hand, this is overwritten so we cant
- # call AEParams(c_hid=50). __postinit__ is intended for any
- # checks, etc.
- def __postinit__(self):
- print(f"Autoencoder fixed parameter: \n "
- f"Parameters in hidden layer: {self.c_hid} \n "
- f"Dimension latent space: {self.bottleneck_size}")
-
+class TrainerConfig:
+ optimizer_hparams: OptimizerHParams
+ model: AutoEncoder
+ loss: LossCalculator
+
+ seed: int = 0
+ epochs: int = 50
+ check_val_every_n_epoch: int = 1
+
@dataclass
class RepoPaths:
"""Path names for 01_faston_converging anomaly detection."""
diff --git a/run/train_model.py b/run/train_model.py
index 18baac1..6d580ae 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -18,7 +18,7 @@ def train_model(train_loader: DynSliceDataloader,
val_loader: DynSliceDataloader,
logger_params: LoggerParams,
n_epochs: int,
- purpose: str = 'Reconstruction'
+ model:
):
# --------- initialize autoencoder ---------------------------------------------------------------------------------
--
GitLab
From b57b12083a4395997f087537bb504b30ae5e62de Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 18:04:01 +0100
Subject: [PATCH 02/12] most of first draft sketched. Still need to think
about handling load/save (do we need to load loggerparams if we get all the
model params from the model? How to get name?). Few issues with circular
commits (might be better to put params with their corresponding class.
---
dataclass_creator/create_dataloader.py | 2 +-
dataclass_creator/dyn_slice_dataloader.py | 2 +-
dataclass_creator/dyn_slice_dataset.py | 2 +-
dataclass_creator/franka_dataclass.py | 2 +-
dataclass_creator/synth_dataclass.py | 2 +-
.../utils_dataclass/real_data_gen.py | 2 +-
.../utils_dataclass/synth_data_gen.py | 2 +-
model/__init__.py | 2 +-
model/autoencoder.py | 2 +-
model/trainer.py | 14 ++---
opt_tools | 2 +-
run/load_save_model.py | 4 +-
run/parameter.py | 59 +++++-------------
run/train_model.py | 60 ++++++++++++-------
utils/loss_calculation.py | 2 +-
15 files changed, 73 insertions(+), 86 deletions(-)
diff --git a/dataclass_creator/create_dataloader.py b/dataclass_creator/create_dataloader.py
index 14c3e8a..956facd 100644
--- a/dataclass_creator/create_dataloader.py
+++ b/dataclass_creator/create_dataloader.py
@@ -8,7 +8,7 @@ from .synth_dataclass import SynthData
from .utils_dataclass.synth_data_gen import DataGeneration
from .dyn_slice_dataloader import DynSliceDataloader, split_dataloader
-from opt_tools import TrajBatch, Traj
+from opt_tools.jax_tools import TrajBatch, Traj
def generate_synth_data(amp_rest_pos: float) -> Iterable[Traj]:
diff --git a/dataclass_creator/dyn_slice_dataloader.py b/dataclass_creator/dyn_slice_dataloader.py
index 87fef67..c38f1fe 100644
--- a/dataclass_creator/dyn_slice_dataloader.py
+++ b/dataclass_creator/dyn_slice_dataloader.py
@@ -6,7 +6,7 @@ import jax.numpy as jnp
from jax import Array
from jax.flatten_util import ravel_pytree
-from opt_tools import Traj, Batch
+from opt_tools.jax_tools import Traj, Batch
from dataclass_creator.dyn_slice_dataset import DynSliceDataset, PredDynSliceDataset
def split_dataloader(data: Iterable[Traj], train=0.7, val=0.2, **kwargs
diff --git a/dataclass_creator/dyn_slice_dataset.py b/dataclass_creator/dyn_slice_dataset.py
index 030b1c2..db0ba8b 100644
--- a/dataclass_creator/dyn_slice_dataset.py
+++ b/dataclass_creator/dyn_slice_dataset.py
@@ -7,7 +7,7 @@ from functools import partial
from typing import Iterable, Tuple
from jax.tree_util import tree_map
-from opt_tools import Traj, Batch
+from opt_tools.jax_tools import Traj, Batch
class DynSliceDataset:
diff --git a/dataclass_creator/franka_dataclass.py b/dataclass_creator/franka_dataclass.py
index 4d4e1df..40d6b1d 100644
--- a/dataclass_creator/franka_dataclass.py
+++ b/dataclass_creator/franka_dataclass.py
@@ -2,7 +2,7 @@ import jax.numpy as jnp
from flax.struct import dataclass
-from opt_tools import SysVarSet
+from opt_tools.jax_tools import SysVarSet
@dataclass
class FrankaData(SysVarSet):
diff --git a/dataclass_creator/synth_dataclass.py b/dataclass_creator/synth_dataclass.py
index 389fafd..f5d1a16 100644
--- a/dataclass_creator/synth_dataclass.py
+++ b/dataclass_creator/synth_dataclass.py
@@ -2,7 +2,7 @@ import jax.numpy as jnp
from flax.struct import dataclass
-from opt_tools import SysVarSet
+from opt_tools.jax_tools import SysVarSet
@dataclass
class SynthData(SysVarSet):
diff --git a/dataclass_creator/utils_dataclass/real_data_gen.py b/dataclass_creator/utils_dataclass/real_data_gen.py
index 77216d0..a173358 100644
--- a/dataclass_creator/utils_dataclass/real_data_gen.py
+++ b/dataclass_creator/utils_dataclass/real_data_gen.py
@@ -11,7 +11,7 @@ from os.path import join
from pathlib import Path
from typing import Any, Iterable
-from opt_tools import SysVarSet, Traj, Batch, tree_map_relaxed
+from opt_tools.jax_tools import SysVarSet, Traj, Batch, tree_map_relaxed
from ros_datahandler import get_bags_iter
from ros_datahandler import JointPosData, ForceData, JointState, WrenchStamped
from dataclass_creator.franka_dataclass import FrankaData
diff --git a/dataclass_creator/utils_dataclass/synth_data_gen.py b/dataclass_creator/utils_dataclass/synth_data_gen.py
index 2ea4dd2..312306b 100644
--- a/dataclass_creator/utils_dataclass/synth_data_gen.py
+++ b/dataclass_creator/utils_dataclass/synth_data_gen.py
@@ -5,7 +5,7 @@ import jax.random as random
import matplotlib.pyplot as plt
from opt_tools.jax_tools.mds_sys import MDSSys, StepParams
-from opt_tools import to_batch
+from opt_tools.jax_tools import to_batch
class DataGeneration():
''' This class is used for generating synthetic data.
diff --git a/model/__init__.py b/model/__init__.py
index 4910b8d..852ade4 100644
--- a/model/__init__.py
+++ b/model/__init__.py
@@ -1,2 +1,2 @@
from .autoencoder import AutoEncoder
-from .trainer import Trainer, TrainerConfig
+from .trainer import Trainer
diff --git a/model/autoencoder.py b/model/autoencoder.py
index 78619f8..de073c7 100644
--- a/model/autoencoder.py
+++ b/model/autoencoder.py
@@ -3,7 +3,7 @@ import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
from flax import linen as nn
-from opt_tools import SysVarSet
+from opt_tools.jax_tools import SysVarSet
class Encoder(nn.Module):
diff --git a/model/trainer.py b/model/trainer.py
index 216b7c6..06ab5c9 100644
--- a/model/trainer.py
+++ b/model/trainer.py
@@ -14,10 +14,10 @@ from flax.training.train_state import TrainState
from tqdm.auto import tqdm
from .autoencoder import AutoEncoder
-from run.parameter import LoggerParams, TrainerConfig
+from run.parameter import TrainerParams
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from utils.loss_calculation import LossCalculator, ReconstructionLoss, PredictionLoss
-from opt_tools import SysVarSet, Batch
+from opt_tools.jax_tools import SysVarSet, Batch
def train_step(loss_grad_fn: LossCalculator.loss,
opt_state: optax.OptState,
@@ -33,7 +33,7 @@ class Trainer:
Train the model to the dataset, using optimizer and logger settings from config.
"""
def __init__(self,
- config: TrainerConfig,
+ config: TrainerParams,
train_dataloader: DynSliceDataloader,
val_dataloader: DynSliceDataloader = None,
prep_batches:bool = False
@@ -42,7 +42,7 @@ class Trainer:
Init to call the Trainer and enfold the magic.
Args:
- config (TrainerConfig): Configuration
+ config (TrainerParams): Configuration
train_dataloader (Dataloader): Dataset with iterable of batches
prep_batches: if true, compile batches before starting
"""
@@ -59,7 +59,7 @@ class Trainer:
self.train_step = self.init_train_step(config.purpose, config.model)
@staticmethod
- def init_state(config: TrainerConfig,
+ def init_state(config: TrainerParams,
init_vec: SysVarSet,
n_steps_per_epoch: int
) -> TrainState:
@@ -91,9 +91,9 @@ class Trainer:
return tf.summary.create_file_writer(log_dir)
@staticmethod
- def init_train_step(config: TrainerConfig,
+ def init_train_step(config: TrainerParams,
model: AutoEncoder,
- loss: Callable[[Model], float]
+ loss: Callable
) -> Callable[[optax.OptState, Batch], float]:
"""
Initializes the loss_fn matching the purpose.
diff --git a/opt_tools b/opt_tools
index 3872400..cfb942f 160000
--- a/opt_tools
+++ b/opt_tools
@@ -1 +1 @@
-Subproject commit 38724006b8735cba3f1fa96669ddc74b767d22a0
+Subproject commit cfb942f0ef3f140013e29b60462f9d0f4fff5964
diff --git a/run/load_save_model.py b/run/load_save_model.py
index f29f8ae..c9a00c9 100644
--- a/run/load_save_model.py
+++ b/run/load_save_model.py
@@ -7,7 +7,7 @@ from typing import Callable, Any, Tuple
import dill
import jax
-from .parameter import LoggerParams
+#from .parameter import LoggerParams
logger = logging.getLogger()
@@ -60,7 +60,7 @@ def load_trained_model(path_to_models: Path, test_name: str):
return (ae_dll, params_dll, logger_params_dll)
-def load_trained_model_jit(path_to_models: Path, test_name: str) -> Tuple[Callable[[Any], Any], LoggerParams]:
+def load_trained_model_jit(path_to_models: Path, test_name: str) -> Tuple[Callable[[Any], Any], Any]:
"""Directly return a callable with your precious model."""
path_to_model = os.path.join(path_to_models, test_name)
diff --git a/run/parameter.py b/run/parameter.py
index 899064c..64771d4 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -1,30 +1,31 @@
-import optax
import os
from pathlib import Path
import logging
-from typing import Iterable
-
-from dataclasses import dataclass
+from typing import Iterable, Any
+from dataclasses import dataclass, field
from datetime import datetime
+import optax
+
from opt_tools.jax_tools import SysVarSet, Traj
-from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+
logging.addLevelName(logging.INFO,
"{}{}{}".format('\033[92m',
logging.getLevelName(logging.INFO),
'\033[0m'))
@dataclass
class DatasetParams:
+ """Parameters to define the data source, batching, etc."""
# Dataset
datatype: SysVarSet
sync_stream: str
- root: Path = 'experiments/example_data/train'
+ root: Path# = Path('experiments/example_data/train')
glob_filter: str = "*SUCCESS*.bag"
# Dataloader
- batch_size: int
+ batch_size: int = 50
# Dataloader, also shared with model
window_size: int = None
pred_offset: int = 0
@@ -40,38 +41,7 @@ class DatasetParams:
window_size=self.window_size,
batch_size=self.batch_size,
pred_offset=self.pred_offset)
-
-@dataclass
-class ModelParams:
- root: Path # previously path_to_models
- latent_dim: int
- window_size: int
- pred_offset: int = 0 # only relevant for prediction models
-
- @property
- def path_to_model(self):
- return os.path.join(self.root, Path(f"{self}"))
-
- def save(self, model, optimized_params):
- if not os.path.exists(self.path_to_model):
- os.makedirs(self.path_to_model)
- time_stamp = datetime.now().strftime("%d_%m_%Y-%H_%M_%S")
- logging.info(f"Path to logging directory: {self.path_to_model} \n Time: {time_stamp}")
-
- to_dill((model, optimized_params, self), path_to_model, "model.dill")
-
- def load(self):
- assert os.path.exists(os.path.join(self.path_to_model, "model.dill")), \
- f"No model.dill found in {self.path_to_model}, did you train?"
- return from_dill(
- def __str__(self):
- return f"w{self.window_size}_l{self.latent_dim}"
-
-@dataclass
-class AEParams(ModelParams):
- c_hid: int = 50 # parameters in hidden layer
-
-
+
@dataclass
class OptimizerHParams:
optimizer: optax = optax.adam
@@ -99,13 +69,14 @@ class OptimizerHParams:
end_value=0.01 * lr
)
return optimizer_class(lr)
-
@dataclass
-class TrainerConfig:
- optimizer_hparams: OptimizerHParams
- model: AutoEncoder
- loss: LossCalculator
+class TrainerParams:
+ """Define how the traiing process should happen."""
+ model: Any #AutoEncoder
+ loss: Any #LossCalculator
+
+ optimizer_hparams: OptimizerHParams = field(default_factory=lambda: OptimizerHParams())
seed: int = 0
epochs: int = 50
diff --git a/run/train_model.py b/run/train_model.py
index 6d580ae..8d087e5 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -1,13 +1,14 @@
import os
-import jax
-
from itertools import product
from pathlib import Path
+import jax
+
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
-from model import AutoEncoder, Trainer, TrainerConfig
-from run.parameter import LoggerParams, AEParams, OptimizerHParams, RepoPaths
+from dataclass_creator.franka_dataclass import FrankaData
+from model import AutoEncoder, Trainer
+from run.parameter import OptimizerHParams, RepoPaths, TrainerParams, DatasetParams
from run.load_save_model import to_dill
print("Default device:", jax.default_backend())
@@ -16,9 +17,9 @@ print("Available devices:", jax.devices())
def train_model(train_loader: DynSliceDataloader,
val_loader: DynSliceDataloader,
- logger_params: LoggerParams,
+ logger_params,
n_epochs: int,
- model:
+ model
):
# --------- initialize autoencoder ---------------------------------------------------------------------------------
@@ -26,8 +27,8 @@ def train_model(train_loader: DynSliceDataloader,
ae = AutoEncoder(AEParams.c_hid, logger_params.latent_dim, train_loader.out_size)
# --------- initialize trainer -------------------------------------------------------------------------------------
- config = TrainerConfig(
- purpose=purpose,
+ config = TrainerParams(
+ #purpose=purpose,
optimizer_hparams=OptimizerHParams(),
model=ae,
logger_params=logger_params,
@@ -49,12 +50,13 @@ def train(path_to_data: Path, window_size: int, batch_size: int, epochs: int, la
train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- logger_params = LoggerParams(window_size=window_size,
- batch_size=batch_size,
- epochs=epochs,
- latent_dim=latent_dim,
- path_to_models=RepoPaths.trained_models_val)
+# logger_params = LoggerParams(window_size=window_size,
+# batch_size=batch_size,
+# epochs=epochs,
+# latent_dim=latent_dim,
+# path_to_models=RepoPaths.trained_models_val)
+ logger_params = None
ae, optimized_params, threshold = train_model(train_loader,
val_loader,
logger_params,
@@ -83,14 +85,28 @@ def train_loop(path_to_data: Path, window_size: list, batch_size: list, epochs:
train(path_to_data=path_to_data, window_size=window_size, batch_size=batch_size, epochs=epochs, latent_dim=latent_dim, purpose=purpose, train_data=train_data)
if __name__ == "__main__":
- train_config = {
- "path_to_data": RepoPaths.data_train,
- "window_size": [100],
- "batch_size": [20,50],
- "epochs": [30,50],
- "latent_dim": [15,20],
- "purpose": "Reconstruction"
- }
- train_loop(**train_config)
+ data = DatasetParams(
+ datatype=FrankaData,
+ sync_stream='force',
+ root = Path('experiments/example_data/mytest'),
+ window_size=100,
+ ).get_dataloader()
+
+ model = AE(
+ c_hid=50,
+ latent_dim=15,
+ c_out=data.out_size
+ )
+
+ to_dill(model, ".", "test.dill")
+ model = from_dill("./test.dill")
+
+ trainer_params = TrainerParams(
+ model=model,
+ loss=test_loss
+ )
+
+ train(data, model, trainer_params)
+
os._exit(1)
diff --git a/utils/loss_calculation.py b/utils/loss_calculation.py
index c589bda..d7f0ab9 100644
--- a/utils/loss_calculation.py
+++ b/utils/loss_calculation.py
@@ -6,7 +6,7 @@ import optax
from typing import Protocol, Union, Tuple
from model.autoencoder import AutoEncoder
-from opt_tools import SysVarSet, Batch, tree_map_relaxed
+from opt_tools.jax_tools import SysVarSet, Batch, tree_map_relaxed
# Protocol for the Loss Fn
--
GitLab
From e53f7be8a48cff9d2b0768f6788752c435fe8f9a Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 20:50:44 +0100
Subject: [PATCH 03/12] triggering through to training, shape error, looks like
batch is getting rolled into next dim
---
.../utils_dataclass/real_data_gen.py | 17 +++++----
model/trainer.py | 19 +++-------
run/load_save_model.py | 5 +++
run/parameter.py | 20 +++++++----
run/train_model.py | 35 ++++++++++---------
5 files changed, 54 insertions(+), 42 deletions(-)
diff --git a/dataclass_creator/utils_dataclass/real_data_gen.py b/dataclass_creator/utils_dataclass/real_data_gen.py
index a173358..939c5a4 100644
--- a/dataclass_creator/utils_dataclass/real_data_gen.py
+++ b/dataclass_creator/utils_dataclass/real_data_gen.py
@@ -15,7 +15,7 @@ from opt_tools.jax_tools import SysVarSet, Traj, Batch, tree_map_relaxed
from ros_datahandler import get_bags_iter
from ros_datahandler import JointPosData, ForceData, JointState, WrenchStamped
from dataclass_creator.franka_dataclass import FrankaData
-from run.parameter import DatasetParams
+
def min_max_normalize_data(data: jnp.ndarray):
for i in range(data.shape[-1]):
@@ -103,7 +103,12 @@ def bag_info(root='data/without_search_strategy/', element: int = 0):
print('Topics of one Rosbag:', topics)
-def load_dataset(dp: DatasetParams) -> Iterable[Traj]:
+def load_dataset(
+ datatype: SysVarSet,
+ sync_stream: str,
+ root: Path,
+ glob_filter: str = "",
+) -> Iterable[Traj]:
"""collects rosbags in specified dir and return a formatted dataclass containg ragged list*3
Args:
@@ -114,13 +119,13 @@ def load_dataset(dp: DatasetParams) -> Iterable[Traj]:
FormattedDataclass: data containing the modalities corresponding to the attributes initialized for RosDatahandler
"""
- bags = glob.glob(join(dp.root, dp.glob_filter))
- assert len(bags) > 0, f"I expected bags! No bags on {dp.root} with filter {dp.glob_filter}"
+ bags = glob.glob(join(root, glob_filter))
+ assert len(bags) > 0, f"I expected bags! No bags on {root} with filter {glob_filter}"
bags = [Path(bag) for bag in bags]
- logging.info(f'Containing # {len(bags)} rosbags for filter {dp.glob_filter}')
+ logging.info(f'Containing # {len(bags)} rosbags for filter {glob_filter}')
- data = get_bags_iter(bags, dp.datatype, sync_stream=dp.sync_stream)
+ data = get_bags_iter(bags, datatype, sync_stream=sync_stream)
data = [d.cast_traj() for d in data]
return data
diff --git a/model/trainer.py b/model/trainer.py
index 06ab5c9..d3a7b7f 100644
--- a/model/trainer.py
+++ b/model/trainer.py
@@ -54,9 +54,12 @@ class Trainer:
if val_dataloader:
raise NotImplementedError("val_dataloader not implemented")
+
self.state = self.init_state(config, train_dataloader.get_init_vec(), len(train_dataloader))
#self.logger = self.init_logger(config.logger_params)
- self.train_step = self.init_train_step(config.purpose, config.model)
+
+ loss_grad_fn = jax.value_and_grad(config.loss)
+ self.train_step = jit(partial(train_step, loss_grad_fn))
@staticmethod
def init_state(config: TrainerParams,
@@ -90,19 +93,6 @@ class Trainer:
log_dir = os.path.join(logger_params.path_to_model + "/log_file")
return tf.summary.create_file_writer(log_dir)
- @staticmethod
- def init_train_step(config: TrainerParams,
- model: AutoEncoder,
- loss: Callable
- ) -> Callable[[optax.OptState, Batch], float]:
- """
- Initializes the loss_fn matching the purpose.
- """
- loss_grad_fn = jax.value_and_grad(config.loss)
- # We attach the loss_grad_fn into train_step so we dont have to keep track of it
- return jit(partial(train_step, loss_grad_fn))
-
-
def train_jit(self) -> Tuple[optax.OptState, jnp.ndarray]:
"""
Train, but using a fori_loop over batches. This puts the batch fetching and resulting
@@ -171,3 +161,4 @@ class Trainer:
# at this point evaluation, check point saving and early break could be added
return self.state.params, loss_array
+
diff --git a/run/load_save_model.py b/run/load_save_model.py
index c9a00c9..e34275a 100644
--- a/run/load_save_model.py
+++ b/run/load_save_model.py
@@ -53,6 +53,11 @@ def from_dill(path: str, file_name: str):
print("Undilling is dead!")
return loaded_data
+def get_config_string(model, dataset_params, trainer_params) -> str:
+ return f"test_w{dataset_params.window_size}\
+ _b{dataset_params.batch_size}\
+ _e{trainer_params.epochs}\
+ _l{model.latent_dim}"
def load_trained_model(path_to_models: Path, test_name: str):
path_to_model = os.path.join(path_to_models, test_name)
diff --git a/run/parameter.py b/run/parameter.py
index 64771d4..0e89c1d 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -9,6 +9,7 @@ import optax
from opt_tools.jax_tools import SysVarSet, Traj
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
logging.addLevelName(logging.INFO,
@@ -36,12 +37,19 @@ class DatasetParams:
got sync_stream {self.sync_stream} and datatype {self.datatype}"
def get_dataloader(self) -> DynSliceDataloader:
- data = load_dataset(self)
- return DynSliceDataloader(data,
- window_size=self.window_size,
- batch_size=self.batch_size,
- pred_offset=self.pred_offset)
-
+ data = load_dataset(
+ self.datatype,
+ self.sync_stream,
+ self.root,
+ self.glob_filter
+ )
+ return DynSliceDataloader(
+ data,
+ window_size=self.window_size,
+ batch_size=self.batch_size,
+ pred_offset=self.pred_offset
+ )
+
@dataclass
class OptimizerHParams:
optimizer: optax = optax.adam
diff --git a/run/train_model.py b/run/train_model.py
index 8d087e5..7764e01 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -8,8 +8,11 @@ from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from dataclass_creator.franka_dataclass import FrankaData
from model import AutoEncoder, Trainer
+
+from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
from run.parameter import OptimizerHParams, RepoPaths, TrainerParams, DatasetParams
-from run.load_save_model import to_dill
+from run.load_save_model import to_dill, from_dill
+from utils.loss_calculation import ReconstructionLoss
print("Default device:", jax.default_backend())
print("Available devices:", jax.devices())
@@ -84,29 +87,29 @@ def train_loop(path_to_data: Path, window_size: list, batch_size: list, epochs:
window_size, batch_size, epochs, latent_dim = hparams
train(path_to_data=path_to_data, window_size=window_size, batch_size=batch_size, epochs=epochs, latent_dim=latent_dim, purpose=purpose, train_data=train_data)
-if __name__ == "__main__":
-
- data = DatasetParams(
- datatype=FrankaData,
- sync_stream='force',
- root = Path('experiments/example_data/mytest'),
+if __name__ == "__main__":
+ train_data = DatasetParams(
+ datatype=FrankaData(
+ jt=JointPosData('/joint_states', JointState),
+ force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
+ ),
+ sync_stream = 'force',
window_size=100,
+ root = Path('experiments/example_data/train'),
).get_dataloader()
-
- model = AE(
+
+ model = AutoEncoder(
c_hid=50,
latent_dim=15,
- c_out=data.out_size
+ c_out=train_data.out_size
)
-
- to_dill(model, ".", "test.dill")
- model = from_dill("./test.dill")
trainer_params = TrainerParams(
model=model,
- loss=test_loss
+ loss=ReconstructionLoss(model).loss,
+ epochs=50
)
- train(data, model, trainer_params)
-
+ optimized_params, loss_array = Trainer(trainer_params, train_data).train_jit()
+
os._exit(1)
--
GitLab
From 4a58510773b820328d17db345caa820ff5f8cffb Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 20:54:58 +0100
Subject: [PATCH 04/12] this mfer is mfing training
---
run/train_model.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/run/train_model.py b/run/train_model.py
index 7764e01..c3805f6 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -99,14 +99,14 @@ if __name__ == "__main__":
).get_dataloader()
model = AutoEncoder(
- c_hid=50,
+ c_hid=30,
latent_dim=15,
c_out=train_data.out_size
)
trainer_params = TrainerParams(
model=model,
- loss=ReconstructionLoss(model).loss,
+ loss=ReconstructionLoss(model).batch_loss,
epochs=50
)
--
GitLab
From e5e91f3886695008430cbf51406a4a4052894202 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 21:17:49 +0100
Subject: [PATCH 05/12] add pickling class
---
run/load_save_model.py | 10 ++--------
run/parameter.py | 22 +++++++++++++++++++++-
run/train_model.py | 37 +++++++++++++++++++++++++++----------
3 files changed, 50 insertions(+), 19 deletions(-)
diff --git a/run/load_save_model.py b/run/load_save_model.py
index e34275a..da60d0b 100644
--- a/run/load_save_model.py
+++ b/run/load_save_model.py
@@ -39,8 +39,8 @@ def to_dill(data_in: any, path: str, file_name: str):
try:
with open(name, 'wb') as f:
dill.dump(data_in, f)
- except:
- print("Dill is dead!")
+ except Exception as e:
+ logging.error(f"Dill is dead! {e}")
def from_dill(path: str, file_name: str):
@@ -53,12 +53,6 @@ def from_dill(path: str, file_name: str):
print("Undilling is dead!")
return loaded_data
-def get_config_string(model, dataset_params, trainer_params) -> str:
- return f"test_w{dataset_params.window_size}\
- _b{dataset_params.batch_size}\
- _e{trainer_params.epochs}\
- _l{model.latent_dim}"
-
def load_trained_model(path_to_models: Path, test_name: str):
path_to_model = os.path.join(path_to_models, test_name)
ae_dll, params_dll, logger_params_dll = from_dill(path_to_model, "model.dill")
diff --git a/run/parameter.py b/run/parameter.py
index 0e89c1d..674ed0f 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -10,7 +10,7 @@ import optax
from opt_tools.jax_tools import SysVarSet, Traj
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
-
+from run.load_save_model import to_dill, from_dill
logging.addLevelName(logging.INFO,
"{}{}{}".format('\033[92m',
@@ -90,6 +90,26 @@ class TrainerParams:
epochs: int = 50
check_val_every_n_epoch: int = 1
+@dataclass
+class ExportParams:
+ """Super-class to catch all those params we want to serialize / save."""
+ model: Any
+ model_params: Any
+ dataset_params: DatasetParams
+ trainer_params: TrainerParams
+
+ def __str__(self):
+ return f"test_w{self.dataset_params.window_size}"\
+ f"_b{self.dataset_params.batch_size}"\
+ f"_e{self.trainer_params.epochs}"\
+ f"_l{self.model.latent_dim}"
+
+
+ def save(self, root: Path=Path(".")):
+ path = os.path.join(root, Path(f"{self}"))
+ os.makedirs(path)
+ to_dill(self, path, "model.dill")
+
@dataclass
class RepoPaths:
diff --git a/run/train_model.py b/run/train_model.py
index c3805f6..2acbbbe 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -10,7 +10,13 @@ from dataclass_creator.franka_dataclass import FrankaData
from model import AutoEncoder, Trainer
from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
-from run.parameter import OptimizerHParams, RepoPaths, TrainerParams, DatasetParams
+from run.parameter import (
+ OptimizerHParams,
+ RepoPaths,
+ TrainerParams,
+ DatasetParams,
+ ExportParams
+)
from run.load_save_model import to_dill, from_dill
from utils.loss_calculation import ReconstructionLoss
@@ -87,8 +93,15 @@ def train_loop(path_to_data: Path, window_size: list, batch_size: list, epochs:
window_size, batch_size, epochs, latent_dim = hparams
train(path_to_data=path_to_data, window_size=window_size, batch_size=batch_size, epochs=epochs, latent_dim=latent_dim, purpose=purpose, train_data=train_data)
-if __name__ == "__main__":
- train_data = DatasetParams(
+def lazytrainer():
+ """
+ This refactor tries to:
+ - provide a separation of concerns (data, model, training)
+ - remove redundancy to the degree possible
+ - make it easier to try new models / loss fns, making only local changes
+ """
+
+ dataset_params = DatasetParams(
datatype=FrankaData(
jt=JointPosData('/joint_states', JointState),
force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
@@ -96,20 +109,24 @@ if __name__ == "__main__":
sync_stream = 'force',
window_size=100,
root = Path('experiments/example_data/train'),
- ).get_dataloader()
-
- model = AutoEncoder(
- c_hid=30,
- latent_dim=15,
- c_out=train_data.out_size
)
+
+ train_data = dataset_params.get_dataloader()
+
+ model = AutoEncoder(c_hid=30, latent_dim=15, c_out=train_data.out_size )
trainer_params = TrainerParams(
model=model,
loss=ReconstructionLoss(model).batch_loss,
- epochs=50
+ epochs=2
)
optimized_params, loss_array = Trainer(trainer_params, train_data).train_jit()
+
+ ExportParams(model, optimized_params, dataset_params, trainer_params).save()
+
+if __name__ == "__main__":
+ lazytrainer()
+
os._exit(1)
--
GitLab
From 20474e2149f504a4b9495e32946e80dcd67c1a68 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Sun, 15 Dec 2024 21:28:38 +0100
Subject: [PATCH 06/12] in the final pytest fixes
---
run/valmodel.py | 2 +-
tests/test_data_load.py | 13 +++++++++++--
tests/test_flattening.py | 6 +++---
tests/test_model_load.py | 2 +-
tests/test_synth_data.py | 4 ++--
tests/test_training.py | 4 ++--
utils/validation_fn.py | 2 +-
utils/visualize_fn.py | 2 +-
8 files changed, 22 insertions(+), 13 deletions(-)
diff --git a/run/valmodel.py b/run/valmodel.py
index 7c5b72b..bbd6c20 100644
--- a/run/valmodel.py
+++ b/run/valmodel.py
@@ -18,7 +18,7 @@ from run.load_save_model import load_trained_model
#for WindowsPath
#from run.load_save_test import load_trained_model
# Types
-from opt_tools import Batch, SysVarSet, Traj
+from opt_tools.jax_tools import Batch, SysVarSet, Traj
from typing import Iterable
diff --git a/tests/test_data_load.py b/tests/test_data_load.py
index b588ca2..160b509 100644
--- a/tests/test_data_load.py
+++ b/tests/test_data_load.py
@@ -10,7 +10,7 @@ from dataclass_creator.franka_dataclass import FrankaData
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
-from opt_tools import Batch, to_batch, tree_map_relaxed
+from opt_tools.jax_tools import Batch, to_batch, tree_map_relaxed
logger = logging.getLogger()
@@ -18,6 +18,8 @@ logger = logging.getLogger()
#root = os.fspath(Path(curr_dir.parent, 'experiments/example_data/train').resolve())
root = Path('experiments/example_data/train')
+datatype =
+
window_size = 20
batch_size = 10
n_epochs = 20
@@ -50,7 +52,14 @@ def test_save_load():
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
logger.info('Loading data from bags')
- raw_data = load_dataset(root, "SUCCESS")[0]
+ raw_data = load_dataset(
+ datatype=FrankaData(
+ jt=JointPosData('/joint_states', JointState),
+ force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
+ ),
+ sync_stream = 'force',
+ root,
+ "SUCCESS")[0]
# Try with some transformed data
treemapped_data = tree_map_relaxed(lambda l: l, raw_data).cast_batch()
diff --git a/tests/test_flattening.py b/tests/test_flattening.py
index 1c36bea..8ec1c39 100644
--- a/tests/test_flattening.py
+++ b/tests/test_flattening.py
@@ -6,12 +6,12 @@ from flax import linen as nn
from jax.tree_util import tree_leaves, tree_map
from jax.flatten_util import ravel_pytree
-from run.parameter import LoggerParams, OptimizerHParams, AEParams
+from run.parameter import OptimizerHParams
from dataclass_creator import create_synth_nominal_data
from dataclass_creator.synth_dataclass import SynthData
from dataclass_creator.utils_dataclass.synth_data_gen import DataGeneration
-from model import AutoEncoder, Trainer, TrainerConfig
-from opt_tools import SysVarSet
+from model import AutoEncoder, Trainer
+from opt_tools.jax_tools import SysVarSet
batch_size = 50
window_size = 80
diff --git a/tests/test_model_load.py b/tests/test_model_load.py
index 392eb69..138743e 100644
--- a/tests/test_model_load.py
+++ b/tests/test_model_load.py
@@ -5,7 +5,7 @@ import jax
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
-from run.parameter import LoggerParams, RepoPaths
+from run.parameter import RepoPaths
from run.train_model import train_model
from run.load_save_model import load_trained_model_jit, load_trained_model, to_dill
from tests.test_data_load import round_trip_dill, round_trip_pickle
diff --git a/tests/test_synth_data.py b/tests/test_synth_data.py
index d590289..826793d 100644
--- a/tests/test_synth_data.py
+++ b/tests/test_synth_data.py
@@ -12,11 +12,11 @@ from run.load_save_model import to_dill
from run.train_model import train_model
from run.valmodel import ValModel
from model import AutoEncoder
-from run.parameter import LoggerParams, RepoPaths
+from run.parameter import RepoPaths
# Types
from dataclass_creator.synth_dataclass import SynthData
-from opt_tools import Traj
+from opt_tools.jax_tools import Traj
curr_dir = Path(os.path.dirname(os.path.abspath(__file__)))
root_pretrained_models = os.fspath(Path(curr_dir.parent, 'trained_models/synth_data').resolve())
diff --git a/tests/test_training.py b/tests/test_training.py
index 791c8ab..91170f0 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -6,8 +6,8 @@ from pathlib import Path
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
-from model import AutoEncoder, Trainer, TrainerConfig
-from run.parameter import LoggerParams, AEParams, OptimizerHParams, RepoPaths
+from model import AutoEncoder, Trainer
+from run.parameter import OptimizerHParams, RepoPaths
from run.train_model import train_model
from run.load_save_model import to_dill
diff --git a/utils/validation_fn.py b/utils/validation_fn.py
index a6e85e2..ef7fa81 100644
--- a/utils/validation_fn.py
+++ b/utils/validation_fn.py
@@ -8,7 +8,7 @@ from dataclass_creator.franka_dataclass import FrankaData
# Types
from typing import List, Union, Tuple
-from opt_tools import Batch, tree_map_relaxed
+from opt_tools.jax_tools import Batch, tree_map_relaxed
def classify_windows(loss: Batch,
threshold: Batch):
diff --git a/utils/visualize_fn.py b/utils/visualize_fn.py
index 7fcb429..c396760 100644
--- a/utils/visualize_fn.py
+++ b/utils/visualize_fn.py
@@ -5,7 +5,7 @@ import jax
# Types
from typing import Iterable
from jax import Array
-from opt_tools import Batch
+from opt_tools.jax_tools import Batch
def plot_loss(loss_array: jnp.ndarray,
--
GitLab
From 6479a3c1a4d982b3fa28b86f1826f791ab09ad46 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Tue, 17 Dec 2024 18:30:34 +0100
Subject: [PATCH 07/12] pytests passing. Mostly there, just need to add some
loaders and think about if model should be in TrainerConfig or separate
---
dataclass_creator/create_dataloader.py | 3 +-
...nth_dataclass.py => custom_dataclasses.py} | 6 +-
dataclass_creator/franka_dataclass.py | 10 ---
.../utils_dataclass/real_data_gen.py | 10 +--
ros_adetector.py | 2 +-
run/construct_thresholds.py | 2 +-
run/parameter.py | 31 +++----
run/train_model.py | 90 ++++++++-----------
tests/test_data_load.py | 21 ++---
tests/test_defaults.py | 53 +++++++++++
tests/test_flattening.py | 2 +-
tests/test_model_load.py | 73 +++++----------
tests/test_synth_data.py | 15 ++--
tests/test_training.py | 66 ++++++--------
utils/validation_fn.py | 3 +-
15 files changed, 178 insertions(+), 209 deletions(-)
rename dataclass_creator/{synth_dataclass.py => custom_dataclasses.py} (69%)
delete mode 100644 dataclass_creator/franka_dataclass.py
create mode 100644 tests/test_defaults.py
diff --git a/dataclass_creator/create_dataloader.py b/dataclass_creator/create_dataloader.py
index 956facd..757b152 100644
--- a/dataclass_creator/create_dataloader.py
+++ b/dataclass_creator/create_dataloader.py
@@ -3,8 +3,7 @@ import jax.numpy as jnp
from typing import Tuple, Union, Iterable
from jax import tree_util
-from .franka_dataclass import FrankaData
-from .synth_dataclass import SynthData
+from .custom_dataclasses import SynthData
from .utils_dataclass.synth_data_gen import DataGeneration
from .dyn_slice_dataloader import DynSliceDataloader, split_dataloader
diff --git a/dataclass_creator/synth_dataclass.py b/dataclass_creator/custom_dataclasses.py
similarity index 69%
rename from dataclass_creator/synth_dataclass.py
rename to dataclass_creator/custom_dataclasses.py
index f5d1a16..f7e89ab 100644
--- a/dataclass_creator/synth_dataclass.py
+++ b/dataclass_creator/custom_dataclasses.py
@@ -1,9 +1,13 @@
import jax.numpy as jnp
-
from flax.struct import dataclass
from opt_tools.jax_tools import SysVarSet
+@dataclass
+class FrankaData(SysVarSet):
+ jt: jnp.ndarray = jnp.zeros(7)
+ force: jnp.ndarray = jnp.zeros(3)
+
@dataclass
class SynthData(SysVarSet):
pos: jnp.array = jnp.zeros(1)
diff --git a/dataclass_creator/franka_dataclass.py b/dataclass_creator/franka_dataclass.py
deleted file mode 100644
index 40d6b1d..0000000
--- a/dataclass_creator/franka_dataclass.py
+++ /dev/null
@@ -1,10 +0,0 @@
-import jax.numpy as jnp
-
-from flax.struct import dataclass
-
-from opt_tools.jax_tools import SysVarSet
-
-@dataclass
-class FrankaData(SysVarSet):
- jt: jnp.ndarray = jnp.zeros(7)
- force: jnp.ndarray = jnp.zeros(3)
diff --git a/dataclass_creator/utils_dataclass/real_data_gen.py b/dataclass_creator/utils_dataclass/real_data_gen.py
index 939c5a4..da2680e 100644
--- a/dataclass_creator/utils_dataclass/real_data_gen.py
+++ b/dataclass_creator/utils_dataclass/real_data_gen.py
@@ -14,7 +14,7 @@ from typing import Any, Iterable
from opt_tools.jax_tools import SysVarSet, Traj, Batch, tree_map_relaxed
from ros_datahandler import get_bags_iter
from ros_datahandler import JointPosData, ForceData, JointState, WrenchStamped
-from dataclass_creator.franka_dataclass import FrankaData
+from dataclass_creator.custom_dataclasses import FrankaData
def min_max_normalize_data(data: jnp.ndarray):
@@ -107,7 +107,7 @@ def load_dataset(
datatype: SysVarSet,
sync_stream: str,
root: Path,
- glob_filter: str = "",
+ name: str = "SUCCESS",
) -> Iterable[Traj]:
"""collects rosbags in specified dir and return a formatted dataclass containg ragged list*3
@@ -119,11 +119,11 @@ def load_dataset(
FormattedDataclass: data containing the modalities corresponding to the attributes initialized for RosDatahandler
"""
- bags = glob.glob(join(root, glob_filter))
- assert len(bags) > 0, f"I expected bags! No bags on {root} with filter {glob_filter}"
+ bags = glob.glob(join(root, f"*{name}*"))
+ assert len(bags) > 0, f"I expected bags! No bags on {root} with filter {name}"
bags = [Path(bag) for bag in bags]
- logging.info(f'Containing # {len(bags)} rosbags for filter {glob_filter}')
+ logging.info(f'Containing # {len(bags)} rosbags for filter {name}')
data = get_bags_iter(bags, datatype, sync_stream=sync_stream)
data = [d.cast_traj() for d in data]
diff --git a/ros_adetector.py b/ros_adetector.py
index fa8aa05..44b976d 100644
--- a/ros_adetector.py
+++ b/ros_adetector.py
@@ -20,7 +20,7 @@ import numpy as np
from ros_datahandler import RosDatahandler, JointPosData, ForceData
from run.parameter import RepoPaths
from run.load_save_model import load_trained_model_jit, from_dill
-from dataclass_creator.franka_dataclass import FrankaData
+from dataclass_creator.custom_dataclasses import FrankaData
@dataclass
class FrankaDataExp:
diff --git a/run/construct_thresholds.py b/run/construct_thresholds.py
index 9c5951b..cacf251 100644
--- a/run/construct_thresholds.py
+++ b/run/construct_thresholds.py
@@ -17,7 +17,7 @@ def construct_threshold(name):
print(f"Threshold calculation of model {name}")
# load test bags based on bag_name
- dataset = load_dataset(root=RepoPaths.threshold, name="SUCCESS")
+ dataset = load_dataset_test(root=RepoPaths.threshold, name="SUCCESS")
ae_params, model_params, logger_params = load_trained_model(path_to_models=RepoPaths.trained_models_val,
test_name=name)
diff --git a/run/parameter.py b/run/parameter.py
index 674ed0f..2ab2d51 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -9,6 +9,7 @@ import optax
from opt_tools.jax_tools import SysVarSet, Traj
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from run.load_save_model import to_dill, from_dill
@@ -20,10 +21,10 @@ logging.addLevelName(logging.INFO,
class DatasetParams:
"""Parameters to define the data source, batching, etc."""
# Dataset
- datatype: SysVarSet
- sync_stream: str
- root: Path# = Path('experiments/example_data/train')
- glob_filter: str = "*SUCCESS*.bag"
+ datatype: SysVarSet # Instance of dataclass to pass to ros_datahandler
+ sync_stream: str # which attribute in `datatype` to use to synchronize topics
+ root: Path # root for the data directory
+ name: str = "SUCCESS" # used to get bags as glob(f"*{name}*")
# Dataloader
batch_size: int = 50
@@ -36,15 +37,14 @@ class DatasetParams:
f"The sync_stream should be an attribute of datatype, \
got sync_stream {self.sync_stream} and datatype {self.datatype}"
- def get_dataloader(self) -> DynSliceDataloader:
- data = load_dataset(
- self.datatype,
- self.sync_stream,
- self.root,
- self.glob_filter
- )
+ def get_dataset(self) -> DynSliceDataset:
+ return load_dataset(self.datatype, self.sync_stream, self.root, self.name)
+
+ def get_dataloader(self, dataset:DynSliceDataset=None) -> DynSliceDataloader:
+ if dataset is None:
+ dataset = self.get_dataset()
return DynSliceDataloader(
- data,
+ data=dataset,
window_size=self.window_size,
batch_size=self.batch_size,
pred_offset=self.pred_offset
@@ -93,21 +93,22 @@ class TrainerParams:
@dataclass
class ExportParams:
"""Super-class to catch all those params we want to serialize / save."""
- model: Any
model_params: Any
dataset_params: DatasetParams
trainer_params: TrainerParams
+ threshold: float = None # Anomaly detection threshold
+
def __str__(self):
return f"test_w{self.dataset_params.window_size}"\
f"_b{self.dataset_params.batch_size}"\
f"_e{self.trainer_params.epochs}"\
- f"_l{self.model.latent_dim}"
+ f"_l{self.trainer_params.model.latent_dim}"
def save(self, root: Path=Path(".")):
path = os.path.join(root, Path(f"{self}"))
- os.makedirs(path)
+ os.makedirs(path, exist_ok=True)
to_dill(self, path, "model.dill")
diff --git a/run/train_model.py b/run/train_model.py
index 2acbbbe..ffe4bae 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -1,15 +1,16 @@
import os
from itertools import product
from pathlib import Path
+from dataclasses import replace
import jax
+
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
-from dataclass_creator.franka_dataclass import FrankaData
+from dataclass_creator.custom_dataclasses import FrankaData
from model import AutoEncoder, Trainer
-
-from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
from run.parameter import (
OptimizerHParams,
RepoPaths,
@@ -24,60 +25,36 @@ print("Default device:", jax.default_backend())
print("Available devices:", jax.devices())
-def train_model(train_loader: DynSliceDataloader,
- val_loader: DynSliceDataloader,
- logger_params,
- n_epochs: int,
- model
- ):
-
- # --------- initialize autoencoder ---------------------------------------------------------------------------------
-
- ae = AutoEncoder(AEParams.c_hid, logger_params.latent_dim, train_loader.out_size)
+def train_model(dataset_params: DatasetParams,
+ trainer_params: TrainerParams,
+ train_data: DynSliceDataset=None, # optionally re-use the loaded training data
+ ) -> ExportParams:
- # --------- initialize trainer -------------------------------------------------------------------------------------
- config = TrainerParams(
- #purpose=purpose,
- optimizer_hparams=OptimizerHParams(),
- model=ae,
- logger_params=logger_params,
- n_epochs=n_epochs
- )
-
- #with jax.profiler.trace("/tmp/jax_jit", create_perfetto_link=True):
- optimized_params, loss_array = Trainer(config, train_loader).train_jit()
- #optimized_params, loss_array = Trainer(config, train_loader, prep_batches=True).train()
+ # Build Dataloader. If train_data is none, will load dataset
+ train_loader = dataset_params.get_dataloader(train_data)
+
+ #with jax.profiler.trace("/tmp/jax_jit", create_perfetto_link=True)
+ optimized_params, loss_array = Trainer(trainer_params, train_loader).train_jit()
random_threshold = 0.001
- return ae, optimized_params, random_threshold
-
-
-def train(path_to_data: Path, window_size: int, batch_size: int, epochs: int, latent_dim: int, purpose: str, train_data=None):
- if not train_data:
- train_data = load_dataset(root=path_to_data, name="SUCCESS")
- train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
-
-# logger_params = LoggerParams(window_size=window_size,
-# batch_size=batch_size,
-# epochs=epochs,
-# latent_dim=latent_dim,
-# path_to_models=RepoPaths.trained_models_val)
-
- logger_params = None
- ae, optimized_params, threshold = train_model(train_loader,
- val_loader,
- logger_params,
- n_epochs=epochs,
- purpose=purpose)
+ export_params = ExportParams(
+ model_params = optimized_params,
+ threshold = random_threshold,
+ dataset_params=dataset_params,
+ trainer_params=trainer_params
+ )
- to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
- assert isinstance (ae, AutoEncoder)
- assert isinstance(optimized_params, dict)
+ export_params.save(root=dataset_params.root)
+ return export_params
-def train_loop(path_to_data: Path, window_size: list, batch_size: list, epochs: list, latent_dim:list, purpose: str):
+def train_loop(dataset_params: DatasetParams,
+ trainer_params: TrainerParams,
+ window_size: list,
+ batch_size: list,
+ epochs: list,
+ latent_dim:list):
"""
Train multiple models by setting different window, batch sizes and number of epochs.
@@ -88,11 +65,17 @@ def train_loop(path_to_data: Path, window_size: list, batch_size: list, epochs:
epochs: List of numbers of epochs the model will be trained for.
"""
model_params = product(window_size, batch_size, epochs, latent_dim)
- train_data = load_dataset(root=path_to_data, name="SUCCESS")
+ train_data = dataset_params.get_dataset()
for hparams in model_params:
window_size, batch_size, epochs, latent_dim = hparams
- train(path_to_data=path_to_data, window_size=window_size, batch_size=batch_size, epochs=epochs, latent_dim=latent_dim, purpose=purpose, train_data=train_data)
-
+ model = replace(trainer_params.model, latent_dim=latent_dim)
+ replace(dataset_params,
+ window_size=window_size,
+ batch_size=batch_size,
+ model=model)
+ replace(trainer_params, epochs=epochs)
+ train_model(dataset_params, trainer_params, train_data=train_data)
+
def lazytrainer():
"""
This refactor tries to:
@@ -128,5 +111,4 @@ def lazytrainer():
if __name__ == "__main__":
lazytrainer()
-
os._exit(1)
diff --git a/tests/test_data_load.py b/tests/test_data_load.py
index 160b509..d168014 100644
--- a/tests/test_data_load.py
+++ b/tests/test_data_load.py
@@ -2,23 +2,20 @@ import logging
import os
import pickle
import dill
-import jax
-
from pathlib import Path
-from dataclass_creator.franka_dataclass import FrankaData
-from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
+import jax
+
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from opt_tools.jax_tools import Batch, to_batch, tree_map_relaxed
+from tests.test_defaults import get_dataset_test
+
logger = logging.getLogger()
#curr_dir = Path(os.path.dirname(os.path.abspath(__file__)))
#root = os.fspath(Path(curr_dir.parent, 'experiments/example_data/train').resolve())
-root = Path('experiments/example_data/train')
-
-datatype =
window_size = 20
batch_size = 10
@@ -52,15 +49,7 @@ def test_save_load():
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
logger.info('Loading data from bags')
- raw_data = load_dataset(
- datatype=FrankaData(
- jt=JointPosData('/joint_states', JointState),
- force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
- ),
- sync_stream = 'force',
- root,
- "SUCCESS")[0]
-
+ raw_data = get_dataset_test()[0]
# Try with some transformed data
treemapped_data = tree_map_relaxed(lambda l: l, raw_data).cast_batch()
diff --git a/tests/test_defaults.py b/tests/test_defaults.py
new file mode 100644
index 0000000..e04a49f
--- /dev/null
+++ b/tests/test_defaults.py
@@ -0,0 +1,53 @@
+"""Helper functions to reduce boilerplate for tests."""
+from pathlib import Path
+
+from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
+from dataclass_creator.custom_dataclasses import FrankaData
+from run.parameter import RepoPaths, DatasetParams, TrainerParams
+from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
+from model.autoencoder import AutoEncoder
+from utils.loss_calculation import ReconstructionLoss
+from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
+
+def get_dataset_params(pred_offset=0):
+ return DatasetParams(
+ datatype=FrankaData(
+ jt=JointPosData('/joint_states', JointState),
+ force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
+ ),
+ sync_stream = 'force',
+ pred_offset=pred_offset,
+ window_size=100,
+ root = Path('experiments/example_data/train'),
+ )
+
+def get_trainer_params():
+ dataloader = get_dataset_params().get_dataloader()
+
+ model = AutoEncoder(c_hid=10, latent_dim=15, c_out=dataloader.out_size)
+ return TrainerParams(
+ model=model,
+ loss=ReconstructionLoss(model).batch_loss,
+ epochs=1
+ )
+
+def get_dataset_test():
+ datatype = FrankaData(jt = JointPosData('/joint_states', JointState),
+ force = ForceData('/franka_state_controller/F_ext', WrenchStamped))
+
+ return load_dataset(
+ datatype=datatype,
+ sync_stream='force',
+ root=RepoPaths.example_data_train,
+ name="SUCCESS"
+ )
+
+
+def get_dataloader_test(pred_offset=0):
+ data = get_dataset_test()
+ return DynSliceDataloader(
+ data,
+ window_size=30,
+ batch_size=10,
+ pred_offset=pred_offset
+ )
diff --git a/tests/test_flattening.py b/tests/test_flattening.py
index 8ec1c39..4f641a2 100644
--- a/tests/test_flattening.py
+++ b/tests/test_flattening.py
@@ -8,7 +8,7 @@ from jax.flatten_util import ravel_pytree
from run.parameter import OptimizerHParams
from dataclass_creator import create_synth_nominal_data
-from dataclass_creator.synth_dataclass import SynthData
+from dataclass_creator.custom_dataclasses import SynthData
from dataclass_creator.utils_dataclass.synth_data_gen import DataGeneration
from model import AutoEncoder, Trainer
from opt_tools.jax_tools import SysVarSet
diff --git a/tests/test_model_load.py b/tests/test_model_load.py
index 138743e..e6f86b6 100644
--- a/tests/test_model_load.py
+++ b/tests/test_model_load.py
@@ -3,75 +3,42 @@ import time
import jax
-from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
+from tests.test_defaults import get_dataloader_test, get_trainer_params, get_dataset_params
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from run.parameter import RepoPaths
from run.train_model import train_model
from run.load_save_model import load_trained_model_jit, load_trained_model, to_dill
from tests.test_data_load import round_trip_dill, round_trip_pickle
-
-
-def test_model_load():
-
- window_size = 10
- batch_size = 10
- n_epochs = 1
-
- train_data = load_dataset(root=RepoPaths.example_data_train, name="SUCCESS")
- train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- logger_params = LoggerParams(window_size=window_size,
- batch_size=batch_size,
- epochs=n_epochs,
- path_to_models=RepoPaths.trained_models_val)
-
- ae, optimized_params, threshold = train_model(train_loader,
- train_loader,
- logger_params,
- n_epochs=n_epochs)
-
-
- (ae_pkl, optimized_params_pkl) = round_trip_pickle((ae, optimized_params))
- (ae_dill, optimized_params_dill) = round_trip_dill((ae, optimized_params))
+def test_model_load():
+ dataset_params = get_dataset_params()
+ trainer_params = get_trainer_params()
+ res = train_model(dataset_params, trainer_params)
- val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
+ val_loader = get_dataloader_test()
reference_batch = next(iter(val_loader))
- reconstruction_batch = jax.vmap(ae_pkl.apply, in_axes=(None, 0))(optimized_params_pkl, reference_batch)
- reconstruction_batch = jax.vmap(ae_dill.apply, in_axes=(None, 0))(optimized_params_dill, reference_batch)
+ #reconstruction_batch = jax.vmap(ae_pkl.apply, in_axes=(None, 0))(optimized_params_pkl, reference_batch)
+ #reconstruction_batch = jax.vmap(ae_dill.apply, in_axes=(None, 0))(optimized_params_dill, reference_batch)
-def test_model_load_jit():
-
- window_size = 10
- batch_size = 10
- n_epochs = 1
-
- train_data = load_dataset(root=RepoPaths.example_data_train, name="SUCCESS")
- train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- logger_params = LoggerParams(window_size=window_size,
- batch_size=batch_size,
- epochs=n_epochs,
- path_to_models=RepoPaths.trained_models_val)
-
- ae, optimized_params, threshold = train_model(train_loader,
- train_loader,
- logger_params,
- n_epochs=n_epochs)
-
+def test_model_load_jit():
+ dataset_params = get_dataset_params()
+ trainer_params = get_trainer_params()
+ res = train_model(dataset_params, trainer_params)
#(ae_pkl, optimized_params_pkl) = round_trip_pickle((ae, optimized_params))
#(ae_dill, optimized_params_dill) = round_trip_dill((ae, optimized_params))
- to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
- ae_jit, _ = load_trained_model_jit(logger_params.path_to_model, "")
- ae, params, _ = load_trained_model(logger_params.path_to_model, "")
+ #to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
+ #ae_jit, _ = load_trained_model_jit(logger_params.path_to_model, "")
+ #ae, params, _ = load_trained_model(logger_params.path_to_model, "")
- val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- reference_batch = next(iter(val_loader))
- print(reference_batch.shape)
-
+ #val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
+ #reference_batch = next(iter(val_loader))
+ #print(reference_batch.shape)
+ """
times = []
for sample in reference_batch:
tic = time.time()
@@ -87,7 +54,7 @@ def test_model_load_jit():
print(f"NOJIT Took {times}, \n final result {res_jit.force}")
print(res_jit.force - res_nojit.force)
-
+ """
if __name__ == "__main__":
#test_model_load()
diff --git a/tests/test_synth_data.py b/tests/test_synth_data.py
index 826793d..6f0f4c0 100644
--- a/tests/test_synth_data.py
+++ b/tests/test_synth_data.py
@@ -4,6 +4,7 @@ import os
from pathlib import Path
+import pytest
import jax
from dataclass_creator.create_dataloader import generate_synth_data
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader, split_dataloader
@@ -12,20 +13,20 @@ from run.load_save_model import to_dill
from run.train_model import train_model
from run.valmodel import ValModel
from model import AutoEncoder
-from run.parameter import RepoPaths
+from run.parameter import RepoPaths, TrainerParams
# Types
-from dataclass_creator.synth_dataclass import SynthData
+from dataclass_creator.custom_dataclasses import SynthData
from opt_tools.jax_tools import Traj
curr_dir = Path(os.path.dirname(os.path.abspath(__file__)))
root_pretrained_models = os.fspath(Path(curr_dir.parent, 'trained_models/synth_data').resolve())
-
window_size = 50
batch_size = 10
n_epochs = 15
+@pytest.mark.skip(reason="Fixed in branch 48")
def test_datagen():
train = generate_synth_data(amp_rest_pos = 0.3)
assert isinstance(train[0], SynthData)
@@ -34,7 +35,7 @@ def test_datagen():
train, val, test = split_dataloader(train, window_size=50)
assert isinstance(train, DynSliceDataloader)
-
+@pytest.mark.skip(reason="Fixed in branch 48")
def test_train():
from run.parameter import RepoPaths
data = generate_synth_data(amp_rest_pos = 0.8)
@@ -42,10 +43,7 @@ def test_train():
anomaly_data = generate_synth_data(amp_rest_pos = 0.6)
anomaly = DynSliceDataloader(anomaly_data)
-
- logger_params = LoggerParams(window_size, batch_size, n_epochs,
- path_to_models=RepoPaths.trained_models_val)
-
+
ae, optimized_params, threshold = train_model(train,
val,
logger_params,
@@ -53,6 +51,7 @@ def test_train():
to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
nominal_eval(data[-1])
+@pytest.mark.skip(reason="Fixed in branch 48")
def nominal_eval(test_traj: Traj):
val_synth = ValModel(path_to_models=RepoPaths.trained_models_val,
name_model=f"test_w{window_size}_b{batch_size}_e{n_epochs}",
diff --git a/tests/test_training.py b/tests/test_training.py
index 91170f0..ef87268 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -1,51 +1,37 @@
import os
-import jax
+from pathlib import Path
+from dataclasses import replace
+import jax
from jax.flatten_util import ravel_pytree
-from pathlib import Path
+
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
-from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
+from tests.test_defaults import get_dataloader_test
from model import AutoEncoder, Trainer
from run.parameter import OptimizerHParams, RepoPaths
from run.train_model import train_model
from run.load_save_model import to_dill
-
-
-def test_recon(window_size: int=30, batch_size: int=10, epochs: int=1):
- train_data = load_dataset(root=RepoPaths.example_data_train, name="SUCCESS")
- train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
-
- logger_params = LoggerParams(window_size=window_size, batch_size=batch_size, epochs=epochs, path_to_models=RepoPaths.trained_models_val)
-
- ae, optimized_params, threshold = train_model(train_loader,
- train_loader,
- logger_params,
- n_epochs=epochs)
-
- to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
-
- assert isinstance (ae, AutoEncoder)
- assert isinstance(optimized_params, dict)
-
-def test_pred(window_size: int=30, batch_size: int=10, epochs: int=1):
- train_data = load_dataset(root=RepoPaths.example_data_train, name="SUCCESS")
- train_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size, pred_offset=10)
-
- logger_params = LoggerParams(window_size=window_size,
- batch_size=batch_size,
- epochs=epochs,
- path_to_models=RepoPaths.trained_models_val)
-
- ae, optimized_params, threshold = train_model(train_loader,
- train_loader,
- logger_params,
- n_epochs=epochs,
- purpose='Prediction')
-
- assert isinstance (ae, AutoEncoder)
- assert isinstance(optimized_params, dict)
+from tests.test_defaults import get_trainer_params, get_dataset_params
+from utils.loss_calculation import PredictionLoss
+
+def test_recon():
+ res = train_model(get_dataset_params(),
+ get_trainer_params())
+
+ assert isinstance (res.trainer_params.model, AutoEncoder)
+ assert isinstance(res.model_params, dict)
+
+def test_pred():
+ trainer_params = get_trainer_params()
+ trainer_params = replace(trainer_params,
+ loss=PredictionLoss(trainer_params.model).batch_loss)
+ res = train_model(get_dataset_params(pred_offset=10),
+ trainer_params)
+
+ assert isinstance (res.trainer_params.model, AutoEncoder)
+ assert isinstance(res.model_params, dict)
if __name__ == "__main__":
- test_pred(epochs=5)
- test_recon(epochs=5)
+ test_pred()
+ test_recon()
diff --git a/utils/validation_fn.py b/utils/validation_fn.py
index ef7fa81..49e718c 100644
--- a/utils/validation_fn.py
+++ b/utils/validation_fn.py
@@ -3,8 +3,7 @@ import jax
from sklearn.metrics import auc, confusion_matrix
-from dataclass_creator.synth_dataclass import SynthData
-from dataclass_creator.franka_dataclass import FrankaData
+from dataclass_creator.custom_dataclasses import SynthData, FrankaData
# Types
from typing import List, Union, Tuple
--
GitLab
From c4798592087cc5a6953d3d848da57d72d4acb480 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Wed, 18 Dec 2024 03:40:33 +0100
Subject: [PATCH 08/12] added model loader
---
.gitignore | 3 ++-
run/construct_thresholds.py | 2 +-
run/load_save_model.py | 4 ++--
run/parameter.py | 37 ++++++++++++++++++++++++--------
tests/test_defaults.py | 23 +++++---------------
tests/test_model_load.py | 42 ++++++++++++++++++-------------------
6 files changed, 58 insertions(+), 53 deletions(-)
diff --git a/.gitignore b/.gitignore
index 694c39b..473368a 100755
--- a/.gitignore
+++ b/.gitignore
@@ -12,4 +12,5 @@ archiv
experiments/01_faston_converging/trained_models_val
experiments/synth_data/
*/log_file
-*/model.dill
\ No newline at end of file
+*/model.dill
+**/result_*
\ No newline at end of file
diff --git a/run/construct_thresholds.py b/run/construct_thresholds.py
index cacf251..9c5951b 100644
--- a/run/construct_thresholds.py
+++ b/run/construct_thresholds.py
@@ -17,7 +17,7 @@ def construct_threshold(name):
print(f"Threshold calculation of model {name}")
# load test bags based on bag_name
- dataset = load_dataset_test(root=RepoPaths.threshold, name="SUCCESS")
+ dataset = load_dataset(root=RepoPaths.threshold, name="SUCCESS")
ae_params, model_params, logger_params = load_trained_model(path_to_models=RepoPaths.trained_models_val,
test_name=name)
diff --git a/run/load_save_model.py b/run/load_save_model.py
index da60d0b..b487e3e 100644
--- a/run/load_save_model.py
+++ b/run/load_save_model.py
@@ -49,8 +49,8 @@ def from_dill(path: str, file_name: str):
try:
with open(name, 'rb') as fopen:
loaded_data = dill.load(fopen)
- except:
- print("Undilling is dead!")
+ except Exception as e:
+ logging.error(f"Undilling is dead! {e}")
return loaded_data
def load_trained_model(path_to_models: Path, test_name: str):
diff --git a/run/parameter.py b/run/parameter.py
index 2ab2d51..1975f73 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -1,10 +1,11 @@
import os
from pathlib import Path
import logging
-from typing import Iterable, Any
+from typing import Iterable, Any, Self, Callable
from dataclasses import dataclass, field
from datetime import datetime
+import jax
import optax
from opt_tools.jax_tools import SysVarSet, Traj
@@ -37,6 +38,9 @@ class DatasetParams:
f"The sync_stream should be an attribute of datatype, \
got sync_stream {self.sync_stream} and datatype {self.datatype}"
+ def __str__(self):
+ return f"w{self.window_size}_b{self.batch_size}"
+
def get_dataset(self) -> DynSliceDataset:
return load_dataset(self.datatype, self.sync_stream, self.root, self.name)
@@ -90,28 +94,43 @@ class TrainerParams:
epochs: int = 50
check_val_every_n_epoch: int = 1
+ def __str__(self):
+ return f"e{self.epochs}_l{self.model.latent_dim}"
+
+
@dataclass
class ExportParams:
"""Super-class to catch all those params we want to serialize / save."""
- model_params: Any
dataset_params: DatasetParams
trainer_params: TrainerParams
+ model_params: Any = None # trained model parameters
+ # @TODO maybe dont need b/c it's exported separately in construct_thresholds?
threshold: float = None # Anomaly detection threshold
-
def __str__(self):
- return f"test_w{self.dataset_params.window_size}"\
- f"_b{self.dataset_params.batch_size}"\
- f"_e{self.trainer_params.epochs}"\
- f"_l{self.trainer_params.model.latent_dim}"
+ return f"result_{self.dataset_params}_{self.trainer_params}"
-
- def save(self, root: Path=Path(".")):
+ def save(self, root: Path = None):
+ if root == None:
+ logging.info("No path provided, using dataset.root")
+ root = Path(self.dataset_params.root)
path = os.path.join(root, Path(f"{self}"))
os.makedirs(path, exist_ok=True)
to_dill(self, path, "model.dill")
+
+ @classmethod
+ def load(cls, dataset_params, trainer_params, root: Path=Path(".")) -> Self:
+ export_params = ExportParams(dataset_params, trainer_params)
+ return from_dill(os.path.join(root, Path(f"{export_params}")), "model.dill")
+
+ def get_apply_jit(self, root: Path=Path(".")) -> Callable[[SysVarSet], SysVarSet]:
+ return jax.jit(lambda input: self.trainer_params.model.apply(self.model_params, input))
+
+ def get_encode_jit(self, root: Path=Path(".")) -> Callable[[SysVarSet], SysVarSet]:
+ return jax.jit(lambda input: self.trainer_params.model.encode(self.model_params, input))
+
@dataclass
class RepoPaths:
"""Path names for 01_faston_converging anomaly detection."""
diff --git a/tests/test_defaults.py b/tests/test_defaults.py
index e04a49f..733e202 100644
--- a/tests/test_defaults.py
+++ b/tests/test_defaults.py
@@ -13,7 +13,7 @@ def get_dataset_params(pred_offset=0):
return DatasetParams(
datatype=FrankaData(
jt=JointPosData('/joint_states', JointState),
- force=ForceData('/franka_state_controller/F_ext', WrenchStamped)#, is_sync=True)
+ force=ForceData('/franka_state_controller/F_ext', WrenchStamped)
),
sync_stream = 'force',
pred_offset=pred_offset,
@@ -32,22 +32,9 @@ def get_trainer_params():
)
def get_dataset_test():
- datatype = FrankaData(jt = JointPosData('/joint_states', JointState),
- force = ForceData('/franka_state_controller/F_ext', WrenchStamped))
-
- return load_dataset(
- datatype=datatype,
- sync_stream='force',
- root=RepoPaths.example_data_train,
- name="SUCCESS"
- )
-
+ dp = get_dataset_params()
+ return dp.get_dataset()
def get_dataloader_test(pred_offset=0):
- data = get_dataset_test()
- return DynSliceDataloader(
- data,
- window_size=30,
- batch_size=10,
- pred_offset=pred_offset
- )
+ dp = get_dataset_params()
+ return dp.get_dataloader()
diff --git a/tests/test_model_load.py b/tests/test_model_load.py
index e6f86b6..b17291e 100644
--- a/tests/test_model_load.py
+++ b/tests/test_model_load.py
@@ -1,11 +1,12 @@
from pathlib import Path
import time
+import os
import jax
from tests.test_defaults import get_dataloader_test, get_trainer_params, get_dataset_params
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
-from run.parameter import RepoPaths
+from run.parameter import RepoPaths, ExportParams
from run.train_model import train_model
from run.load_save_model import load_trained_model_jit, load_trained_model, to_dill
from tests.test_data_load import round_trip_dill, round_trip_pickle
@@ -14,48 +15,45 @@ def test_model_load():
dataset_params = get_dataset_params()
trainer_params = get_trainer_params()
res = train_model(dataset_params, trainer_params)
-
+
val_loader = get_dataloader_test()
reference_batch = next(iter(val_loader))
-
- #reconstruction_batch = jax.vmap(ae_pkl.apply, in_axes=(None, 0))(optimized_params_pkl, reference_batch)
- #reconstruction_batch = jax.vmap(ae_dill.apply, in_axes=(None, 0))(optimized_params_dill, reference_batch)
-
-
+ reference_sample = jax.tree_util.tree_map(lambda l: l[0], reference_batch)
+
+ res_load = ExportParams.load(dataset_params, trainer_params, root=dataset_params.root)
+
+ apply = res_load.get_apply_jit()
+ recon_sample = res_load.get_apply_jit()(reference_sample)
+ #encode_sample = res_load.get_encode_jit()(reference_sample)
+
def test_model_load_jit():
dataset_params = get_dataset_params()
trainer_params = get_trainer_params()
res = train_model(dataset_params, trainer_params)
- #(ae_pkl, optimized_params_pkl) = round_trip_pickle((ae, optimized_params))
- #(ae_dill, optimized_params_dill) = round_trip_dill((ae, optimized_params))
- #to_dill((ae, optimized_params, logger_params), logger_params.path_to_model, "model.dill")
- #ae_jit, _ = load_trained_model_jit(logger_params.path_to_model, "")
- #ae, params, _ = load_trained_model(logger_params.path_to_model, "")
+ val_loader = get_dataloader_test()
+ reference_batch = next(iter(val_loader))
+
+ res_load = ExportParams.load(dataset_params, trainer_params, root=dataset_params.root)
+
+ apply = res_load.get_apply_jit()
- #val_loader = DynSliceDataloader(train_data, window_size=window_size, batch_size=batch_size)
- #reference_batch = next(iter(val_loader))
- #print(reference_batch.shape)
-
- """
times = []
for sample in reference_batch:
tic = time.time()
- res_jit = ae_jit(sample)
+ res_jit = apply(sample)
times.append(time.time()-tic)
print(f"JIT Took {times}, \n final result {res_jit.force}")
times = []
for sample in reference_batch:
tic = time.time()
- res_nojit = ae.apply(params, sample)
+ res_nojit = res.trainer_params.model.apply(res.model_params, sample)
times.append(time.time() - tic)
print(f"NOJIT Took {times}, \n final result {res_jit.force}")
- print(res_jit.force - res_nojit.force)
- """
if __name__ == "__main__":
- #test_model_load()
+ test_model_load()
test_model_load_jit()
--
GitLab
From b656864dde16dcc26b2e7dc3276eabaa6af6f0f4 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Wed, 18 Dec 2024 03:53:03 +0100
Subject: [PATCH 09/12] loading added, tested
---
run/construct_thresholds.py | 19 ++++++++-----------
run/parameter.py | 15 +++++++++++----
tests/test_model_load.py | 1 +
3 files changed, 20 insertions(+), 15 deletions(-)
diff --git a/run/construct_thresholds.py b/run/construct_thresholds.py
index 9c5951b..c7df0f2 100644
--- a/run/construct_thresholds.py
+++ b/run/construct_thresholds.py
@@ -6,6 +6,7 @@ import numpy as np
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from run.load_save_model import to_dill, from_dill, load_trained_model
+from parameter import ExportParams
#for WindowsPath
# Types
@@ -19,27 +20,23 @@ def construct_threshold(name):
# load test bags based on bag_name
dataset = load_dataset(root=RepoPaths.threshold, name="SUCCESS")
- ae_params, model_params, logger_params = load_trained_model(path_to_models=RepoPaths.trained_models_val,
- test_name=name)
-
- # we can load all test SUCCESS/FAILURE bags but validate only one at a time
- test_loader = DynSliceDataloader(data=dataset,
- window_size=logger_params.window_size,
- batch_size=logger_params.batch_size)
-
+ export_params = ExportParams.load_from_full_path(RepoPaths.trained_models_val+name)
+ test_loader = export_params.dataset_params.get_dataloader(dataset)
+ window_size = export_params.dataset_params.window_size
+
start = 0
end = len(test_loader.dset)
- all_timesteps = [i * logger_params.window_size for i in range(start // logger_params.window_size, end // logger_params.window_size + 1)]
+ all_timesteps = [i * window_size for i in range(start // window_size, end // window_size + 1)]
windows = test_loader.dset.get_batch(all_timesteps)
print(f"num of windows {windows.batch_size}")
- recon_windows = jax.vmap(ae_params.apply, in_axes=(None, 0))(model_params, windows)
+ recon_windows = export_params.get_apply_vmap_jit()(windows)
loss = tree_map(lambda w, rw: jnp.mean(jnp.abs(w - rw), axis=1),windows, recon_windows)
threshold = tree_map(lambda l: jnp.max(l, axis=0), loss)
threshold = tree_map(lambda l: l + 0.6*l, threshold)
thresholds = [tree_map(lambda l: l + i*l, threshold) for i in np.arange(-0.9, 0.95, 0.05).tolist()]
- to_dill(threshold, logger_params.path_to_model, "threshold.dill")
+ to_dill(threshold, str(export_params), "threshold.dill")
if __name__ == "__main__":
diff --git a/run/parameter.py b/run/parameter.py
index 1975f73..13c9148 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -8,7 +8,7 @@ from datetime import datetime
import jax
import optax
-from opt_tools.jax_tools import SysVarSet, Traj
+from opt_tools.jax_tools import SysVarSet, Batch
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
@@ -109,12 +109,12 @@ class ExportParams:
def __str__(self):
return f"result_{self.dataset_params}_{self.trainer_params}"
-
+
def save(self, root: Path = None):
if root == None:
logging.info("No path provided, using dataset.root")
root = Path(self.dataset_params.root)
- path = os.path.join(root, Path(f"{self}"))
+ path = os.path.join(root, Path(str(self)))
os.makedirs(path, exist_ok=True)
to_dill(self, path, "model.dill")
@@ -123,9 +123,16 @@ class ExportParams:
export_params = ExportParams(dataset_params, trainer_params)
return from_dill(os.path.join(root, Path(f"{export_params}")), "model.dill")
- def get_apply_jit(self, root: Path=Path(".")) -> Callable[[SysVarSet], SysVarSet]:
+ @classmethod
+ def load_from_full_path(cls, root: Path) -> Self:
+ return from_dill(root, "model.dill")
+
+ def get_apply_jit(self) -> Callable[[SysVarSet], SysVarSet]:
return jax.jit(lambda input: self.trainer_params.model.apply(self.model_params, input))
+ def get_apply_vmap_jit(self) -> Callable[[Batch], Batch]:
+ apply_vmap = jax.vmap(self.trainer_params.model.apply, in_axes=(None, 0))
+ return jax.jit(lambda input: apply_vmap(self.model_params, input))
def get_encode_jit(self, root: Path=Path(".")) -> Callable[[SysVarSet], SysVarSet]:
return jax.jit(lambda input: self.trainer_params.model.encode(self.model_params, input))
diff --git a/tests/test_model_load.py b/tests/test_model_load.py
index b17291e..6b5fa5f 100644
--- a/tests/test_model_load.py
+++ b/tests/test_model_load.py
@@ -16,6 +16,7 @@ def test_model_load():
trainer_params = get_trainer_params()
res = train_model(dataset_params, trainer_params)
+
val_loader = get_dataloader_test()
reference_batch = next(iter(val_loader))
reference_sample = jax.tree_util.tree_map(lambda l: l[0], reference_batch)
--
GitLab
From a029b5aa61d8e557f815336fe10cbc28d0cf62c7 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Wed, 18 Dec 2024 15:38:09 +0100
Subject: [PATCH 10/12] add purpose
---
dataclass_creator/dyn_slice_dataloader.py | 8 ++++
main.py | 50 -----------------------
run/parameter.py | 40 ++++++++++++++++--
run/train_model.py | 23 +++++++----
4 files changed, 60 insertions(+), 61 deletions(-)
delete mode 100755 main.py
diff --git a/dataclass_creator/dyn_slice_dataloader.py b/dataclass_creator/dyn_slice_dataloader.py
index c38f1fe..eb51bf4 100644
--- a/dataclass_creator/dyn_slice_dataloader.py
+++ b/dataclass_creator/dyn_slice_dataloader.py
@@ -102,3 +102,11 @@ class DynSliceDataloader:
def out_size(self) -> int:
"""Return length of output vector."""
return ravel_pytree(self.get_init_vec())[0].shape[0]
+
+
+ @property
+ def purpose(self) -> str:
+ if type(self.dset) == DynSliceDataset:
+ return "reconstruction"
+ elif type(self.dset) == PredDynSliceDataset:
+ return "prediction"
diff --git a/main.py b/main.py
deleted file mode 100755
index 1f5d086..0000000
--- a/main.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# Sketch of main fucntionality
-
-from TargetSys import * # get the typedefs for State, Input, Param and Output
-from TargetSys import SysName as DynSys
-# I think it might be acceptable to switch the types and
-# datastream_config by the include... it's good to define
-# those both in one place, as they are pretty closely
-# related.
-
-
-class Encoder(DynSys):
- def detect(state: State, measured_output: Output):
- expected_output = self.output(state)
- return jnp.norm(expected_output - measured_output)
-
-
-class AnomalyDetector:
- def __init__(self):
- self._encoder = Encoder("checkpoint.pkl")
- self.ros_init()
-
- def ros_init(self):
- self._err_pub = None # Topic to publish result of anomaly detection
- self._listener = RosListener(datastream_config)
- self._listener.block_until_all_topics_populated()
-
- def eval(self):
- err = self._encoder.detect(
- self._listener.get_input(),
- self._listener.get_state(),
- self._listener.get_output(),
- )
- self._err_pub.publish(err)
-
-
-class Planner(AnomalyDetector):
- # Inherit so we avoid re-implementation of encoder-related things
- def __init__(self):
- super().__init__()
-
- self._planner = MPC(sys=self._encoder, config="mpc_params.yaml")
- self._cmd_pub = None # Topic to publish commands
-
- def plan(self):
- cmd = self._planner.plan(self._listener.get_current_state())
- self._cmd_pub.publish(cmd)
-
-
-if __name__ == "__main__":
- anomaly_detector = AnomalyDetector()
diff --git a/run/parameter.py b/run/parameter.py
index 13c9148..fc9363a 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -1,4 +1,5 @@
import os
+from enum import Enum
from pathlib import Path
import logging
from typing import Iterable, Any, Self, Callable
@@ -13,11 +14,17 @@ from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from run.load_save_model import to_dill, from_dill
+from utils.loss_calculation import PredictionLoss, ReconstructionLoss
logging.addLevelName(logging.INFO,
"{}{}{}".format('\033[92m',
logging.getLevelName(logging.INFO),
'\033[0m'))
+class Purpose(Enum):
+ """Purpose of a dataset / model / loss fn."""
+ RECONSTRUCTION = "reconstruction"
+ PREDICTION = "prediction"
+
@dataclass
class DatasetParams:
"""Parameters to define the data source, batching, etc."""
@@ -29,7 +36,6 @@ class DatasetParams:
# Dataloader
batch_size: int = 50
- # Dataloader, also shared with model
window_size: int = None
pred_offset: int = 0
@@ -53,7 +59,16 @@ class DatasetParams:
batch_size=self.batch_size,
pred_offset=self.pred_offset
)
-
+
+ @property
+ def purpose(self) -> str:
+ if self.pred_offset == 0:
+ return Purpose.RECONSTRUCTION
+ elif self.pred_offset > 0:
+ return Purpose.PREDICTION
+ else:
+ raise TypeError(f"Pred_offset is {self.pred_offset}<0 in DatasetParams, this is not supported.")
+
@dataclass
class OptimizerHParams:
optimizer: optax = optax.adam
@@ -82,7 +97,7 @@ class OptimizerHParams:
)
return optimizer_class(lr)
-@dataclass
+@dataclass(kw_only=True)
class TrainerParams:
"""Define how the traiing process should happen."""
model: Any #AutoEncoder
@@ -97,6 +112,16 @@ class TrainerParams:
def __str__(self):
return f"e{self.epochs}_l{self.model.latent_dim}"
+ @classmethod
+ def from_purpose(cls, purpose: Purpose, model: Any, **kwargs):
+ """Helper to automagically derive loss function from purpose."""
+ if purpose == Purpose.RECONSTRUCTION:
+ loss = ReconstructionLoss(model).batch_loss
+ elif purpose == Purpose.PREDICTION:
+ loss = PredictionLoss(model).batch_loss
+ else:
+ raise TypeError(f"Unknown purpose {purpose}, can't guess your loss fn.")
+ return cls(model=model, loss=loss, **kwargs)
@dataclass
class ExportParams:
@@ -111,6 +136,10 @@ class ExportParams:
return f"result_{self.dataset_params}_{self.trainer_params}"
def save(self, root: Path = None):
+ """Save model with path root/str(self)/model.dill.
+ IN:
+ root (Path): root to save. If none, uses dataset_params.root
+ """
if root == None:
logging.info("No path provided, using dataset.root")
root = Path(self.dataset_params.root)
@@ -121,20 +150,23 @@ class ExportParams:
@classmethod
def load(cls, dataset_params, trainer_params, root: Path=Path(".")) -> Self:
export_params = ExportParams(dataset_params, trainer_params)
- return from_dill(os.path.join(root, Path(f"{export_params}")), "model.dill")
+ return from_dill(os.path.join(root, Path(str(export_params))), "model.dill")
@classmethod
def load_from_full_path(cls, root: Path) -> Self:
return from_dill(root, "model.dill")
def get_apply_jit(self) -> Callable[[SysVarSet], SysVarSet]:
+ """Returns fn which evals model on single sample input."""
return jax.jit(lambda input: self.trainer_params.model.apply(self.model_params, input))
def get_apply_vmap_jit(self) -> Callable[[Batch], Batch]:
+ """Returns fn which evals model on batch input."""
apply_vmap = jax.vmap(self.trainer_params.model.apply, in_axes=(None, 0))
return jax.jit(lambda input: apply_vmap(self.model_params, input))
def get_encode_jit(self, root: Path=Path(".")) -> Callable[[SysVarSet], SysVarSet]:
+ """Returns fn which evals model encoder on single input."""
return jax.jit(lambda input: self.trainer_params.model.encode(self.model_params, input))
diff --git a/run/train_model.py b/run/train_model.py
index ffe4bae..e8a6b5f 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -5,7 +5,7 @@ from dataclasses import replace
import jax
-
+from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
@@ -44,7 +44,7 @@ def train_model(dataset_params: DatasetParams,
dataset_params=dataset_params,
trainer_params=trainer_params
)
-
+
export_params.save(root=dataset_params.root)
return export_params
@@ -96,19 +96,28 @@ def lazytrainer():
train_data = dataset_params.get_dataloader()
- model = AutoEncoder(c_hid=30, latent_dim=15, c_out=train_data.out_size )
+ model = AutoEncoder(c_hid=30, latent_dim=15, c_out=train_data.out_size)
- trainer_params = TrainerParams(
+ #trainer_params = TrainerParams(
+ # model=model,
+ # loss=ReconstructionLoss(model).batch_loss,
+ # epochs=2
+ #)
+
+ # Even lazier way!
+ trainer_params = TrainerParams.from_purpose(
+ purpose=dataset_params.purpose,
model=model,
- loss=ReconstructionLoss(model).batch_loss,
epochs=2
)
optimized_params, loss_array = Trainer(trainer_params, train_data).train_jit()
-
- ExportParams(model, optimized_params, dataset_params, trainer_params).save()
+
+ ExportParams(dataset_params, trainer_params, optimized_params).save()
if __name__ == "__main__":
lazytrainer()
+
+
os._exit(1)
--
GitLab
From 84ce54d0a1f33c78a97c7262c7934729cf7e8ba7 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Thu, 19 Dec 2024 15:02:32 +0100
Subject: [PATCH 11/12] cleanup, fix circular import, pytests passing
---
dataclass_creator/dyn_slice_dataloader.py | 7 -------
model/__init__.py | 2 --
ros_adetector.py | 2 +-
run/parameter.py | 25 ++++++++++++++++-------
run/train_model.py | 3 ++-
tests/test_flattening.py | 3 ++-
tests/test_synth_data.py | 3 ++-
tests/test_training.py | 3 ++-
8 files changed, 27 insertions(+), 21 deletions(-)
diff --git a/dataclass_creator/dyn_slice_dataloader.py b/dataclass_creator/dyn_slice_dataloader.py
index eb51bf4..368f114 100644
--- a/dataclass_creator/dyn_slice_dataloader.py
+++ b/dataclass_creator/dyn_slice_dataloader.py
@@ -103,10 +103,3 @@ class DynSliceDataloader:
"""Return length of output vector."""
return ravel_pytree(self.get_init_vec())[0].shape[0]
-
- @property
- def purpose(self) -> str:
- if type(self.dset) == DynSliceDataset:
- return "reconstruction"
- elif type(self.dset) == PredDynSliceDataset:
- return "prediction"
diff --git a/model/__init__.py b/model/__init__.py
index 852ade4..e69de29 100644
--- a/model/__init__.py
+++ b/model/__init__.py
@@ -1,2 +0,0 @@
-from .autoencoder import AutoEncoder
-from .trainer import Trainer
diff --git a/ros_adetector.py b/ros_adetector.py
index 44b976d..d9b7164 100644
--- a/ros_adetector.py
+++ b/ros_adetector.py
@@ -140,7 +140,7 @@ def start_node():
def test(len_test: int=15, window_size: int=10):
try:
anomaly_detector = AnomalyDetector(data_streams=FrankaData(jt = JointPosData('/joint_states', JointState),
- force = ForceData('/franka_state_controller/F_ext', WrenchStamped)),
+ force = ForceData('/franka_state_controller/F_ext', WrenchStamped)),
model_name="test_w100_b20_e50_l15")
rospub_process = subprocess.Popen( ['rostopic', 'pub', '/start_anomaly', 'std_msgs/Bool', 'True', '--once'] )
diff --git a/run/parameter.py b/run/parameter.py
index fc9363a..144ed55 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -61,7 +61,7 @@ class DatasetParams:
)
@property
- def purpose(self) -> str:
+ def purpose(self) -> Purpose:
if self.pred_offset == 0:
return Purpose.RECONSTRUCTION
elif self.pred_offset > 0:
@@ -99,7 +99,7 @@ class OptimizerHParams:
@dataclass(kw_only=True)
class TrainerParams:
- """Define how the traiing process should happen."""
+ """Define how the training process should happen."""
model: Any #AutoEncoder
loss: Any #LossCalculator
@@ -141,20 +141,31 @@ class ExportParams:
root (Path): root to save. If none, uses dataset_params.root
"""
if root == None:
- logging.info("No path provided, using dataset.root")
+ logging.info("No root provided, using dataset.root")
root = Path(self.dataset_params.root)
path = os.path.join(root, Path(str(self)))
os.makedirs(path, exist_ok=True)
to_dill(self, path, "model.dill")
@classmethod
- def load(cls, dataset_params, trainer_params, root: Path=Path(".")) -> Self:
+ def load(cls,
+ dataset_params: DatasetParams,
+ trainer_params:TrainerParams,
+ root:Path=None) -> Self:
+ """Load from path root/str(self)/model.dill.
+ IN:
+ root (Path): root to load. If none, uses dataset_params.root"""
+ if root == None:
+ logging.info("No root provided, using dataset.root")
+ root = Path(dataset_params.root)
export_params = ExportParams(dataset_params, trainer_params)
- return from_dill(os.path.join(root, Path(str(export_params))), "model.dill")
-
+ return cls.load_from_full_path(os.path.join(root, Path(str(export_params))))
+
@classmethod
def load_from_full_path(cls, root: Path) -> Self:
- return from_dill(root, "model.dill")
+ obj = from_dill(root, "model.dill")
+ assert isinstance(obj, ExportParams), f"The object at {root}/model.dill was not an ExportParams class."
+ return obj
def get_apply_jit(self) -> Callable[[SysVarSet], SysVarSet]:
"""Returns fn which evals model on single sample input."""
diff --git a/run/train_model.py b/run/train_model.py
index e8a6b5f..72b1d41 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -10,7 +10,8 @@ from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
from dataclass_creator.custom_dataclasses import FrankaData
-from model import AutoEncoder, Trainer
+from model.autoencoder import AutoEncoder
+from model.trainer import Trainer
from run.parameter import (
OptimizerHParams,
RepoPaths,
diff --git a/tests/test_flattening.py b/tests/test_flattening.py
index 4f641a2..01408c3 100644
--- a/tests/test_flattening.py
+++ b/tests/test_flattening.py
@@ -10,7 +10,8 @@ from run.parameter import OptimizerHParams
from dataclass_creator import create_synth_nominal_data
from dataclass_creator.custom_dataclasses import SynthData
from dataclass_creator.utils_dataclass.synth_data_gen import DataGeneration
-from model import AutoEncoder, Trainer
+from model.autoencoder import AutoEncoder
+from model.trainer import Trainer
from opt_tools.jax_tools import SysVarSet
batch_size = 50
diff --git a/tests/test_synth_data.py b/tests/test_synth_data.py
index 6f0f4c0..248bf61 100644
--- a/tests/test_synth_data.py
+++ b/tests/test_synth_data.py
@@ -12,7 +12,8 @@ from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from run.load_save_model import to_dill
from run.train_model import train_model
from run.valmodel import ValModel
-from model import AutoEncoder
+from model.autoencoder import AutoEncoder
+from model.trainer import Trainer
from run.parameter import RepoPaths, TrainerParams
# Types
diff --git a/tests/test_training.py b/tests/test_training.py
index ef87268..232aac2 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -8,7 +8,8 @@ from jax.flatten_util import ravel_pytree
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from tests.test_defaults import get_dataloader_test
-from model import AutoEncoder, Trainer
+from model.autoencoder import AutoEncoder
+from model.trainer import Trainer
from run.parameter import OptimizerHParams, RepoPaths
from run.train_model import train_model
from run.load_save_model import to_dill
--
GitLab
From c260af56209c9b75ec881e84fa54538b10907e08 Mon Sep 17 00:00:00 2001
From: Kevin Haninger <khaninger@gmail.com>
Date: Thu, 19 Dec 2024 15:26:06 +0100
Subject: [PATCH 12/12] ros_adetector running with new loader. Threshold
commented out.
---
ros_adetector.py | 33 +++++++++++++++------------------
run/parameter.py | 9 ++++++++-
run/train_model.py | 4 +++-
3 files changed, 26 insertions(+), 20 deletions(-)
diff --git a/ros_adetector.py b/ros_adetector.py
index d9b7164..8409a1a 100644
--- a/ros_adetector.py
+++ b/ros_adetector.py
@@ -18,7 +18,7 @@ from numpy import ndarray as Array
import numpy as np
from ros_datahandler import RosDatahandler, JointPosData, ForceData
-from run.parameter import RepoPaths
+from run.parameter import RepoPaths, ExportParams
from run.load_save_model import load_trained_model_jit, from_dill
from dataclass_creator.custom_dataclasses import FrankaData
@@ -44,8 +44,6 @@ def build_wrench_msg(F = None, frame_id = 'panda_link8') -> WrenchStamped:
# Boundary value for classifying anomaly
anomaly_bound = -100.
-
-
class AnomalyDetector(RosDatahandler):
"""Anomaly detector class which wraps the online RosDatahandler to inherit the ROS subscribers.
@@ -61,14 +59,15 @@ class AnomalyDetector(RosDatahandler):
- makes a classification based on anomaly_bound
"""
- def __init__(self, data_streams, model_name: str):
- self.ae_jit, logger_params = load_trained_model_jit(path_to_models=RepoPaths.trained_models_val, test_name=model_name)
- super().__init__(data_streams,
- sync_stream="force", # when a message arrives on this topic, all the last msgs are added to the data buffer
- window_size=logger_params.window_size)
- self.threshold = from_dill(path=os.path.join(RepoPaths.trained_models_val, Path(model_name)),
- file_name="threshold.dill")
- print("Model: ", model_name)
+ def __init__(self, path_to_model: Path):
+ params = ExportParams.load_from_full_path(path_to_model)
+ self.ae_jit = params.get_apply_jit()
+
+ print(params.dataset_params.datatype)
+ super().__init__(**params.dataset_params.get_datahandler_args())
+ #self.threshold = from_dill(path=path_to_model,
+ # file_name="threshold.dill")
+ print(f"Model: {params}")
self.start = True # True means get_data() and False nothing, currently for testing
self.start_datastreams(ros_version = "ros1") # This sets up subscribers for the data_streams we have in the super().__init__
@@ -80,7 +79,7 @@ class AnomalyDetector(RosDatahandler):
# If we insert data_streams.force.topic_name we only subscribe the forces - why???
# @kevin 21.10: I guess this is b/c the force topic is faster and we werent waiting for all topics
# to be recv'd before. You probably want your callback on the slowest topic.
- self.msg_callback_sub_force = self.subscriber_factory(data_streams.jt.topic_name,
+ self.msg_callback_sub_force = self.subscriber_factory(params.dataset_params.datatype.jt.topic_name,
JointState,
self.msg_callback)
self.start_sub = self.subscriber_factory('/start_anomaly',
@@ -134,14 +133,12 @@ def start_node():
# have the right topic names, etc.
anomaly_detector = AnomalyDetector(data_streams=FrankaData(jt = JointPosData('/joint_states', JointState),
force = ForceData('/franka_state_controller/F_ext', WrenchStamped)),
- model_name="test_w100_b20_e50_l15")
+ model_name="result_w100_b50_e1_l15")
rospy.spin()
def test(len_test: int=15, window_size: int=10):
try:
- anomaly_detector = AnomalyDetector(data_streams=FrankaData(jt = JointPosData('/joint_states', JointState),
- force = ForceData('/franka_state_controller/F_ext', WrenchStamped)),
- model_name="test_w100_b20_e50_l15")
+ anomaly_detector = AnomalyDetector(os.path.join(RepoPaths.example_data_train, "result_w100_b50_e2_l15"))
rospub_process = subprocess.Popen( ['rostopic', 'pub', '/start_anomaly', 'std_msgs/Bool', 'True', '--once'] )
#rosbag_process = subprocess.Popen( ['rosbag', 'play', f'-u {len_test}', 'experiments/01_faston_converging/data/with_search_strategy/test/214_FAILURE_+0,0033_not_inserted.bag'] )
@@ -162,7 +159,7 @@ def test(len_test: int=15, window_size: int=10):
if __name__ == "__main__":
- start_node()
+ #start_node()
#tested with following exp_script
#~/converging/converging-kabelstecker-stecken/exp_scripts$ python3 exp_cable_insert.py
- #test()
+ test()
diff --git a/run/parameter.py b/run/parameter.py
index 144ed55..4f28c8c 100644
--- a/run/parameter.py
+++ b/run/parameter.py
@@ -2,7 +2,7 @@ import os
from enum import Enum
from pathlib import Path
import logging
-from typing import Iterable, Any, Self, Callable
+from typing import Iterable, Any, Self, Callable, Dict
from dataclasses import dataclass, field
from datetime import datetime
@@ -60,6 +60,13 @@ class DatasetParams:
pred_offset=self.pred_offset
)
+ def get_datahandler_args(self) -> Dict:
+ return dict(
+ data_streams = self.datatype,
+ sync_stream = self.sync_stream,
+ window_size = self.window_size
+ )
+
@property
def purpose(self) -> Purpose:
if self.pred_offset == 0:
diff --git a/run/train_model.py b/run/train_model.py
index 72b1d41..f3f419d 100644
--- a/run/train_model.py
+++ b/run/train_model.py
@@ -5,7 +5,9 @@ from dataclasses import replace
import jax
-from ros_datahandler import JointPosData, JointState, ForceData, WrenchStamped
+from ros_datahandler import JointPosData, ForceData
+from sensor_msgs.msg import JointState
+from geometry_msgs.msg import WrenchStamped
from dataclass_creator.dyn_slice_dataloader import DynSliceDataloader
from dataclass_creator.dyn_slice_dataset import DynSliceDataset
from dataclass_creator.utils_dataclass.real_data_gen import load_dataset
--
GitLab