refactor to use torch device instead of cuda boolean

projects/DML/dml.py

@@ -22,11 +22,11 @@ def __init__(
         weight_t=0.5,
         weight_x=1.,
         weight_ml=1.,
-        use_gpu=True,
+        device='cuda',
         label_smooth=True,
         deploy='model1'
     ):
-        super(ImageDMLEngine, self).__init__(datamanager, use_gpu)
+        super(ImageDMLEngine, self).__init__(datamanager, device)
 
         self.model1 = model1
         self.optimizer1 = optimizer1
@@ -48,16 +48,15 @@ def __init__(
         self.criterion_t = TripletLoss(margin=margin)
         self.criterion_x = CrossEntropyLoss(
             num_classes=self.datamanager.num_train_pids,
-            use_gpu=self.use_gpu,
+            device=self.device,
             label_smooth=label_smooth
         )
 
     def forward_backward(self, data):
         imgs, pids = self.parse_data_for_train(data)
 
-        if self.use_gpu:
-            imgs = imgs.cuda()
-            pids = pids.cuda()
+        imgs = imgs.to(self.device)
+        pids = pids.to(self.device)
 
         outputs1, features1 = self.model1(imgs)
         loss1_x = self.compute_loss(self.criterion_x, outputs1, pids)

projects/DML/main.py

@@ -92,7 +92,7 @@ def main():
         num_classes=datamanager.num_train_pids,
         loss=cfg.loss.name,
         pretrained=cfg.model.pretrained,
-        use_gpu=cfg.use_gpu
+        device='cuda' if cfg.use_gpu else 'cpu'
     )
     num_params, flops = compute_model_complexity(
         model1, (1, 3, cfg.data.height, cfg.data.width)
@@ -155,7 +155,7 @@ def main():
         weight_t=cfg.loss.triplet.weight_t,
         weight_x=cfg.loss.triplet.weight_x,
         weight_ml=cfg.loss.dml.weight_ml,
-        use_gpu=cfg.use_gpu,
+        device='cuda' if cfg.use_gpu else 'cpu',
         label_smooth=cfg.loss.softmax.label_smooth,
         deploy=cfg.model.deploy
     )

projects/OSNet_AIN/main.py

@@ -123,7 +123,7 @@ def main():
         model,
         optimizer,
         scheduler=scheduler,
-        use_gpu=cfg.use_gpu,
+        device='cuda' if cfg.use_gpu else 'cpu',
         label_smooth=cfg.loss.softmax.label_smooth,
         mc_iter=cfg.nas.mc_iter,
         init_lmda=cfg.nas.init_lmda,

projects/OSNet_AIN/softmax_nas.py

@@ -13,7 +13,7 @@ def __init__(
         model,
         optimizer,
         scheduler=None,
-        use_gpu=False,
+        device='cuda',
         label_smooth=True,
         mc_iter=1,
         init_lmda=1.,
@@ -22,7 +22,7 @@ def __init__(
         lmda_decay_rate=0.5,
         fixed_lmda=False
     ):
-        super(ImageSoftmaxNASEngine, self).__init__(datamanager, use_gpu)
+        super(ImageSoftmaxNASEngine, self).__init__(datamanager, device)
         self.mc_iter = mc_iter
         self.init_lmda = init_lmda
         self.min_lmda = min_lmda
@@ -37,16 +37,15 @@ def __init__(
 
         self.criterion = CrossEntropyLoss(
             num_classes=self.datamanager.num_train_pids,
-            use_gpu=self.use_gpu,
+            device=self.device,
             label_smooth=label_smooth
         )
 
     def forward_backward(self, data):
         imgs, pids = self.parse_data_for_train(data)
 
-        if self.use_gpu:
-            imgs = imgs.cuda()
-            pids = pids.cuda()
+        imgs = imgs.to(self.device)
+        pids = pids.to(self.device)
 
         # softmax temporature
         if self.fixed_lmda or self.lmda_decay_step == -1:

projects/attribute_recognition/main.py

@@ -152,7 +152,7 @@ def main():
         args.arch,
         num_attrs,
         pretrained=not args.no_pretrained,
-        use_gpu=use_gpu
+        device='cuda' if use_gpu else 'cpu'
     )
     num_params, flops = compute_model_complexity(
         model, (1, 3, args.height, args.width)

projects/attribute_recognition/models/__init__.py

@@ -8,10 +8,10 @@
 }
 
 
-def build_model(name, num_classes, pretrained=True, use_gpu=True):
+def build_model(name, num_classes, pretrained=True, device='cuda'):
     avai_models = list(__model_factory.keys())
     if name not in avai_models:
         raise KeyError
     return __model_factory[name](
-        num_classes=num_classes, pretrained=pretrained, use_gpu=use_gpu
+        num_classes=num_classes, pretrained=pretrained, use_gpu=device.startswith('cuda')
     )

scripts/default_config.py

@@ -120,7 +120,7 @@ def imagedata_kwargs(cfg):
         'k_tfm': cfg.data.k_tfm,
         'norm_mean': cfg.data.norm_mean,
         'norm_std': cfg.data.norm_std,
-        'use_gpu': cfg.use_gpu,
+        'device': 'cuda' if cfg.use_gpu else 'cpu',
         'split_id': cfg.data.split_id,
         'combineall': cfg.data.combineall,
         'load_train_targets': cfg.data.load_train_targets,
@@ -149,7 +149,7 @@ def videodata_kwargs(cfg):
         'transforms': cfg.data.transforms,
         'norm_mean': cfg.data.norm_mean,
         'norm_std': cfg.data.norm_std,
-        'use_gpu': cfg.use_gpu,
+        'device': 'cuda' if cfg.use_gpu else 'cpu',
         'split_id': cfg.data.split_id,
         'combineall': cfg.data.combineall,
         'batch_size_train': cfg.train.batch_size,

scripts/main.py

@@ -32,7 +32,7 @@ def build_engine(cfg, datamanager, model, optimizer, scheduler):
                 model,
                 optimizer=optimizer,
                 scheduler=scheduler,
-                use_gpu=cfg.use_gpu,
+                device='cuda' if cfg.use_gpu else 'cpu',
                 label_smooth=cfg.loss.softmax.label_smooth
             )
 
