main.py 10.7 KB
Newer Older
1
import argparse
2

3
4
5
6
7
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from EEGGraphDataset import EEGGraphDataset
8
from joblib import load
9
from sklearn import preprocessing
10
11
12
from sklearn.metrics import balanced_accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split
from torch.utils.data import WeightedRandomSampler
13

14
from dgl.dataloading import GraphDataLoader
15
16
17

if __name__ == "__main__":
    # argparse commandline args
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
    parser = argparse.ArgumentParser(
        description="Execute training pipeline on a given train/val subjects"
    )
    parser.add_argument(
        "--num_feats",
        type=int,
        default=6,
        help="Number of features per node for the graph",
    )
    parser.add_argument(
        "--num_nodes", type=int, default=8, help="Number of nodes in the graph"
    )
    parser.add_argument(
        "--gpu_idx",
        type=int,
        default=0,
        help="index of GPU device that should be used for this run, defaults to 0.",
    )
    parser.add_argument(
        "--num_epochs",
        type=int,
        default=40,
        help="Number of epochs used to train",
    )
    parser.add_argument(
        "--exp_name", type=str, default="default", help="Name for the test."
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=512,
        help="Batch Size. Default is 512.",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="shallow",
        help="type shallow to use shallow_EEGGraphDataset; "
        "type deep to use deep_EEGGraphDataset. Default is shallow",
    )
58
59
60
    args = parser.parse_args()

    # choose model
61
    if args.model == "shallow":
62
63
        from shallow_EEGGraphConvNet import EEGGraphConvNet

64
    if args.model == "deep":
65
66
67
68
69
70
71
72
        from deep_EEGGraphConvNet import EEGGraphConvNet

    # set the random seed so that we can reproduce the results
    np.random.seed(42)
    torch.manual_seed(42)

    # use GPU when available
    _GPU_IDX = args.gpu_idx
73
74
75
    _DEVICE = torch.device(
        f"cuda:{_GPU_IDX}" if torch.cuda.is_available() else "cpu"
    )
76
    torch.cuda.set_device(_DEVICE)
77
    print(f" Using device: {_DEVICE} {torch.cuda.get_device_name(_DEVICE)}")
78
79
80
81
82
83
84
85
86
87
88
89
90
91

    # load patient level indices
    _DATASET_INDEX = pd.read_csv("master_metadata_index.csv")
    all_subjects = _DATASET_INDEX["patient_ID"].astype("str").unique()
    print(f"Subject list fetched! Total subjects are {len(all_subjects)}.")

    # retrieve inputs
    num_nodes = args.num_nodes
    _NUM_EPOCHS = args.num_epochs
    _EXPERIMENT_NAME = args.exp_name
    _BATCH_SIZE = args.batch_size
    num_feats = args.num_feats

    # set up input and targets from files
92
93
94
95
    memmap_x = f"psd_features_data_X"
    memmap_y = f"labels_y"
    x = load(memmap_x, mmap_mode="r")
    y = load(memmap_y, mmap_mode="r")
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

    # normalize psd features data
    normd_x = []
    for i in range(len(y)):
        arr = x[i, :]
        arr = arr.reshape(1, -1)
        arr2 = preprocessing.normalize(arr)
        arr2 = arr2.reshape(48)
        normd_x.append(arr2)

    norm = np.array(normd_x)
    x = norm.reshape(len(y), 48)
    # map 0/1 to diseased/healthy
    label_mapping, y = np.unique(y, return_inverse=True)
    print(f"Unique labels 0/1 mapping: {label_mapping}")

    # split the dataset to train and test. The ratio of test is 0.3.
113
114
115
    train_and_val_subjects, heldout_subjects = train_test_split(
        all_subjects, test_size=0.3, random_state=42
    )
116
117
118

    # split the dataset using patient indices
    train_window_indices = _DATASET_INDEX.index[
119
120
        _DATASET_INDEX["patient_ID"].astype("str").isin(train_and_val_subjects)
    ].tolist()
121
    heldout_test_window_indices = _DATASET_INDEX.index[
122
123
        _DATASET_INDEX["patient_ID"].astype("str").isin(heldout_subjects)
    ].tolist()
124
125
126
127
128

    # define model, optimizer, scheduler
    model = EEGGraphConvNet(num_feats)
    loss_function = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
129
130
131
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=[i * 10 for i in range(1, 26)], gamma=0.1
    )
132
133

    model = model.to(_DEVICE).double()
134
135
136
137
138
139
    num_trainable_params = np.sum(
        [
            np.prod(p.size()) if p.requires_grad else 0
            for p in model.parameters()
        ]
    )
