Alex Blank
128 lines
4.0 KiB
Python
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,
|
|
)
|
|
]
|
|
}
|