|
| 1 | +# Finetuning |
| 2 | + |
| 3 | +## Overview |
| 4 | + |
| 5 | +Since v1.2.0, the complete training code is available at [hsf_train](https://github.com/clementpoiret/hsf_train). |
| 6 | +The `hsf_train` repository also contains easy to use scripts to train OR **finetune your own models**. |
| 7 | + |
| 8 | +## Purpose & Use Cases |
| 9 | + |
| 10 | +The goal of HSF is to provide foundational models that can be used as a starting point for further development and customization in Hippocampal Subfields segmentation. |
| 11 | + |
| 12 | +We are aware that the provided models may not be perfect for every use case, and that's why we provide the training code and the possibility to finetune the models. |
| 13 | + |
| 14 | +If you want to use HSF on MRIs that are very different from the ones used for training, or if you want to segment the hippocampal subfields in a specific way, you are at the right place. |
| 15 | + |
| 16 | +## Configuration |
| 17 | + |
| 18 | +The `hsf_train` repository contains a `conf` directory, analogous to the `conf` directory in the `hsf` repository. This directory contains the configuration files for the training and finetuning scripts. |
| 19 | + |
| 20 | +Here is the default configuration file for the finetuning script: |
| 21 | + |
| 22 | +```yaml |
| 23 | +mode: decoder # encoder or decoder, defines which part of the model to finetune (see below) |
| 24 | +depth: -1 # -1 for all layers, 0 for the first layer, 1 for the second layer, etc. |
| 25 | +unfreeze_frequency: 4 # how often to unfreeze a layer |
| 26 | +out_channels: 6 # number of output channels (subfields) in case you want to segment a different number of subfields |
| 27 | +``` |
| 28 | +
|
| 29 | +## Getting Started |
| 30 | +
|
| 31 | +### Installation |
| 32 | +
|
| 33 | +You will need to clone the repository, install PyTorch, and install the required packages: |
| 34 | +
|
| 35 | +```bash |
| 36 | +git clone https://github.com/clementpoiret/hsf_train.git |
| 37 | +cd hsf_train |
| 38 | +conda create -n hsf_train python=3.10 |
| 39 | +conda activate hsf_train |
| 40 | +conda install pytorch torchvision torchaudio cudatoolkit=12.1 -c pytorch -c nvidia |
| 41 | +pip install -r requirements.txt |
| 42 | +``` |
| 43 | + |
| 44 | +### Custom Dataset |
| 45 | + |
| 46 | +Whoever wants to finetune the models will need to provide their own dataset. The heavier the changes are between the training dataset and the custom dataset, the more data you will need. |
| 47 | + |
| 48 | +You will need to adapt the example `custom_dataset.yaml` file to suit your needs, such as: |
| 49 | + |
| 50 | +```yaml |
| 51 | +main_path: "/mnt/data/hsf/" |
| 52 | +output_path: "/mnt/hsf/models/" |
| 53 | +batch_size: 1 |
| 54 | +num_workers: 16 |
| 55 | +pin_memory: True |
| 56 | +train_ratio: .9 |
| 57 | +replace: False |
| 58 | +k_sample: Null # i.e. k = train_ratio * num_samples |
| 59 | +train_val_test_idx: Null |
| 60 | +train_on_all: False |
| 61 | + |
| 62 | +datasets: |
| 63 | + clark: |
| 64 | + path: "hippocampus_clark_3T" |
| 65 | + ca_type: "1/23" |
| 66 | + patterns: |
| 67 | + right_t2: |
| 68 | + mri: "**/t2w_Hippocampus_right_ElasticSyN_crop.nii.gz" |
| 69 | + label: "t2w_Hippocampus_right_ElasticSyN_seg_crop.nii.gz" |
| 70 | + left_t2: |
| 71 | + mri: "**/t2w_Hippocampus_left_ElasticSyN_crop.nii.gz" |
| 72 | + label: "t2w_Hippocampus_left_ElasticSyN_seg_crop.nii.gz" |
| 73 | + averaged_right_t2: |
| 74 | + mri: "**/averaged_t2w_Hippocampus_right_ElasticSyN_crop.nii.gz" |
| 75 | + label: "averaged_t2w_Hippocampus_right_ElasticSyN_seg_crop.nii.gz" |
| 76 | + averaged_left_t2: |
| 77 | + mri: "**/averaged_t2w_Hippocampus_left_ElasticSyN_crop.nii.gz" |
| 78 | + label: "averaged_t2w_Hippocampus_left_ElasticSyN_seg_crop.nii.gz" |
| 79 | + labels: |
| 80 | + 1: 1 |
| 81 | + 2: 2 |
| 82 | + 3: 3 |
| 83 | + 4: 4 |
| 84 | + 5: 5 |
| 85 | + 6: 6 |
| 86 | + 7: 7 |
| 87 | + labels_names: |
| 88 | + 1: "DG" |
| 89 | + 2: "CA2/3" |
| 90 | + 3: "CA1" |
| 91 | + 4: "PRESUB" |
| 92 | + 5: "UNCUS" |
| 93 | + 6: "PARASUB" |
| 94 | + 7: "KYST" |
| 95 | +``` |
| 96 | +
|
| 97 | +??? info "Comment on the `out_channels` parameter" |
| 98 | + You can see in the example above that we have 7 subfields. However, the `out_channels` parameter also includes the background (class 0), so you should set it to 8. |
| 99 | + |
| 100 | +### Finetuning |
| 101 | + |
| 102 | +Then, you can run the finetuning script: |
| 103 | + |
| 104 | +```bash |
| 105 | +python finetune.py \ |
| 106 | + datasets=custom_dataset \ |
| 107 | + finetuning.out_channels=8 \ |
| 108 | + models.lr=1e-3 \ |
| 109 | + models.use_forgiving_loss=False |
| 110 | +``` |
| 111 | + |
| 112 | +## Training Tips |
| 113 | + |
| 114 | +Numerous factors can influence the quality of the finetuned model. Here are some tips to help you get the best results: |
| 115 | + |
| 116 | +- **Number of manual segmentations**: The more manual segmentations you have, the better. The more subtle the differences between our original dataset and your custom dataset, the less data you will need, |
| 117 | +- **Learning rate**: We recommend starting with a learning rate of 1e-3, |
| 118 | +- **Depth**: We recommend starting with a depth of -1, which means that all layers will be finetuned. If your changes are small, you can finetune less layers, or even only the final layer (depth = 0), |
| 119 | +- **Number of epochs**: We recommend starting with 16 epochs if the depth is -1, otherwise a rule of thumb might be `epochs = (depth + 1) * unfreeze_frequency`, |
| 120 | +- **Mode**: It depends on the type of changes you want to make. If you change the image modality (e.g. another contrast than T1w or T2w, a magnetic field intensity larger than 7T, etc.), you might want to finetune the encoder. If you want to segment a different number of subfields or use another segmentation guideline, you might want to finetune the decoder. |
| 121 | +- **Unfreeze frequency**: We recommend starting with an unfreeze frequency of 4. Check the learning curves, you should reach a plateau before unfreezing the next layer, |
0 commit comments