140
141
142
143
144
145
146
147
148
149
150
151
152
153

    # Dataloader========================================================================================================

    # use WeightedRandomSampler to balance the training dataset
    NUM_WORKERS = 4

    labels_unique, counts = np.unique(y, return_counts=True)

    class_weights = np.array([1.0 / x for x in counts])
    # provide weights for samples in the training set only
    sample_weights = class_weights[y[train_window_indices]]
    # sampler needs to come up with training set size number of samples
    weighted_sampler = WeightedRandomSampler(
        weights=sample_weights,
154
155
        num_samples=len(train_window_indices),
        replacement=True,
156
157
158
159
160
161
162
163
    )

    # train data loader
    train_dataset = EEGGraphDataset(
        x=x, y=y, num_nodes=num_nodes, indices=train_window_indices
    )

    train_loader = GraphDataLoader(
164
165
        dataset=train_dataset,
        batch_size=_BATCH_SIZE,
166
167
        sampler=weighted_sampler,
        num_workers=NUM_WORKERS,
168
        pin_memory=True,
169
170
171
172
    )

    # this loader is used without weighted sampling, to evaluate metrics on full training set after each epoch
    train_metrics_loader = GraphDataLoader(
173
174
175
176
177
        dataset=train_dataset,
        batch_size=_BATCH_SIZE,
        shuffle=False,
        num_workers=NUM_WORKERS,
        pin_memory=True,
178
179
180
181
182
183
184
185
    )

    # test data loader
    test_dataset = EEGGraphDataset(
        x=x, y=y, num_nodes=num_nodes, indices=heldout_test_window_indices
    )

    test_loader = GraphDataLoader(
186
187
188
189
190
        dataset=test_dataset,
        batch_size=_BATCH_SIZE,
        shuffle=False,
        num_workers=NUM_WORKERS,
        pin_memory=True,
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
    )

    auroc_train_history = []
    auroc_test_history = []
    balACC_train_history = []
    balACC_test_history = []
    loss_train_history = []
    loss_test_history = []

    # training=========================================================================================================
    for epoch in range(_NUM_EPOCHS):
        model.train()
        train_loss = []

        for batch_idx, batch in enumerate(train_loader):
            # send batch to GPU
            g, dataset_idx, y = batch
            g_batch = g.to(device=_DEVICE, non_blocking=True)
            y_batch = y.to(device=_DEVICE, non_blocking=True)
            optimizer.zero_grad()

            # forward pass
            outputs = model(g_batch)
            loss = loss_function(outputs, y_batch)
            train_loss.append(loss.item())

            # backward pass
            loss.backward()
            optimizer.step()

        # update learning rate
        scheduler.step()

224
        # evaluate model after each epoch for train-metric data============================================================
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
        model.eval()
        with torch.no_grad():
            y_probs_train = torch.empty(0, 2).to(_DEVICE)
            y_true_train, y_pred_train = [], []

            for i, batch in enumerate(train_metrics_loader):
                g, dataset_idx, y = batch
                g_batch = g.to(device=_DEVICE, non_blocking=True)
                y_batch = y.to(device=_DEVICE, non_blocking=True)

                # forward pass
                outputs = model(g_batch)

                _, predicted = torch.max(outputs.data, 1)
                y_pred_train += predicted.cpu().numpy().tolist()
                # concatenate along 0th dimension
                y_probs_train = torch.cat((y_probs_train, outputs.data), 0)
                y_true_train += y_batch.cpu().numpy().tolist()

        # returning prob distribution over target classes, take softmax over the 1st dimension
245
246
247
        y_probs_train = (
            nn.functional.softmax(y_probs_train, dim=1).cpu().numpy()
        )
248
249
        y_true_train = np.array(y_true_train)

250
        # evaluate model after each epoch for validation data ==============================================================
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        y_probs_test = torch.empty(0, 2).to(_DEVICE)
        y_true_test, minibatch_loss, y_pred_test = [], [], []

        for i, batch in enumerate(test_loader):
            g, dataset_idx, y = batch
            g_batch = g.to(device=_DEVICE, non_blocking=True)
            y_batch = y.to(device=_DEVICE, non_blocking=True)

            # forward pass
            outputs = model(g_batch)
            _, predicted = torch.max(outputs.data, 1)
            y_pred_test += predicted.cpu().numpy().tolist()

            loss = loss_function(outputs, y_batch)
            minibatch_loss.append(loss.item())
            y_probs_test = torch.cat((y_probs_test, outputs.data), 0)
            y_true_test += y_batch.cpu().numpy().tolist()

        # returning prob distribution over target classes, take softmax over the 1st dimension
270
271
272
        y_probs_test = (
            torch.nn.functional.softmax(y_probs_test, dim=1).cpu().numpy()
        )
273
274
275
        y_true_test = np.array(y_true_test)

        # record training auroc and testing auroc
276
277
278
279
280
281
        auroc_train_history.append(
            roc_auc_score(y_true_train, y_probs_train[:, 1])
        )
        auroc_test_history.append(
            roc_auc_score(y_true_test, y_probs_test[:, 1])
        )
282
283

        # record training balanced accuracy and testing balanced accuracy
284
285
286
287
288
289
        balACC_train_history.append(
            balanced_accuracy_score(y_true_train, y_pred_train)
        )
        balACC_test_history.append(
            balanced_accuracy_score(y_true_test, y_pred_test)
        )
290
291
292
293
294
295

        # LOSS - epoch loss is defined as mean of minibatch losses within epoch
        loss_train_history.append(np.mean(train_loss))
        loss_test_history.append(np.mean(minibatch_loss))

        # print the metrics
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
        print(
            "Train loss: {}, test loss: {}".format(
                loss_train_history[-1], loss_test_history[-1]
            )
        )
        print(
            "Train AUC: {}, test AUC: {}".format(
                auroc_train_history[-1], auroc_test_history[-1]
            )
        )
        print(
            "Train Bal.ACC: {}, test Bal.ACC: {}".format(
                balACC_train_history[-1], balACC_test_history[-1]
            )
        )
311
312
313

        # save model from each epoch====================================================================================
        state = {
314
315
316
317
318
            "epochs": _NUM_EPOCHS,
            "experiment_name": _EXPERIMENT_NAME,
            "model_description": str(model),
            "state_dict": model.state_dict(),
            "optimizer": optimizer.state_dict(),
319
320
        }
        torch.save(state, f"{_EXPERIMENT_NAME}_Epoch_{epoch}.ckpt")