@@ -45,7 +45,7 @@ def build_engine(cfg, datamanager, model, optimizer, scheduler):
                 weight_t=cfg.loss.triplet.weight_t,
                 weight_x=cfg.loss.triplet.weight_x,
                 scheduler=scheduler,
-                use_gpu=cfg.use_gpu,
+                device='cuda' if cfg.use_gpu else 'cpu',
                 label_smooth=cfg.loss.softmax.label_smooth
             )
 
@@ -56,7 +56,7 @@ def build_engine(cfg, datamanager, model, optimizer, scheduler):
                 model,
                 optimizer=optimizer,
                 scheduler=scheduler,
-                use_gpu=cfg.use_gpu,
+                device='cuda' if cfg.use_gpu else 'cpu',
                 label_smooth=cfg.loss.softmax.label_smooth,
                 pooling_method=cfg.video.pooling_method
             )
@@ -70,7 +70,7 @@ def build_engine(cfg, datamanager, model, optimizer, scheduler):
                 weight_t=cfg.loss.triplet.weight_t,
                 weight_x=cfg.loss.triplet.weight_x,
                 scheduler=scheduler,
-                use_gpu=cfg.use_gpu,
+                device='cuda' if cfg.use_gpu else 'cpu',
                 label_smooth=cfg.loss.softmax.label_smooth
             )
 
@@ -157,7 +157,7 @@ def main():
         num_classes=datamanager.num_train_pids,
         loss=cfg.loss.name,
         pretrained=cfg.model.pretrained,
-        use_gpu=cfg.use_gpu
+        device='cuda' if cfg.use_gpu else 'cpu'
     )
     num_params, flops = compute_model_complexity(
         model, (1, 3, cfg.data.height, cfg.data.width)

tools/visualize_actmap.py

@@ -155,7 +155,7 @@ def main():
     model = torchreid.models.build_model(
         name=args.model,
         num_classes=datamanager.num_train_pids,
-        use_gpu=use_gpu
+        device='cuda' if use_gpu else 'cpu'
     )
 
     if use_gpu:

torchreid/data/datamanager.py

@@ -19,7 +19,7 @@ class DataManager(object):
             Default is 'random_flip'.
         norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).
         norm_std (list or None, optional): data std. Default is None (use imagenet std).
-        use_gpu (bool, optional): use gpu. Default is True.
+        device (str, optional): device to use. Default is 'cuda'.
     """
 
     def __init__(
@@ -31,7 +31,7 @@ def __init__(
         transforms='random_flip',
         norm_mean=None,
         norm_std=None,
-        use_gpu=False
+        device='cuda'
     ):
         self.sources = sources
         self.targets = targets
@@ -58,7 +58,8 @@ def __init__(
             norm_std=norm_std
         )
 
-        self.use_gpu = (torch.cuda.is_available() and use_gpu)
+        self.device = device
+        self.use_gpu = (torch.cuda.is_available() and device.startswith('cuda'))
 
     @property
     def num_train_pids(self):
@@ -104,7 +105,7 @@ class ImageDataManager(DataManager):
             useful for training and is currently valid for image datasets only.
         norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).
         norm_std (list or None, optional): data std. Default is None (use imagenet std).
-        use_gpu (bool, optional): use gpu. Default is True.
+        device (str, optional): device to use. Default is 'cuda'.
         split_id (int, optional): split id (*0-based*). Default is 0.
         combineall (bool, optional): combine train, query and gallery in a dataset for
             training. Default is False.
@@ -161,7 +162,7 @@ def __init__(
         k_tfm=1,
         norm_mean=None,
         norm_std=None,
-        use_gpu=True,
+        device='cuda',
         split_id=0,
         combineall=False,
         load_train_targets=False,
@@ -186,7 +187,7 @@ def __init__(
             transforms=transforms,
             norm_mean=norm_mean,
             norm_std=norm_std,
-            use_gpu=use_gpu
+            device=device
         )
 
         print('=> Loading train (source) dataset')
@@ -361,7 +362,7 @@ class VideoDataManager(DataManager):
             Default is 'random_flip'.
         norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).
         norm_std (list or None, optional): data std. Default is None (use imagenet std).
-        use_gpu (bool, optional): use gpu. Default is True.
+        device (str, optional): device to use. Default is 'cuda'.
         split_id (int, optional): split id (*0-based*). Default is 0.
         combineall (bool, optional): combine train, query and gallery in a dataset for
             training. Default is False.
@@ -418,7 +419,7 @@ def __init__(
         transforms='random_flip',
         norm_mean=None,
         norm_std=None,
-        use_gpu=True,
+        device='cuda',
         split_id=0,
         combineall=False,
         batch_size_train=3,
@@ -440,7 +441,7 @@ def __init__(
             transforms=transforms,
             norm_mean=norm_mean,
             norm_std=norm_std,
-            use_gpu=use_gpu
+            device=device
         )
 
         print('=> Loading train (source) dataset')