@@ -205,7 +205,8 @@ def update_args(args):
205
205
args .test_add_geo_coords = bool (args .test_add_geo_coords )
206
206
207
207
# set cuda values
208
- if args .gpu >= 0 :
208
+ # if args.gpu >= 0:
209
+ if args .gpu != "-1" :
209
210
args .use_GPU , args .use_CUDNN = 1 , 1
210
211
else :
211
212
args .use_GPU , args .use_CUDNN = 0 , 0
@@ -661,6 +662,7 @@ def yolt_command(framework='yolt2',
661
662
else :
662
663
gpu_cmd = '-i ' + str (gpu )
663
664
# gpu_cmd = '-i ' + str(3-args.gpu) # originally, numbers were reversed
665
+ ngpus = len (gpu .split (',' ))
664
666
665
667
##########################
666
668
# SET VARIABLES ACCORDING TO MODE (SET UNNECCESSARY VALUES TO 0 OR NULL)
@@ -713,6 +715,7 @@ def yolt_command(framework='yolt2',
713
715
str (nbands ),
714
716
yolt_loss_file ,
715
717
str (min_retain_prob ),
718
+ str (ngpus ),
716
719
suffix
717
720
]
718
721
@@ -1732,8 +1735,9 @@ def main():
1732
1735
help = "object detection framework [yolt2, 'yolt3', ssd, faster_rcnn]" )
1733
1736
parser .add_argument ('--mode' , type = str , default = 'test' ,
1734
1737
help = "[compile, test, train, test]" )
1735
- parser .add_argument ('--gpu' , type = int , default = 0 ,
1736
- help = "GPU number, set < 0 to turn off GPU support" )
1738
+ parser .add_argument ('--gpu' , type = str , default = "0" ,
1739
+ help = "GPU number, set < 0 to turn off GPU support " \
1740
+ "to use multiple, use '0,1'" )
1737
1741
parser .add_argument ('--single_gpu_machine' , type = int , default = 0 ,
1738
1742
help = "Switch to use a machine with just one gpu" )
1739
1743
parser .add_argument ('--nbands' , type = int , default = 3 ,
0 commit comments