Brain Age Estimation on T2-FLAIR Scans for Application to Multiple Sclerosis.

1. import os
2. import sys
3. import pandas as pd
4. import numpy as np
5. import torch
6. from torch.utils.data import DataLoader
7. import monai
8. from monai.data import NiftiDataset
9. from monai.transforms import AddChannel, Compose, RandFlip, Resize, ScaleIntensity, ToTensor, SqueezeDim, AsChannelFirst, NormalizeIntensity, RandScaleIntensity, RandShiftIntensity, RandAdjustContrast, RandGaussianSmooth, RandHistogramShift, RandGaussianNoise, RandSpatialCrop, CenterSpatialCrop
11. def main():
13. ## load list of image files and the file with the coresponding ages in the same order. must be a text file with no header and a unique image file name on each row with ages in corresponding file in the exact same order.
14. File_list_in = pd.read_csv('traning_imgs_flair.txt',header=None)
15. Age_list_in = pd.read_csv('traning_imgs_ages.txt',header=None)
16. File_list = np.array(File_list_in[0])
17. Age_list = np.array(Age_list_in[0]).astype(float)
20. ## print list lengths to ensure expected number of files
21. print(len(File_list))
22. print(len(Age_list))
24. ## set up image transfomrations and data augmentation in this case the 1mm3 images are resized to 1.45mm isotropic
25. ## various data agmentations used which are outlined in monai documentation. The image is randomly cropped to 3 voxels smaller in each dimension to allow for random shift of voxels within convolutional kernels
26. train_transforms = Compose([AddChannel(), NormalizeIntensity(nonzero=True), Resize((121,145,121),mode="trilinear"), ScaleIntensity(), RandScaleIntensity(factors=0.2, prob=0.25), RandShiftIntensity(offsets=0.2, prob=0.25), RandFlip(prob=0.5,spatial_axis=0), RandAdjustContrast(), RandGaussianSmooth(), RandHistogramShift(), RandGaussianNoise(), RandSpatialCrop([118,142,118],random_size=False), ToTensor()])
28. ## set up validation input trasforms without data augmentation. centre croped used instead of random crop
29. val_transforms = Compose([AddChannel(), NormalizeIntensity(nonzero=True), Resize((121,145,121),mode="trilinear"), ScaleIntensity(), CenterSpatialCrop([118,142,118]), ToTensor()])
31. # Define image dataset data loader and test
32. check_ds = NiftiDataset(image_files=File_list, labels=Age_list, transform=train_transforms)
33. check_loader = DataLoader(check_ds, batch_size=10, num_workers=1, pin_memory=torch.cuda.is_available())
34. im, label = monai.utils.misc.first(check_loader)
36. # create a training data loader. leave a set number of images out from the end of the list out for validation in this case 174 images. batch size set at 10 for a 12Gb GPU.
37. train_ds = NiftiDataset(image_files=File_list[:-174], labels=Age_list[:-174], transform=train_transforms)
38. train_loader = DataLoader(train_ds, batch_size=10, shuffle=True, num_workers=1, pin_memory=torch.cuda.is_available())
40. # create a validation data loader. input only the left out images at the end of the list.
41. val_ds = NiftiDataset(image_files=File_list[-174:], labels=Age_list[-174:], transform=val_transforms)
42. val_loader = DataLoader(val_ds, batch_size=10, num_workers=1, pin_memory=torch.cuda.is_available())
44. # Create DenseNet169 model using monai pre build. dropout of 0.1 added.
45. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
46. model_1 = monai.networks.nets.densenet.densenet169(spatial_dims=3, in_channels=1, out_channels=1, dropout_prob=0.1)
48. ## add a linear regression layer onto the end of the model to directly predict a continous variable, in this case age (without this the output will be limited from -1 to 1).
49. model = torch.nn.Sequential(model_1,torch.nn.Linear(1,1)).to(device)
50. ## define loss function and metric of the model in this case mean square error loss and mean absolute error.
51. loss_function = torch.nn.MSELoss()
52. MAE_metric = torch.nn.L1Loss()
53. ## define the optimiser, ADAM is is used in this case with the commonly used learning rate of 1e-4 with a weight decay of 1e-5
54. optimizer = torch.optim.Adam(model.parameters(), lr=1e-4,weight_decay=1e-5)
55. ## set maximum number of epochs of the model
56. max_epochs = 200
57. ## additional learning rate decay schduling is added for further regulization of the model. Schduling used as recomended by DeepLab and commonly used in the liturature.
58. lambda1 = lambda epoch: (1 - (epoch / max_epochs))**0.9
59. scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)
61. ## define and create text files to be used for output of model traning moitoring. A file for training epoch average and validation performance is defined
62. txt_file_train_av_name = 'densenet_flair_train_epoch_average.txt'
63. txt_file_val_name = 'densenet_flair_validation.txt'
65. txt_file_train = open(txt_file_train_av_name, 'a')
66. txt_file_train.write("densenet_flair")
67. txt_file_train.write("\n")
68. txt_file_train.close()
70. txt_file_train = open(txt_file_val_name, 'a')
71. txt_file_train.write("densenet_flair")
72. txt_file_train.write("\n")
73. txt_file_train.close()
76. # start a typical PyTorch training
77. # early stopping is utilized with a patients of 20 epochs defined here.
78. val_interval = 1
79. best_metric = 100000
80. wait_num = 0
81. patience = 20
82. best_metric_epoch = -1
83. epoch_loss_values = list()
84. metric_values = list()
85. ## start loop to run model for each epoch
86. for epoch in range(max_epochs):
87. print("-" * 10)
88. print(f"epoch {epoch + 1}/{200}")
89. model.train()
90. epoch_loss = 0
91. epoch_metric = 0
92. step = 0
94. for batch_data in train_loader:
95. step += 1
96. ## prepare input tensor
97. inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
98. optimizer.zero_grad()
99. ## run model
100. outputs = torch.squeeze(model(inputs))
101. print(outputs,labels)
103. ## calculate loss and back propagate
104. loss = loss_function(outputs.float(), labels.float())
105. metric = MAE_metric(outputs.float(), labels.float())
106. loss.backward()
107. optimizer.step()
109. ## add batch loss and metric to epoch runing total
110. epoch_loss += loss.item()
111. epoch_metric += metric.item()
112. epoch_len = len(train_ds) // train_loader.batch_size
113. print(f"{step}/{epoch_len}, metric: {metric} train_loss: {loss.item():.4f}")
115. ## calculate mean epoch loss and metric and write to monitoring file
116. epoch_loss /= step
117. epoch_metric /= step
118. epoch_loss_values.append(epoch_loss)
119. print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
120. txt_file_train_av = open(txt_file_train_av_name, 'a')
121. txt_file_train_av.write(f"epoch {epoch + 1} average metric: {epoch_metric} average loss: {epoch_loss:.4f}")
122. txt_file_train_av.write("\n")
123. txt_file_train_av.close()
125. ## step forward scheduler for LR decay
126. scheduler.step()
129. ## evaluate model performance on validation files at end of epoch
130. model.eval()
131. with torch.no_grad():
132. num_correct = 0.0
133. metric_count = 0
134. loss_all = 0
135. for val_data in val_loader:
136. val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
137. val_outputs = torch.squeeze(model(val_images))
138. loss = loss_function(val_outputs.float(), val_labels.float())
139. metric_val = MAE_metric(val_outputs.float(), val_labels.float())
140. metric_count += 1
141. num_correct += metric_val.item()
142. loss_all += loss.item()
143. metric_av = num_correct / metric_count
144. loss_av = loss_all / metric_count
145. print('VAL LOSS',metric_av)
147. ## write validation loss and metric to monitoring file
148. txt_file_val = open(txt_file_val_name, 'a')
149. txt_file_val.write(f"epoch {epoch + 1}, Val_metric: {metric_av} Val_loss: {loss_av:.4f}")
150. txt_file_val.write("\n")
151. txt_file_val.close()
153. metric_values.append(metric)
155. ## early stopping if current validation perfomrance not an improvment will add to wait with patience of 20 to stop model early.
156. if metric_av > best_metric:
157. wait_num += 1
159. if wait_num >= patience:
160. txt_file_val = open(txt_file_val_name, 'a')
161. txt_file_val.write(f"stopped early at epoch {epoch + 1}, best val metric: {best_metric} at epoch {best_metric_epoch}")
162. txt_file_val.write("\n")
163. txt_file_val.close()
164. break
165. else:
166. wait_num = 0
168. #update saved weights if current epoch has the best metric so far
169. if metric_av < best_metric:
170. best_metric = metric_av
171. best_metric_epoch = epoch + 1
172. torch.save(model.state_dict(), "best_model_densenet_flair.pth")
173. print("saved new best metric model")
174. print(
175. "current epoch: {} current accuracy: {:.4f} best accuracy: {:.4f} at epoch {}".format(
176. epoch + 1, metric_av, best_metric, best_metric_epoch
177. )
178. )
179. #writer.add_scalar("val_accuracy", metric, epoch + 1)
180. print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
183. if __name__ == "__main__":
184. main()

densenet_training.py at commit 84822e8, under MIT · at the source