-
Notifications
You must be signed in to change notification settings - Fork 0
/
main_executor.py
350 lines (319 loc) · 15.3 KB
/
main_executor.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
#!/usr/bin/env python
"""
Purpose: Entry Point for the Vessel Segmentation Solution
"""
import os
import argparse
import random
import numpy as np
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter
from Utils.logger import Logger
from Utils.model_manager import get_model
from Utils.vessel_utils import load_model_huggingface
from pipeline import Pipeline
from cross_validation_pipeline import CrossValidationPipeline
__author__ = "Kartik Prabhu, Mahantesh Pattadkal, and Soumick Chatterjee"
__copyright__ = "Copyright 2020, Faculty of Computer Science, Otto von Guericke University Magdeburg, Germany"
__credits__ = ["Kartik Prabhu", "Mahantesh Pattadkal", "Soumick Chatterjee"]
__license__ = "GPL"
__version__ = "1.0.0"
__maintainer__ = "Soumick Chatterjee"
__email__ = "[email protected]"
__status__ = "Production"
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(2020)
np.random.seed(2020)
random.seed(2020)
# torch.autograd.set_detect_anomaly(True)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-model",
type=int,
default=2,
help="1{U-Net}; \n"
"2{U-Net_Deepsup}; \n"
"3{Attention-U-Net}; \n"
"4{Probabilistic-U-Net};")
parser.add_argument("-model_name",
default="Model_MIP_v1",
help="Name of the model")
parser.add_argument("-dataset_path",
default="",
help="Path to folder containing dataset."
"If cross validation is to be performed, please create the dataset folder consisting of"
"train and train_label folders along with test and test_label folders"
"e.g."
"/sample_dataset"
" /train"
" /train_label"
" /test"
" /test_label"
"else if training is to be performed, prepare the dataset folder consisting of"
"train, train_label, validate, validate_label, test and test_label folders"
"e.g.,"
"/sample_dataset"
" /train"
" /train_label"
" /validate"
" /validate_label"
" /test"
" /test_label"
"each folder must contain at least one 3D MRA volume in nifti .nii or nii.gz formats"
"else if inference is to be performed, specify folder consisting of"
"3D MRA volume in nifti .nii or nii.gz formats"
"Example: /home/dataset/")
parser.add_argument("-output_path",
default="",
help="Folder path to store output "
"Example: /home/output/")
parser.add_argument('-train',
default=False,
help="To train the model")
parser.add_argument('-test',
default=False,
help="To test the model")
parser.add_argument('-eval',
default=False,
help="To render inference of specified nifti 3D volumes")
parser.add_argument('-cross_validate',
default=False,
help="To train with k-fold cross validation")
parser.add_argument('-predict',
default=False,
help="To predict a segmentation output of the model and to get a diff between label and output")
parser.add_argument('-predictor_path',
default="",
help="Path to the input image to predict an output, ex:/home/test/ww25.nii ")
parser.add_argument('-predictor_label_path',
default="",
help="Path to the label image to find the diff between label an output, "
"e.g.,:/home/test/ww25_label.nii ")
parser.add_argument('-load_huggingface',
default="",
help="Load model from huggingface model hub ex: 'soumickmj/DS6_UNetMSS3D_wDeform' [model param will be ignored]")
parser.add_argument('-load_path',
default="",
help="Path to checkpoint of existing model to load, ex:/home/model/checkpoint/ [If this is supplied, load_huggingface will be ignored]")
parser.add_argument('-load_best',
default=True,
help="Specifiy whether to load the best checkpoint or the last. "
"Also to be used if Train and Test both are true. [Only if load_path is supplied]")
parser.add_argument('-pre_train',
default=True,
help="Specifiy whether to load the pre trained weights for training/tuning.")
parser.add_argument('-deform',
default=False,
help="To use deformation for training")
parser.add_argument('-clip_grads',
default=True,
help="To use deformation for training")
parser.add_argument('-apex',
default=True,
help="To use half precision on model weights.")
parser.add_argument('-with_mip',
default=True,
help="Train with MIP Loss")
parser.add_argument('-use_madam',
default=False,
help="Set this to True to use madam optimizer")
parser.add_argument("-batch_size",
type=int,
default=15,
help="Batch size for training")
parser.add_argument("-num_epochs",
type=int,
default=50,
help="Number of epochs for training")
parser.add_argument("-learning_rate",
type=float,
default=0.0001,
help="Learning rate")
parser.add_argument("-patch_size",
type=int,
default=64,
help="Patch size of the input volume")
parser.add_argument("-stride_depth",
type=int,
default=16,
help="Strides for dividing the input volume into patches in depth dimension "
"(To be used during validation and inference)")
parser.add_argument("-stride_width",
type=int,
default=32,
help="Strides for dividing the input volume into patches in width dimension "
"(To be used during validation and inference)")
parser.add_argument("-stride_length",
type=int,
default=32,
help="Strides for dividing the input volume into patches in length dimension "
"(To be used during validation and inference)")
parser.add_argument("-samples_per_epoch",
type=int,
default=8000,
help="Number of samples per epoch")
parser.add_argument("-num_worker",
type=int,
default=8,
help="Number of worker threads")
parser.add_argument("-floss_coeff",
type=float,
default=0.7,
help="Loss coefficient for floss in total loss")
parser.add_argument("-mip_loss_coeff",
type=float,
default=0.3,
help="Loss coefficient for mip_loss in total loss")
parser.add_argument("-floss_param_smooth",
type=float,
default=1,
help="Loss coefficient for floss_param_smooth")
parser.add_argument("-floss_param_gamma",
type=float,
default=0.75,
help="Loss coefficient for floss_param_gamma")
parser.add_argument("-floss_param_alpha",
type=float,
default=0.7,
help="Loss coefficient for floss_param_alpha")
parser.add_argument("-mip_loss_param_smooth",
type=float,
default=1,
help="Loss coefficient for mip_loss_param_smooth")
parser.add_argument("-mip_loss_param_gamma",
type=float,
default=0.75,
help="Loss coefficient for mip_loss_param_gamma")
parser.add_argument("-mip_loss_param_alpha",
type=float,
default=0.7,
help="Loss coefficient for mip_loss_param_alpha")
parser.add_argument("-mip_axis",
type=str,
default="z",
help="Set projection axis. Default is z-axis. use axis in [x, y, z] or 'multi'")
parser.add_argument("-k_folds",
type=int,
default=5,
help="Set the number of folds for cross validation")
parser.add_argument("-fold_index",
type=str,
default="",
help="Set the number of folds for cross validation")
parser.add_argument("-wandb",
default=False,
help="Set this to true to include wandb logging")
parser.add_argument("-wandb_project",
type=str,
default="",
help="Set this to wandb project name e.g., 'DS6_VesselSeg2'")
parser.add_argument("-wandb_entity",
type=str,
default="",
help="Set this to wandb project name e.g., 'ds6_vessel_seg2'")
parser.add_argument("-wandb_api_key",
type=str,
default="",
help="API Key to login that can be found at https://wandb.ai/authorize")
args = parser.parse_args()
if str(args.deform).lower() == "true":
args.model_name += "_Deform"
MODEL_NAME = args.model_name
DATASET_FOLDER = args.dataset_path
OUTPUT_PATH = args.output_path
os.makedirs(OUTPUT_PATH, exist_ok=True)
LOAD_PATH = args.load_path
CHECKPOINT_PATH = OUTPUT_PATH + "/" + MODEL_NAME + '/checkpoint/'
if str(args.eval).lower() == "true":
TENSORBOARD_PATH_TRAINING = None
TENSORBOARD_PATH_VALIDATION = None
TENSORBOARD_PATH_TESTING = None
LOGGER_PATH = None
logger = None
test_logger = None
writer_training = None
writer_validating = None
else:
TENSORBOARD_PATH_TRAINING = OUTPUT_PATH + "/" + MODEL_NAME + '/tensorboard/tensorboard_training/'
TENSORBOARD_PATH_VALIDATION = OUTPUT_PATH + "/" + MODEL_NAME + '/tensorboard/tensorboard_validation/'
TENSORBOARD_PATH_TESTING = OUTPUT_PATH + "/" + MODEL_NAME + '/tensorboard/tensorboard_testing/'
LOGGER_PATH = OUTPUT_PATH + "/" + MODEL_NAME + '.log'
logger = Logger(MODEL_NAME, LOGGER_PATH).get_logger()
test_logger = Logger(MODEL_NAME + '_test', LOGGER_PATH).get_logger()
writer_training = SummaryWriter(TENSORBOARD_PATH_TRAINING)
writer_validating = SummaryWriter(TENSORBOARD_PATH_VALIDATION)
# Model
if args.load_huggingface:
model = load_model_huggingface(args.load_huggingface)
else:
model = get_model(args.model)
model = torch.nn.DataParallel(model)
model.cuda()
wandb = None
if str(args.wandb).lower() == "true":
import wandb
wandb.login(key=args.wandb_api_key)
wandb.init(project=args.wandb_project, entity=args.wandb_entity, group=args.model_name, notes=args.model_name)
wandb.config = {
"learning_rate": args.learning_rate,
"epochs": args.num_epochs,
"batch_size": args.batch_size,
"patch_size": args.patch_size,
"samples_per_epoch": args.samples_per_epoch,
"mip_loss_coeff": args.mip_loss_coeff,
"floss_coeff": args.floss_coeff
}
# Choose pipeline based on whether or not to perform cross validation
# If cross validation is to be performed, please create the dataset folder consisting of
# train and train_label folders along with test and test_label folders
# e.g.,
# /sample_dataset
# /train
# /train_label
# /test
# /test_label
# Otherwise prepare the dataset folder consisting of
# train, train_label, validate, validate_label, test and test_label folders
# e.g.,
# /sample_dataset
# /train
# /train_label
# /validate
# /validate_label
# /test
# /test_label
# Each folder must contain at least one 3D MRA volume in nifti .nii or nii.gz formats
if str(args.cross_validate).lower() == "true":
pipeline = CrossValidationPipeline(cmd_args=args, model=model, logger=logger,
dir_path=DATASET_FOLDER, checkpoint_path=CHECKPOINT_PATH,
writer_training=writer_training, writer_validating=writer_validating,
test_logger=test_logger, wandb=wandb)
else:
pipeline = Pipeline(cmd_args=args, model=model, logger=logger,
dir_path=DATASET_FOLDER, checkpoint_path=CHECKPOINT_PATH,
writer_training=writer_training, writer_validating=writer_validating, wandb=wandb)
# loading existing checkpoint if supplied
if str(args.pre_train).lower() == "true":
if bool(LOAD_PATH):
pipeline.load(checkpoint_path=LOAD_PATH, load_best=args.load_best, fold_index=args.fold_index)
else:
print("-pre_train is set to True, but -load_path has not been supplied. pre_train loading will be skipped!")
if str(args.train).lower() == "true":
pipeline.train()
torch.cuda.empty_cache() # to avoid memory errors
if str(args.test).lower() == "true":
if str(args.load_best).lower() == "true":
if bool(LOAD_PATH):
pipeline.load(checkpoint_path=LOAD_PATH, load_best=args.load_best, fold_index=args.fold_index)
else:
pipeline.load(load_best=args.load_best)
pipeline.test(test_logger=test_logger, fold_index=args.fold_index)
torch.cuda.empty_cache() # to avoid memory errors
if str(args.eval).lower() == "true":
pipeline.eval(model_name=args.model_name)
if str(args.predict).lower() == "true":
pipeline.predict(args.predictor_path, args.predictor_label_path, predict_logger=test_logger)
if str(args.eval).lower() != "true":
writer_training.close()
writer_validating.close()