From 51ae19245f655cf0ee334db2a945ceb1a4d6df59 Mon Sep 17 00:00:00 2001 From: Jacob Gildenblat Date: Tue, 19 Dec 2023 15:15:16 +0200 Subject: [PATCH] Documentation for device change --- README.md | 8 ++--- cam.py | 12 +++----- pytorch_grad_cam/base_cam.py | 3 +- pytorch_grad_cam/score_cam.py | 2 +- setup.py | 58 +++++++++++++++++------------------ 5 files changed, 41 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index da39e5f9..6a964524 100644 --- a/README.md +++ b/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 --method --output-dir ` -To use with CUDA: -`python cam.py --image-path --use-cuda --output-dir ` +To use with a specific device, like cpu, cuda, cuda:0 or mps: +`python cam.py --image-path --device cuda --output-dir ` ---------- diff --git a/cam.py b/cam.py index 12962490..c8b81274 100644 --- a/cam.py +++ b/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. diff --git a/pytorch_grad_cam/base_cam.py b/pytorch_grad_cam/base_cam.py index 81a38da1..f08cd12b 100644 --- a/pytorch_grad_cam/base_cam.py +++ b/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 diff --git a/pytorch_grad_cam/score_cam.py b/pytorch_grad_cam/score_cam.py index 275509ad..3aa055fc 100644 --- a/pytorch_grad_cam/score_cam.py +++ b/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) diff --git a/setup.py b/setup.py index 516fb26b..d3f0c652 100644 --- a/setup.py +++ b/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='jacob.gildenblat@gmail.com', - 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='jacob.gildenblat@gmail.com', + 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)