Builds a model from a configuration and fits it to a dataset. Args: loader: the loader containing the data to train on/validate with run_config: the model and run configuration task: the task to train the model for device: the torch device to train on (such as "c
(
loader: Loader,
run_config: dict,
task: Task,
device: str | None = "cpu",
gpus: list[int] | None = None,
logger_config: dict | None = None,
snapshot_path: str | Path | None = None,
transform: A.BaseCompose | None = None,
inference_transform: A.BaseCompose | None = None,
max_snapshots_to_keep: int | None = None,
load_head_weights: bool = True,
)
| 42 | |
| 43 | |
| 44 | def train( |
| 45 | loader: Loader, |
| 46 | run_config: dict, |
| 47 | task: Task, |
| 48 | device: str | None = "cpu", |
| 49 | gpus: list[int] | None = None, |
| 50 | logger_config: dict | None = None, |
| 51 | snapshot_path: str | Path | None = None, |
| 52 | transform: A.BaseCompose | None = None, |
| 53 | inference_transform: A.BaseCompose | None = None, |
| 54 | max_snapshots_to_keep: int | None = None, |
| 55 | load_head_weights: bool = True, |
| 56 | ) -> None: |
| 57 | """Builds a model from a configuration and fits it to a dataset. |
| 58 | |
| 59 | Args: |
| 60 | loader: the loader containing the data to train on/validate with |
| 61 | run_config: the model and run configuration |
| 62 | task: the task to train the model for |
| 63 | device: the torch device to train on (such as "cpu", "cuda", "mps") |
| 64 | gpus: the list of GPU indices to use for multi-GPU training |
| 65 | logger_config: the configuration of a logger to use |
| 66 | snapshot_path: if continuing to train from a snapshot, the path containing the |
| 67 | weights to load |
| 68 | transform: if defined, overwrites the transform defined in the model config |
| 69 | inference_transform: if defined, overwrites the inference transform defined in |
| 70 | the model config |
| 71 | max_snapshots_to_keep: the maximum number of snapshots to store for each model |
| 72 | load_head_weights: When `snapshot_path` is not None and a pose model is being |
| 73 | trained, whether to load the head weights from the saved snapshot. |
| 74 | """ |
| 75 | weight_init = None |
| 76 | pretrained = True |
| 77 | |
| 78 | if weight_init_cfg := run_config["train_settings"].get("weight_init"): |
| 79 | weight_init = WeightInitialization.from_dict(weight_init_cfg) |
| 80 | pretrained = False |
| 81 | elif snapshot_path is not None: |
| 82 | # If we're loading from a snapshot, don't use pretrained backbone weights |
| 83 | # since the weights will be loaded from the snapshot |
| 84 | pretrained = False |
| 85 | |
| 86 | if task == Task.DETECT: |
| 87 | model = DETECTORS.build( |
| 88 | run_config["model"], |
| 89 | weight_init=weight_init, |
| 90 | pretrained=pretrained, |
| 91 | ) |
| 92 | |
| 93 | else: |
| 94 | model = PoseModel.build( |
| 95 | run_config["model"], |
| 96 | weight_init=weight_init, |
| 97 | pretrained_backbone=pretrained, |
| 98 | ) |
| 99 | |
| 100 | if max_snapshots_to_keep is not None: |
| 101 | run_config["runner"]["snapshots"]["max_snapshots"] = max_snapshots_to_keep |
no test coverage detected