optuna_tune ¶

def get_metrics(
    self,
    model_instance: StimulusModel,
    data_loader: torch.utils.data.DataLoader,
    loss_dict: dict[str, torch.nn.Module],
    device: torch.device,
) -> dict[str, float]:
    """Compute the objective metric(s) for the tuning process."""

    def update_metric_dict(
        metric_dict: dict[str, torch.Tensor],
        metrics: dict[str, torch.Tensor],
        loss: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        """Update the metric dictionary with the new metrics and loss."""
        for key, value in metrics.items():
            if key not in metric_dict:
                if value.ndim == 0:
                    metric_dict[key] = value.unsqueeze(0)
                else:
                    metric_dict[key] = value
            elif value.ndim == 0:
                metric_dict[key] = torch.cat([metric_dict[key], value.unsqueeze(0)], dim=0)
            else:
                metric_dict[key] = torch.cat([metric_dict[key], value], dim=0)
        if "loss" not in metric_dict:
            if loss.ndim == 0:
                metric_dict["loss"] = loss.unsqueeze(0)
            else:
                metric_dict["loss"] = loss
        elif loss.ndim == 0:
            metric_dict["loss"] = torch.cat([metric_dict["loss"], loss.unsqueeze(0)], dim=0)
        else:
            metric_dict["loss"] = torch.cat([metric_dict["loss"], loss], dim=0)
        return metric_dict

    # set model in eval mode
    model_instance.eval()

    metric_dict: dict = {}

    for batch in data_loader:
        try:
            # Move all tensors to device (sample_id is now an integer tensor)
            device_batch = {}
            for key, value in batch.items():
                try:
                    device_batch[key] = value.to(device, non_blocking=True)
                except AttributeError as e:
                    raise AttributeError(
                        f"Error moving '{key}' to device during inference. Expected tensor but got {type(value)}. "
                        f"This usually happens when dataset columns contain non-tensor data. "
                        f"Original error: {e}",
                    ) from e

            # Perform a batch update
            loss, metrics = model_instance.inference(batch=device_batch, **loss_dict)

        except RuntimeError as e:
            if ("CUDA out of memory" in str(e) and self.device.type == "cuda") or (
                "MPS backend out of memory" in str(e) and self.device.type == "mps"
            ):
                logger.warning(f"{self.device.type.upper()} out of memory during training: {e}")
                logger.warning("Falling back to CPU for this trial")
                temp_device = torch.device("cpu")
                model_instance = model_instance.to(temp_device)
                # Consider adjusting batch size or other parameters
                device_batch = {key: value.to(temp_device) for key, value in batch.items()}
                # Retry the batch
                loss, metrics = model_instance.inference(batch=device_batch, **loss_dict)
            else:
                raise

        metric_dict = update_metric_dict(metric_dict, metrics, loss)

    # devide all metrics by number of batches
    for key in metric_dict:
        metric_dict[key] = metric_dict[key].mean()

    # Convert tensors to floats before returning
    return {k: v.item() if isinstance(v, torch.Tensor) else v for k, v in metric_dict.items()}

def objective(
    self,
    model_instance: StimulusModel,
    train_loader: torch.utils.data.DataLoader,
    val_loader: torch.utils.data.DataLoader,
    loss_dict: dict[str, torch.nn.Module],
    device: torch.device,
) -> dict[str, float]:
    """Compute the objective metric(s) for the tuning process.

    The objectives are outputed by the model's batch function in the form of loss, metric_dictionary.
    """
    train_metrics = self.get_metrics(model_instance, train_loader, loss_dict, device)
    val_metrics = self.get_metrics(model_instance, val_loader, loss_dict, device)

    # add train_ and val_ prefix to related keys.
    return {
        **{f"train_{k}": v for k, v in train_metrics.items()},
        **{f"val_{k}": v for k, v in val_metrics.items()},
    }

def save_checkpoint(
    self,
    trial: optuna.Trial,
    model_instance: StimulusModel,
    optimizer: torch.optim.Optimizer,
    complete_suggestions: dict,
) -> None:
    """Save the model and optimizer to the trial."""
    # Convert model to CPU before saving to avoid device-specific tensors
    model_instance = model_instance.cpu()
    optimizer_state = optimizer.state_dict()

    # Convert optimizer state to CPU tensors
    for param in optimizer_state["state"].values():
        for k, v in param.items():
            if isinstance(v, torch.Tensor):
                param[k] = v.cpu()
    unique_id = str(uuid.uuid4())[:8]
    model_path = f"{trial.number}_{unique_id}_model.safetensors"
    optimizer_path = f"{trial.number}_{unique_id}_optimizer.pt"
    model_suggestions_path = f"{trial.number}_{unique_id}_model_suggestions.json"
    safe_save_model(model_instance, model_path)
    torch.save(optimizer_state, optimizer_path)
    with open(model_suggestions_path, "w") as f:
        json.dump(complete_suggestions, f)
    artifact_id_model = optuna.artifacts.upload_artifact(
        artifact_store=self.artifact_store,
        file_path=model_path,
        study_or_trial=trial.study,
    )
    artifact_id_optimizer = optuna.artifacts.upload_artifact(
        artifact_store=self.artifact_store,
        file_path=optimizer_path,
        study_or_trial=trial.study,
    )
    artifact_id_model_suggestions = optuna.artifacts.upload_artifact(
        artifact_store=self.artifact_store,
        file_path=model_suggestions_path,
        study_or_trial=trial.study,
    )
    # delete the files from the local filesystem
    try:
        os.remove(model_path)
        os.remove(optimizer_path)
        os.remove(model_suggestions_path)
    except FileNotFoundError:
        logger.info(
            f"File was already deleted: {model_path} or {optimizer_path} or {model_suggestions_path}, most likely due to pruning",
        )
    trial.set_user_attr("model_id", artifact_id_model)
    trial.set_user_attr("model_path", model_path)
    trial.set_user_attr("optimizer_id", artifact_id_optimizer)
    trial.set_user_attr("optimizer_path", optimizer_path)
    trial.set_user_attr("model_suggestions_id", artifact_id_model_suggestions)
    trial.set_user_attr("model_suggestions_path", model_suggestions_path)