Как создать собственные сплиты в nnU-Net

Иногда 5-кратная кросс-валидация по умолчанию, которую предлагает nnU-Net, не подходит для проекта. Возможно, вы хотите использовать 3-кратную кросс-валидацию? Или ваши обучающие примеры нельзя разделить случайным образом и требуется тщательная стратификация. Не беспокойтесь, nnU-Net поможет вам (он действительно может всё <3).

Сплиты, которые использует nnU-Net, генерируются в функции do_split класса nnUNetTrainer. Эта функция сначала ищет существующие сплиты, сохранённые в виде файла, и если сплит не существует, она создаёт его. Поэтому, если вы хотите повлиять на сплит, вручную создайте файл сплита, который будет распознан и использован!

Файл сплитов находится в папке nnUNet_preprocessed/DATASETXXX_NAME. Поэтому рекомендуется сначала заполнить эту папку, запустив nnUNetv2_plan_and_preproccess.

Сплиты хранятся в виде файла .json. Они представляют собой простой список Python. Длина этого списка — это количество сплитов, которое он содержит (так что в стандартном nnU-Net их 5). Каждая запись списка представляет собой словарь с ключами «train» и «val». Значениями являются списки с идентификаторами обучающих примеров в каждом наборе. Чтобы проиллюстрировать это, я просто работаю с файлом Dataset002 в качестве примера:

In [1]: from batchgenerators.utilities.file_and_folder_operations import load_json

In [2]: splits = load_json('splits_final.json')

In [3]: len(splits)
Out[3]: 5

In [4]: splits[0].keys()
Out[4]: dict_keys(['train', 'val'])

In [5]: len(splits[0]['train'])
Out[5]: 16

In [6]: len(splits[0]['val'])
Out[6]: 4

In [7]: print(splits[0])
{'train': ['la_003', 'la_004', 'la_005', 'la_009', 'la_010', 'la_011', 'la_014', 'la_017', 'la_018', 'la_019', 'la_020', 'la_022', 'la_023', 'la_026', 'la_029', 'la_030'],
'val': ['la_007', 'la_016', 'la_021', 'la_024']}

Если вы всё ещё не уверены, как должны выглядеть сплиты, просто загрузите какой-нибудь эталонный набор данных с сайта Medical Decathlon, начните обучение (чтобы сгенерировать сплиты) и вручную проверьте файл .json с помощью любого текстового редактора!

Чтобы создать свои собственные сплиты, всё, что вам нужно сделать, это воспроизвести структуру данных, описанную выше, и сохранить её как splits_final.json в папке nnUNet_preprocessed/DATASETXXX_NAME. Затем используйте nnUNetv2_train и т. д., как обычно.