# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Training
#

# ### Building and compiling a model

# +
# %%capture
import h5py
from chimeranet.models import ChimeraPPModel

# probe shape of dataset and set embedding dimension
dataset_path = 'example-dataset-train.hdf5'
with h5py.File(dataset_path, 'r') as f:
    _, T, F, C = f['y/embedding'].shape
D = 20
cm = ChimeraPPModel(T, F, C, D)

# build_model returns Keras' Model object
model = cm.build_model()
model.compile(
    'rmsprop',
    loss={
        'embedding': cm.loss_deepclustering(),
        'mask': cm.loss_mask()
    },
    loss_weights={
        'embedding': 0.9,
        'mask': 0.1
    }
)
# -

# ### Training a model

# +
# load train and validation data
dataset_validation_path = 'example-dataset-validation.hdf5'
with h5py.File(dataset_path, 'r') as f:
    x_train = f['x'][()]
    y_train = {'mask': f['y/mask'][()], 'embedding': f['y/embedding'][()]}
with h5py.File(dataset_validation_path, 'r') as f:
    x_validation = f['x'][()]
    y_validation = {'mask': f['y/mask'][()], 'embedding': f['y/embedding'][()]}

# train model by model.fit function
model.fit(
    x=x_train,
    y=y_train,
    validation_data=(x_validation, y_validation),
    batch_size=32,
    epochs=10
)
# save the model
model_path = 'example-model.hdf5'
model.save(model_path)

StarcoderdataJupyterScriptsDedupFiltered