Commit fb9f2b8d authored by Huaqi Qiu's avatar Huaqi Qiu
Browse files

add gpu selection

parent 996de11b
...@@ -193,9 +193,13 @@ if __name__=='__main__': ...@@ -193,9 +193,13 @@ if __name__=='__main__':
parser.add_argument('--save_name_format', default='seg_sa_{}.nii.gz', help='save mask format. {} for ED/ES') parser.add_argument('--save_name_format', default='seg_sa_{}.nii.gz', help='save mask format. {} for ED/ES')
parser.add_argument('--no_resample', default=False, action='store_true',help='not resample image before prediction') parser.add_argument('--no_resample', default=False, action='store_true',help='not resample image before prediction')
parser.add_argument('--batch_size', default=1, help='how many image slices to be sent to predictor each iteration ') parser.add_argument('--batch_size', default=1, help='how many image slices to be sent to predictor each iteration ')
parser.add_argument('--gpu', default=0, help='select GPU by masking shell environment variable CUDA_VISIBLE_DEVICES')
args = parser.parse_args() args = parser.parse_args()
### GPU CONFIG
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
### DATA CONFIG ### DATA CONFIG
root_dir = os.path.abspath(args.root_dir) root_dir = os.path.abspath(args.root_dir)
sequence_name=args.sequence sequence_name=args.sequence
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment