Files
temperature-based-fertility…/code/configs/cnn_run_config.py
T
Alex Blank c6defa2065 fixes
2025-05-19 13:59:16 +02:00

128 lines
4.0 KiB
Python

from functools import partial
from configs.feature_config import feature_config
from models.cnn import CNNTransformer
from utils.data_utils import *
from utils.training import *
from models.lstm import *
from models.utils import *
from models.collation import *
take_every_nth = int(288 / 12)
shift_in_hours = 12
input_window_length = (288 // take_every_nth) * 80
# output_window_length = (288 // take_every_nth) * 1
output_window_length = 1
output_window_offset = input_window_length + (288 // take_every_nth) * 0
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
max_lr = 1e-5
batch_size = 256
def get_cnn_run_config(
run_name: str,
run_description: str,
input_window_length: int,
output_window_length: int,
take_every_nth: int,
shift_in_hours: int,
output_window_offset: int,
batch_size: int,
model_parameters: dict,
max_lr: float,
num_epochs: int,
patience: int,
feature_config: dict):
"""
Get the configuration for the CNN model
Returns:
run_configuration: configuration for the CNN model
"""
base_model_config = {
"model_name": "cnn_regressor",
"version": "1.0.0",
"model_class": CNNTransformer,
"feature_config": feature_config | {"ignored_features":
[
"fertility_probability",
"ov_over_probability",
# "days_relative_to_ov",
# "is_biphasic",
"ov_day",
]},
"preprocessing": {
"window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
"take_every_nth": take_every_nth,
"min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
},
"batch_fn": produce_window_batches,
"collate_fn": simple_x_y_collate,
"model_creation_fn": simple_model_creation,
"model_save_fn": simple_model_save,
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
"batch_loss_fn": get_model_loss,
"actual_fn": simple_get_y,
"predict_fn": simple_x_y_predict,
"model_parameters": {
**model_parameters | {"seq_len": input_window_length},
},
"input_window_length": input_window_length,
"output_window_length": output_window_length,
"output_window_offset": output_window_offset,
}
base_training_config = {
"batch_size": batch_size,
"model_class": CNNTransformer,
"learning_parameters": {
"learning_rate": max_lr * (batch_size / 4),
"epochs": num_epochs,
"patience": patience
},
"loss_functions": [
nn.MSELoss(),
nn.BCEWithLogitsLoss(),
],
"max_grad_norm": 1.0,
"train_size": 0.7,
"val_size": 0.15,
"test_size": 0.15,
}
return {
"name": run_name,
"description": run_description,
"model_configuration": base_model_config,
"training_configuration": base_training_config,
}
run_configuration = {
"runs": [
get_cnn_run_config(
run_name="cnn_ovulation_regression",
run_description="CNN model for ovulation regression",
input_window_length=input_window_length,
output_window_length=output_window_length,
take_every_nth=take_every_nth,
shift_in_hours=shift_in_hours,
output_window_offset=output_window_offset,
batch_size=batch_size,
model_parameters={
"cnn_channels": 64,
"kernel_size": 3,
"embed_dim": 64,
"num_enc_layers": 2,
"num_heads": 2,
},
max_lr=max_lr,
num_epochs=10,
patience=3,
feature_config=feature_config,
)
]
}