Documentation for device change

Author: jacobgil

README.md

@@ -121,11 +121,11 @@ input_tensor = # Create an input tensor image for your model..
 # Note: input_tensor can be a batch tensor with several images!
 
 # Construct the CAM object once, and then re-use it on many images:
-cam = GradCAM(model=model, target_layers=target_layers, use_cuda=args.use_cuda)
+cam = GradCAM(model=model, target_layers=target_layers)
 
 # You can also use it within a with statement, to make sure it is freed,
 # In case you need to re-create it inside an outer loop:
-# with GradCAM(model=model, target_layers=target_layers, use_cuda=args.use_cuda) as cam:
+# with GradCAM(model=model, target_layers=target_layers) as cam:
 #   ...
 
 # We have to specify the target we want to generate
@@ -244,8 +244,8 @@ two smoothing methods are supported:
 Usage: `python cam.py --image-path <path_to_image> --method <method> --output-dir <output_dir_path> `
 
 
-To use with CUDA:
-`python cam.py --image-path <path_to_image> --use-cuda  --output-dir <output_dir_path> `
+To use with a specific device, like cpu, cuda, cuda:0 or mps:
+`python cam.py --image-path <path_to_image> --device cuda  --output-dir <output_dir_path> `
 
 ----------
 

cam.py

@@ -4,7 +4,6 @@
 import numpy as np
 import torch
 from torchvision import models
-from torchvision.models import ResNet50_Weights
 from pytorch_grad_cam import (
     GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus,
     AblationCAM, XGradCAM, EigenCAM, EigenGradCAM,
@@ -19,7 +18,7 @@
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--device', type=str, default=None,
+    parser.add_argument('--device', type=str, default='cpu',
                         help='Torch device to use')
     parser.add_argument(
         '--image-path',
@@ -77,7 +76,7 @@ def get_args():
         "gradcamelementwise": GradCAMElementWise
     }
 
-    model = models.resnet50(weights=ResNet50_Weights.DEFAULT).to(args.device).eval()
+    model = models.resnet50(pretrained=True).to(torch.device(args.device)).eval()
 
     # Choose the target layer you want to compute the visualization for.
     # Usually this will be the last convolutional layer in the model.
@@ -104,16 +103,15 @@ def get_args():
     # the Class Activation Maps for.
     # If targets is None, the highest scoring category (for every member in the batch) will be used.
     # You can target specific categories by
-    # targets = [e.g ClassifierOutputTarget(281)]
+    # targets = [ClassifierOutputTarget(281)]
+    # targets = [ClassifierOutputTarget(281)]
     targets = None
 
     # Using the with statement ensures the context is freed, and you can
     # recreate different CAM objects in a loop.
     cam_algorithm = methods[args.method]
     with cam_algorithm(model=model,
-                       target_layers=target_layers,
-                       device=args.device) as cam:
-
+                       target_layers=target_layers) as cam:
 
         # AblationCAM and ScoreCAM have batched implementations.
         # You can override the internal batch size for faster computation.

pytorch_grad_cam/base_cam.py

@@ -18,8 +18,9 @@ def __init__(self,
                  tta_transforms: Optional[tta.Compose] = None) -> None:
         self.model = model.eval()
         self.target_layers = target_layers
-        self.device = next(self.model.parameters()).device
 
+        # Use the same device as the model.
+        self.device = next(self.model.parameters()).device
         self.reshape_transform = reshape_transform
         self.compute_input_gradient = compute_input_gradient
         self.uses_gradients = uses_gradients

pytorch_grad_cam/score_cam.py

@@ -24,7 +24,7 @@ def get_cam_weights(self,
             upsample = torch.nn.UpsamplingBilinear2d(
                 size=input_tensor.shape[-2:])
             activation_tensor = torch.from_numpy(activations)
-            activation_tensor = activation_tensor.to(next(self.model.parameters()).device)
+            activation_tensor = activation_tensor.to(self.device)
 
             upsampled = upsample(activation_tensor)
 

setup.py

@@ -1,29 +1,29 @@
-import setuptools
-
-with open('README.md', mode='r', encoding='utf-8') as fh:
-    long_description = fh.read()
-
-with open("requirements.txt", "r") as f:
-    requirements = f.readlines()
-
-setuptools.setup(
-    name='grad-cam',
-    version='1.4.8',
-    author='Jacob Gildenblat',
-    author_email='[email protected]',
-    description='Many Class Activation Map methods implemented in Pytorch for classification, segmentation, object detection and more',
-    long_description=long_description,
-    long_description_content_type='text/markdown',
-    url='https://github.com/jacobgil/pytorch-grad-cam',
-    project_urls={
-        'Bug Tracker': 'https://github.com/jacobgil/pytorch-grad-cam/issues',
-    },
-    classifiers=[
-        'Programming Language :: Python :: 3',
-        'License :: OSI Approved :: MIT License',
-        'Operating System :: OS Independent',
-    ],
-    packages=setuptools.find_packages(
-        exclude=["*tutorials*"]),
-    python_requires='>=3.6',
-    install_requires=requirements)
+import setuptools
+
+with open('README.md', mode='r', encoding='utf-8') as fh:
+    long_description = fh.read()
+
+with open("requirements.txt", "r") as f:
+    requirements = f.readlines()
+
+setuptools.setup(
+    name='grad-cam',
+    version='1.5.0',
+    author='Jacob Gildenblat',
+    author_email='[email protected]',
+    description='Many Class Activation Map methods implemented in Pytorch for classification, segmentation, object detection and more',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/jacobgil/pytorch-grad-cam',
+    project_urls={
+        'Bug Tracker': 'https://github.com/jacobgil/pytorch-grad-cam/issues',
+    },
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: OS Independent',
+    ],
+    packages=setuptools.find_packages(
+        exclude=["*tutorials*"]),
+    python_requires='>=3.8',
+    install_requires=requirements)