Рисунок 15. ECG arrhythmia classification using a 2-D convolutional neural network [21]
Реализация сети на PyTorch принимает следующий вид:
1. class HeartNet(nn.Module):
2. def __init__(self, num_classes=7):
3. super(HeartNet, self).__init__()
4. self.features = nn.Sequential(
5. nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
6. nn.ELU(inplace=True),
7. nn.BatchNorm2d(64, eps=0.001),
8. nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
9. nn.ELU(inplace=True),
10. nn.BatchNorm2d(64, eps=0.001),
11. nn.MaxPool2d(kernel_size=2, stride=2),
12. nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
13. nn.ELU(inplace=True),
14. nn.BatchNorm2d(128, eps=0.001),
15. nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
16. nn.ELU(inplace=True),
17. nn.BatchNorm2d(128, eps=0.001),
18. nn.MaxPool2d(kernel_size=2, stride=2),
19. nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
20. nn.ELU(inplace=True),
21. nn.BatchNorm2d(256, eps=0.001),
22. nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
23. nn.ELU(inplace=True),
24. nn.BatchNorm2d(256, eps=0.001),
25. nn.MaxPool2d(kernel_size=2, stride=2)
26. )
27. self.classifier = nn.Sequential(
28. nn.Linear(16 * 16 * 256, 2048),
29. nn.ELU(inplace=True),
30. nn.BatchNorm1d(2048, eps=0.001),
31. nn.Dropout(0.5),
32. nn.Linear(2048, num_classes)
33. )
34. def forward(self, x):
35. x = self.features(x)
36. x = x.view(x.size(0), 16 * 16 * 256)
37. x = self.classifier(x)
38. return x
Рисунок 16. Imbalanced data sampling
Уже упоминалось, что данные ЭКГ являются несбалансированными, поэтому были применены такие техники обучения, как imbalanced sampling (Рис 16.) и применение весов к кросс энтропии.
Однако, улучшить результат применением данных техник не удалось, и наивысшее значение доли правильных ответов было достигнуто без них.
Параметры и результаты:
· Model - [21]
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001 (0.95 decay каждые 1000 итераций)
· Accuracy - 0,9920
Следующим логическим шагом было попробовать воспроизвести модели, применяемые в задачах классификации изображений.
В каждом из экспериментов применялись следующие параметры:
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001
В результате были получены следующие результаты по каждой из примененных к задаче моделей (метрика accuracy):
· AlexNet - 0,9888
· ResNet18 - 0,9925
· VGG16 - 0,9931
· EfficientNetB3 - 0,9930
· EfficientNetB4 - 0,9935
· ShuffleNet - 0,9927
· MobileNetV2 - 0,9934
Лучший результат показала сеть с наибольшим числом параметров EfficientNetB4, однако, гораздо интереснее то, что сеть намного меньшего размера, которая, более того, создана для выполнения вычислений на мобильных устройствах, справляется не хуже.
4.4 1D CNN
Несмотря на то, что задача решается методами двумерных сверточных сетей, намного интереснее решать ее одномерными, т.к. они требуют на вход данные значительно меньшего размера. Одной из таких сетей является сеть, предложенная Pranav Rajpurkar et al. в [25].
Рисунок 17. Cardiologist-Level Arrhythmia Detection with Convolutional Neural Networks [25]
Сверточная нейронная сеть состоит из 16 модифицированных ResNet блоков. Размер фильтра каждого слоя свертки - 16, и каждый слой имеет 64k фильтров, где k начинается с 1 и увеличивается каждый четвертый блок. Кроме того, каждый 2-й блок сжимает входные данные в 2 раза. Отличие от оригинального ResNet блока в том, что авторы добавляют Max Pooling в Residual connection и Dropout. В выпускной квалификационной работе сеть была заимплементирована без dropout слоя (Приложение 4)
Параметры и результаты:
· Model - [25]
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001
· Accuracy - 0,9827
Однако, данная сеть справилась с задачей хуже, чем 2D сети. Скорее всего, это связано с тем, что сеть, предложенная авторами, применяется к другому набору данных.
Поэтому, логическим продолжением было использование одномерной сверточной сети, которая была применена к набору данных MIT-BIH Arrhythmia Database. Xuexiang Xu et al. в работе [26] 2020 года используют наборы данных, входная длина которых равна 300 для решения задачи классификации на 5 типов аритмий.
Реализация сети на PyTorch:
1. class HeartNetIEEE(nn.Module):
2. def __init__(self, num_classes=8):
3. super().__init__()
4. self.features = nn.Sequential(
5. nn.Conv1d(1, 64, kernel_size=5),
6. nn.ReLU(inplace=True),
7. nn.Conv1d(64, 64, kernel_size=5),
8. nn.ReLU(inplace=True),
9. nn.MaxPool1d(2),
10. nn.Conv1d(64, 128, kernel_size=3),
11. nn.ReLU(inplace=True),
12. nn.Conv1d(128, 128, kernel_size=3),
13. nn.ReLU(inplace=True),
14. nn.MaxPool1d(2)
15. )
16.
17. self.classifier = nn.Sequential(
18. nn.Linear(128 * 28, 256),
19. nn.Linear(256, 128),
20. nn.Linear(128, num_classes)
21. )
22.
23. def forward(self, x):
24. x = self.features(x)
25. x = x.view(x.size(0), 128 * 28)
26. x = self.classifier(x)
27. return x
Рисунок 18. ECG Heartbeat Classification Using Convolutional Neural Networks [26]
Параметры и результаты:
· Model - [26]
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001
· Accuracy - 0,9864
Итоговый результат был улучшен, однако он был все еще хуже двумерных сверточных нейронных сетей.
Еще одной попыткой в решении задачи стала реализация Encoder части из [14].
Рисунок 19. Encoder - Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder [14]
Реализация сети на PyTorch:
1. class ZolotyhNet(nn.Module):
2. def __init__(self, num_classes=8):
3. super().__init__()
4. self.features_up = nn.Sequential(
5. nn.Conv1d(1, 8, kernel_size=3, padding=1),
6. nn.BatchNorm1d(8),
7. nn.ReLU(inplace=True),
8. nn.MaxPool1d(2),
9. nn.Conv1d(8, 16, kernel_size=3, padding=1),
10. nn.BatchNorm1d(16),
11. nn.ReLU(inplace=True),
12. nn.MaxPool1d(2),
13. nn.Conv1d(16, 32, kernel_size=3, padding=1),
14. nn.BatchNorm1d(32),
15. nn.ReLU(inplace=True),
16. nn.MaxPool1d(2),
17. nn.Conv1d(32, 32, kernel_size=3, padding=1),
18. nn.BatchNorm1d(32),
19. nn.ReLU(inplace=True),
20. nn.MaxPool1d(2),
21. nn.Conv1d(32, 1, kernel_size=3, padding=1),
22. Flatten(),
23. )
24. self.features_down = nn.Sequential(
25. Flatten(),
26. nn.Linear(128,64),
27. nn.BatchNorm1d(64),
28. nn.ReLU(inplace=True),
29. nn.Linear(64, 16),
30. nn.BatchNorm1d(16),
31. nn.ReLU(inplace=True),
32. nn.Linear(16, 8)
33. )
34. self.classifier = nn.Linear(8, num_classes)
35.
36. def forward(self, x):
37. out_up = self.features_up(x)
38. out_down = self.features_down(x)
39. out_middle = out_up + out_down
40.
41. out = self.classifier(out_middle)
42.
43. return out
Параметры и результаты:
· Model - [14]
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001
· Accuracy - 0,9886
Эта сеть показала наивысший результат среди остальных одномерных сверточных нейронных сетей.
4.5 Реализация нового решения EcgResNet34
Было принято большое количество попыток в нахождении решения, которое смогло бы не только показать схожий с 2D сетями результат, но и улучшить его.
Финальным решением стала модификация модели ResNet-34. Несмотря на то, что решение, предложенное Pranav Rajpurkar et al. [25] не показало высоких результатов, идея модифицировать ResNet показалась достаточно перспективной. В своей имплементации авторы используют недоступный в открытом доступе набор данных и используют последовательности большей длины, они сжимают входные данные в раз, что не подходит для размера последовательностей небольшой длины, как в MIT-BIH Arrhythmia Database. Кроме того, авторы статьи добавляют пулинг, который в выпускной квалификационной работе был заменен сверточным слоем со страйдом 2. Размер фильтра каждой свертки - 17. В первом блоке количество выходных каналов 64, далее 5 блоков с 128 фильтрами, затем 5 блоков с 256 фильтрами и, наконец, 5 блоков с количеством выходных каналов равным 512. Архитектура модифицированного ResNet блока изображена на Рис 20. В финальной части сверточной нейронной сети применяется AveragePooling и полносвязный слой на 8 классов. Финальная архитектура представлена на Рис 21.
Рисунок 20. Модифицированный ResNet
Параметры и результаты:
· Model - Novel EcgResNet34
· Batch size - 128
· Optimizer - Adam
· Learning rate - 0.001
· Accuracy - 0,9938
Precision / recall per class
N - 0.994670 / 0.998262
L - 0.998758 / 0.996283
V - 0.990099 / 0.983146
\ - 1.000000 / 1.000000
R - 0.994505 / 0.997245
A - 0.957447 / 0.882353
! - 0.955556 / 0.914894
E - 1.000000 / 1.000000
Рисунок 21. Предложенная в ВКР архитектура сверточной нейронной сети
Таким образом, с использованием архитектуры, предложенной в выпускной квалификационной работе, удалось добиться наилучшего результата, как для 1D, так и для 2D моделей. Ниже, на Рис 22. Изображена матрица ошибок наилучшего решения.
Рисунок 22. Матрица ошибок
4.6 Выводы
В выпускной квалификационной работе было реализовано несколько архитектур из последних статей, связанных с классификацией аритмий, а также сделан ряд экспериментов с использованием популярных для решения задач классификации изображений сверточных нейронных сетей. В результате была реализована новая архитектура сверточной нейронный сети, показавшая наивысшее значение целевой метрики на валидационном наборе данных - 99,38%. Все выполненные эксперименты представлены на Рис. 23.
Рисунок 23. Проделанные эксперименты, параметры запуска и результаты
5 Особенности выпускной квалификационной работы
Одной из целей выпускной квалификационной работы было поддержание исследований в области анализа данных ЭКГ, поэтому было уделено особое внимание реализации используемых моделей. Так, был создан репозиторий, доступный по адресу https://github.com/lxdv/ecg-classification, в котором указана инструкция по запуску всех используемых в ВКР решений, от скачивания набора данных, до запуска тренировки, валидации и тестирования. Исходный код каждого доступен в ветках experiments/exp-<>, а таблица с экспериментами доступна по адресу
https://docs.google.com/spreadsheets/d/159OjSlXuItvngeQwBxC5NaQbU9PjaMN4mY--bX26m1o
В репозитории поддерживается парадигма конфигурационных файлов
· train.py --config configs/training/<config>.json
· inference.py -config configs/inference/config.json
· pipeline.py --config configs/pipelines/config/json
1. {
2. "exp_name":"EcgResNet34",
3. "type": "Trainer1D",
4. "device": "cpu",
5. "model": "EcgResNet34",
6.
7. "train_json": "data/train.json",
8. "val_json": "data/val.json",
9. "mapping_json": "data/class-mapper.json",
10. "num_classes": 8,
11.
12. "batch_size": 128,
13. "optim": "Adam",
14. "optim_params":
15. {
16. "lr": 1e-3
17. },
18.
19. "num_workers": 3,
20. "exp_dir" : "experiments/",
21.
22. "model_path": 0,
23. "epochs": 1
24. }
Такой подход дает возможность исследователям переиспользования кода в целях ознакомления с полученными результатами, его улучшения и добавления новых решений. Это также становится возможным благодаря применению GitHub Actions в качестве средства непрерывной интерграции, который запускает тестовую тренировку (на удаленной пустой машине) каждого решения с установлением всех необходимых зависимостей. Вот, например, конфигурация для установки всех зависимостей и запуска тренировки сети, предложенной в выпускной квалификационной работе.
1. name: EcgResNet34
2.
3. on:
4. push:
5. branches: [ dev, feature/actions-ci ]
6. pull_request:
7. branches: [ master ]
8.
9. jobs:
10. Train:
11. runs-on: ubuntu-16.04
12. steps:
13. - uses: actions/checkout@v2
14. - name: Set up Python 3.7
15. uses: actions/setup-python@v1
16. with:
17. python-version: 3.7
18. - name: Install dependencies
19. run: |
20. python -m pip install -r requirements.txt
21. - name: Generate data files
22. run: |
23. python dataset-generation-pool.py
24. working-directory:
25. ./scripts
26. - name: Generate 1D annotation files
27. run: |
28. python annotation-generation-1d.py
29. working-directory:
30. ./scripts
31. - name: Running training
32. run: |
33. python train.py --config configs/training/EcgResNet34.json
Кроме того, во все решения интегрировано логирование с использованием TensorBoard (Рис 24).
Рисунок 24. Примеры визуализации функции потерь и целевой метрики с использованием TensorBoard
6 Автоматический анализ электрокардиограммы
В выпускной квалификационной работе был реализован прототип анализа электрокардиограммы, т.е. возможность применить лучшее решение к ЭКГ в формате WFDB - программа сделает все необходимое - найдет каждый сердечный удар, выполнит по ним классификацию и выдаст результат в удобном для человека формате - HTML страницы, которая может быть открыта любым браузером. Пример результирующего файла представлен на Рис 25.
На Рис 25. в рамке указано настоящее значение, а над ним предсказанное. Получение результатов достигается при помощи использования библиотеки plotly (Приложение 7). Для удобства визуализация неаномальных ударов была выключена. Таким образом, итоговая визуализация изображена на рисунках ниже, а используемый код представлен в Приложении 5.
Рисунок 26. Прототип работы программы
Рисунок 27. Прототип работы программы
Рисунок 28. Прототип работы программы
Рисунок 29. Прототип работы программы
Заключение
Целью выпускной квалификационной работы была реализация алгоритма автоматической классификации аритмий с использованием сверточных нейронных сетей. Эксперименты показали, что использование нейронных сетей в задаче возможно и, более того, является перспективным выбором в решении задачи классификации болезней, связанных с аномальной работой сердца. В работе было проведено сравнение существующих архитектур, представлено новое решение и разработана программа, способная производить автоматический анализ кардиограммы и находить аномальные сердечные удары. Весь исходный код проделанной работы доступен в отрытом доступе, что позволяет исследователям ознакомиться с представленной работой и интегрировать автоматический анализ электрокардиограмм в свой бизнес или исследовательский продукт. Итоговый результат доли правильных ответов (accuracy) на валидационном наборе данных MIT-BIH Arrhythmia Database составляет 99,38%, что является наилучшим по сравнению с другими современными подходами решения задачи.