added strexp

.gitignore

@@ -0,0 +1,224 @@
+**/saved_models/*
+**/data_lmdb_release/*
+**/image_release/*
+**/vitstr_base_patch*
+**/result/*
+**/results/*
+**/oldData/
+*.mdb
+*.xlsx
+*.pth
+*.json
+*.pkl
+*.tar
+*.ipynb
+*.zip
+*.eps
+*.pdf
+**/grcnn_straug/*
+**/augmentation/results/*
+**/tmp/*
+*.sh
+**/__pycache__
+workdir/
+.remote-sync.json
+*.png
+pretrained/
+attributionImgs/
+attributionImgsOld/
+attrSelectivityOld/
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### OSX ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+### Python Patch ###
+.venv/
+
+### Python.VirtualEnv Stack ###
+# Virtualenv
+# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
+[Bb]in
+[Ii]nclude
+[Ll]ib
+[Ll]ib64
+[Ll]ocal
+[Ss]cripts
+pyvenv.cfg
+pip-selfcheck.json
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk

attribution_ops.py

@@ -0,0 +1,87 @@
+import os
+import pickle
+from captum_improve_vitstr import rankedAttributionsBySegm
+import matplotlib.pyplot as plt
+from skimage.color import gray2rgb
+from captum.attr._utils.visualization import visualize_image_attr
+import torch
+import numpy as np
+
+def attr_one_dataset():
+    modelName = "vitstr"
+    datasetName = "IIIT5k_3000"
+
+    rootDir = f"/data/goo/strattr/attributionData/{modelName}/{datasetName}/"
+    attrOutputImgs = f"/data/goo/strattr/attributionDataImgs/{modelName}/{datasetName}/"
+    if not os.path.exists(attrOutputImgs):
+        os.makedirs(attrOutputImgs)
+
+    minNumber = 1000000
+    maxNumber = 0
+    # From a folder containing saved attribution pickle files, convert them into attribution images
+    for path, subdirs, files in os.walk(rootDir):
+        for name in files:
+            fullfilename = os.path.join(rootDir, name) # Value
+            # fullfilename: /data/goo/strattr/attributionData/trba/CUTE80/66_featablt.pkl
+            partfilename = fullfilename[fullfilename.rfind('/')+1:]
+            print("fullfilename: ", fullfilename)
+            imgNum = int(partfilename.split('_')[0])
+            attrImgName = partfilename.replace('.pkl', '.png')
+            minNumber = min(minNumber, imgNum)
+            maxNumber = max(maxNumber, imgNum)
+            with open(fullfilename, 'rb') as f:
+                pklData = pickle.load(f)
+                attributions = pklData['attribution']
+                segmDataNP = pklData['segmData']
+                origImgNP = pklData['origImg']
+            if np.isnan(attributions).any():
+                continue
+            attributions = torch.from_numpy(attributions)
+            rankedAttr = rankedAttributionsBySegm(attributions, segmDataNP)
+            rankedAttr = rankedAttr.detach().cpu().numpy()[0][0]
+            rankedAttr = gray2rgb(rankedAttr)
+            mplotfig, _ = visualize_image_attr(rankedAttr, origImgNP, method='blended_heat_map', cmap='RdYlGn')
+            mplotfig.savefig(attrOutputImgs + attrImgName)
+            mplotfig.clear()
+            plt.close(mplotfig)
+
+def attr_all_dataset():
+    modelName = "vitstr"
+
+    datasetNameList = ['IIIT5k_3000', 'SVT', 'IC03_860', 'IC03_867', 'IC13_857', 'IC13_1015', 'IC15_1811', 'IC15_2077', 'SVTP', 'CUTE80']
+
+    for datasetName in datasetNameList:
+        rootDir = f"/data/goo/strattr/attributionData/{modelName}/{datasetName}/"
+        attrOutputImgs = f"/data/goo/strattr/attributionDataImgs/{modelName}/{datasetName}/"
+        if not os.path.exists(attrOutputImgs):
+            os.makedirs(attrOutputImgs)
+
+        minNumber = 1000000
+        maxNumber = 0
+        # From a folder containing saved attribution pickle files, convert them into attribution images
+        for path, subdirs, files in os.walk(rootDir):
+            for name in files:
+                fullfilename = os.path.join(rootDir, name) # Value
+                # fullfilename: /data/goo/strattr/attributionData/trba/CUTE80/66_featablt.pkl
+                partfilename = fullfilename[fullfilename.rfind('/')+1:]
+                imgNum = int(partfilename.split('_')[0])
+                attrImgName = partfilename.replace('.pkl', '.png')
+                minNumber = min(minNumber, imgNum)
+                maxNumber = max(maxNumber, imgNum)
+                with open(fullfilename, 'rb') as f:
+                    pklData = pickle.load(f)
+                    attributions = pklData['attribution']
+                    segmDataNP = pklData['segmData']
+                    origImgNP = pklData['origImg']
+                attributions = torch.from_numpy(attributions)
+                rankedAttr = rankedAttributionsBySegm(attributions, segmDataNP)
+                rankedAttr = rankedAttr.detach().cpu().numpy()[0][0]
+                rankedAttr = gray2rgb(rankedAttr)
+                mplotfig, _ = visualize_image_attr(rankedAttr, origImgNP, method='blended_heat_map', cmap='RdYlGn')
+                mplotfig.savefig(attrOutputImgs + attrImgName)
+                mplotfig.clear()
+                plt.close(mplotfig)
+
+if __name__ == '__main__':
+    attr_one_dataset()
+    # attr_all_dataset()

augmentation/blur.py

@@ -0,0 +1,189 @@
+
+import cv2
+import numpy as np
+from PIL import Image, ImageOps
+import torchvision.transforms as transforms
+from wand.image import Image as WandImage
+from scipy.ndimage import zoom as scizoom
+from skimage.filters import gaussian
+from wand.api import library as wandlibrary
+from io import BytesIO
+
+#from skimage import color
+from .ops import MotionImage, clipped_zoom, disk, plasma_fractal
+'''
+    PIL resize (W,H)
+'''
+class GaussianBlur:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #kernel = [(31,31)] prev 1 level only
+        kernel = (31, 31)
+        sigmas = [.5, 1, 2]
+        if mag<0 or mag>=len(kernel):
+            index = np.random.randint(0, len(sigmas))
+        else:
+            index = mag
+
+        sigma = sigmas[index]
+        return transforms.GaussianBlur(kernel_size=kernel, sigma=sigma)(img)
+
+
+class DefocusBlur:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+        #c = [(3, 0.1), (4, 0.5), (6, 0.5), (8, 0.5), (10, 0.5)]
+        c = [(2, 0.1), (3, 0.1), (4, 0.1)] #, (6, 0.5)] #prev 2 levels only
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        img = np.array(img) / 255.
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+            n_channels = 3
+        kernel = disk(radius=c[0], alias_blur=c[1])
+
+        channels = []
+        for d in range(n_channels):
+            channels.append(cv2.filter2D(img[:, :, d], -1, kernel))
+        channels = np.array(channels).transpose((1, 2, 0))  # 3x224x224 -> 224x224x3
+
+        #if isgray:
+        #    img = img[:,:,0]
+        #    img = np.squeeze(img)
+
+        img = np.clip(channels, 0, 1) * 255
+        img = Image.fromarray(img.astype(np.uint8))
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img
+
+
+class MotionBlur:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+        #c = [(10, 3), (15, 5), (15, 8), (15, 12), (20, 15)]
+        c = [(10, 3), (12, 4), (14, 5)]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        output = BytesIO()
+        img.save(output, format='PNG')
+        img = MotionImage(blob=output.getvalue())
+
+        img.motion_blur(radius=c[0], sigma=c[1], angle=np.random.uniform(-45, 45))
+        img = cv2.imdecode(np.fromstring(img.make_blob(), np.uint8), cv2.IMREAD_UNCHANGED)
+        if len(img.shape) > 2:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        img = Image.fromarray(img.astype(np.uint8))
+
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img
+
+class GlassBlur:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #c = [(0.7, 1, 2), (0.9, 2, 1), (1, 2, 3), (1.1, 3, 2), (1.5, 4, 2)][severity - 1]
+        c = [(0.7, 1, 2), (0.75, 1, 2), (0.8, 1, 2)] #, (1, 2, 3)] #prev 2 levels only
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+
+        c = c[index]
+
+        img = np.uint8(gaussian(np.array(img) / 255., sigma=c[0], multichannel=True) * 255)
+
+        # locally shuffle pixels
+        for i in range(c[2]):
+            for h in range(H - c[1], c[1], -1):
+                for w in range(W - c[1], c[1], -1):
+                    dx, dy = np.random.randint(-c[1], c[1], size=(2,))
+                    h_prime, w_prime = h + dy, w + dx
+                    # swap
+                    img[h, w], img[h_prime, w_prime] = img[h_prime, w_prime], img[h, w]
+
+        img = np.clip(gaussian(img / 255., sigma=c[0], multichannel=True), 0, 1) * 255
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class ZoomBlur:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        c = [np.arange(1, 1.11, .01),
+             np.arange(1, 1.16, .01),
+             np.arange(1, 1.21, .02)]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        uint8_img = img
+        img = (np.array(img) / 255.).astype(np.float32)
+
+        out = np.zeros_like(img)
+        for zoom_factor in c:
+            ZW = int(W*zoom_factor)
+            ZH = int(H*zoom_factor)
+            zoom_img = uint8_img.resize((ZW, ZH), Image.BICUBIC)
+            x1 = (ZW - W) // 2
+            y1 = (ZH - H) // 2
+            x2 = x1 + W
+            y2 = y1 + H
+            zoom_img = zoom_img.crop((x1,y1,x2,y2))
+            out += (np.array(zoom_img) / 255.).astype(np.float32)
+
+        img = (img + out) / (len(c) + 1)
+
+        img = np.clip(img, 0, 1) * 255
+        img = Image.fromarray(img.astype(np.uint8))
+
+        return img

augmentation/camera.py

@@ -0,0 +1,120 @@
+
+import cv2
+import numpy as np
+import skimage as sk
+from PIL import Image, ImageOps
+from io import BytesIO
+
+from skimage import color
+'''
+    PIL resize (W,H)
+    cv2 image is BGR
+    PIL image is RGB
+'''
+class Contrast:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        #c = [0.4, .3, .2, .1, .05]
+        c = [0.4, .3, .2]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        img = np.array(img) / 255.
+        means = np.mean(img, axis=(0, 1), keepdims=True)
+        img = np.clip((img - means) * c + means, 0, 1) * 255
+
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class Brightness:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        #W, H = img.size
+        #c = [.1, .2, .3, .4, .5]
+        c = [.1, .2, .3]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.array(img) / 255.
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+
+        img = sk.color.rgb2hsv(img)
+        img[:, :, 2] = np.clip(img[:, :, 2] + c, 0, 1)
+        img = sk.color.hsv2rgb(img)
+
+        #if isgray:
+        #    img = img[:,:,0]
+        #    img = np.squeeze(img)
+
+        img = np.clip(img, 0, 1) * 255
+        img = Image.fromarray(img.astype(np.uint8))
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img
+        #if isgray:
+        #if isgray:
+        #    img = color.rgb2gray(img)
+
+        #return Image.fromarray(img.astype(np.uint8))
+
+
+class JpegCompression:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        #c = [25, 18, 15, 10, 7]
+        c = [25, 18, 15]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        output = BytesIO()
+        img.save(output, 'JPEG', quality=c)
+        return Image.open(output)
+
+
+class Pixelate:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #c = [0.6, 0.5, 0.4, 0.3, 0.25]
+        c = [0.6, 0.5, 0.4]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        img = img.resize((int(W* c), int(H * c)), Image.BOX)
+        return img.resize((W, H), Image.BOX)
+

augmentation/geometry.py

@@ -0,0 +1,233 @@
+
+import cv2
+import numpy as np
+from PIL import Image, ImageOps
+
+'''
+    PIL resize (W,H)
+    Torch resize is (H,W)
+'''
+class Shrink:
+    def __init__(self):
+        self.tps = cv2.createThinPlateSplineShapeTransformer()
+        self.translateXAbs = TranslateXAbs()
+        self.translateYAbs = TranslateYAbs()
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        img = np.array(img)
+        srcpt = list()
+        dstpt = list()
+
+        W_33 = 0.33 * W
+        W_50 = 0.50 * W
+        W_66 = 0.66 * W
+
+        H_50 = 0.50 * H
+
+        P = 0
+
+        #frac = 0.4
+
+        b = [.2, .3, .4]
+        if mag<0 or mag>=len(b):
+            index = 0
+        else:
+            index = mag
+        frac = b[index]
+
+        # left-most
+        srcpt.append([P, P])
+        srcpt.append([P, H-P])
+        x = np.random.uniform(frac-.1, frac)*W_33 
+        y = np.random.uniform(frac-.1, frac)*H_50
+        dstpt.append([P+x, P+y])
+        dstpt.append([P+x, H-P-y])
+        
+        # 2nd left-most 
+        srcpt.append([P+W_33, P])
+        srcpt.append([P+W_33, H-P])
+        dstpt.append([P+W_33, P+y])
+        dstpt.append([P+W_33, H-P-y])
+        
+        # 3rd left-most 
+        srcpt.append([P+W_66, P])
+        srcpt.append([P+W_66, H-P])
+        dstpt.append([P+W_66, P+y])
+        dstpt.append([P+W_66, H-P-y])
+        
+        # right-most 
+        srcpt.append([W-P, P])
+        srcpt.append([W-P, H-P])
+        dstpt.append([W-P-x, P+y])
+        dstpt.append([W-P-x, H-P-y])
+
+        N = len(dstpt)
+        matches = [cv2.DMatch(i, i, 0) for i in range(N)]
+        dst_shape = np.array(dstpt).reshape((-1, N, 2))
+        src_shape = np.array(srcpt).reshape((-1, N, 2))
+        self.tps.estimateTransformation(dst_shape, src_shape, matches)
+        img = self.tps.warpImage(img)
+        img = Image.fromarray(img)
+
+        if np.random.uniform(0, 1) < 0.5:
+            img = self.translateXAbs(img, val=x)
+        else:
+            img = self.translateYAbs(img, val=y)
+
+        return img
+
+
+class Rotate:
+    def __init__(self, square_side=224):
+        self.side = square_side
+
+    def __call__(self, img, iscurve=False, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+
+        if H!=self.side or W!=self.side:
+            img = img.resize((self.side, self.side), Image.BICUBIC)
+
+        b = [20., 40, 60]
+        if mag<0 or mag>=len(b):
+            index = 1
+        else:
+            index = mag
+        rotate_angle = b[index]
+
+        angle = np.random.uniform(rotate_angle-20, rotate_angle)
+        if np.random.uniform(0, 1) < 0.5:
+            angle = -angle
+
+        #angle = np.random.normal(loc=0., scale=rotate_angle)
+        #angle = min(angle, 2*rotate_angle)
+        #angle = max(angle, -2*rotate_angle)
+
+        expand = False if iscurve else True
+        img = img.rotate(angle=angle, resample=Image.BICUBIC, expand=expand)
+        img = img.resize((W, H), Image.BICUBIC)
+
+        return img
+
+class Perspective:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+
+        # upper-left, upper-right, lower-left, lower-right
+        src =  np.float32([[0, 0], [W, 0], [0, H], [W, H]])
+        #low = 0.3 
+
+        b = [.1, .2, .3]
+        if mag<0 or mag>=len(b):
+            index = 2
+        else:
+            index = mag
+        low = b[index]
+
+        high = 1 - low
+        if np.random.uniform(0, 1) > 0.5:
+            toprightY = np.random.uniform(low, low+.1)*H
+            bottomrightY = np.random.uniform(high-.1, high)*H
+            dest = np.float32([[0, 0], [W, toprightY], [0, H], [W, bottomrightY]])
+        else:
+            topleftY = np.random.uniform(low, low+.1)*H
+            bottomleftY = np.random.uniform(high-.1, high)*H
+            dest = np.float32([[0, topleftY], [W, 0], [0, bottomleftY], [W, H]])
+        M = cv2.getPerspectiveTransform(src, dest)
+        img = np.array(img)
+        img = cv2.warpPerspective(img, M, (W, H) )
+        img = Image.fromarray(img)
+
+        return img
+
+
+class TranslateX:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        b = [.03, .06, .09]
+        if mag<0 or mag>=len(b):
+            index = 2
+        else:
+            index = mag
+        v = b[index]
+        v = np.random.uniform(v-0.03, v)
+
+        v = v * img.size[0]
+        if np.random.uniform(0,1) > 0.5:
+            v = -v
+        return img.transform(img.size, Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+class TranslateY:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        b = [.07, .14, .21]
+        if mag<0 or mag>=len(b):
+            index = 2
+        else:
+            index = mag
+        v = b[index]
+        v = np.random.uniform(v-0.07, v)
+
+        v = v * img.size[1]
+        if np.random.uniform(0,1) > 0.5:
+            v = -v
+        return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+class TranslateXAbs:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, val=0, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        v = np.random.uniform(0, val)
+
+        if np.random.uniform(0,1) > 0.5:
+            v = -v
+        return img.transform(img.size, Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+class TranslateYAbs:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, val=0, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        v = np.random.uniform(0, val)
+
+        if np.random.uniform(0,1) > 0.5:
+            v = -v
+        return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+
+
+
+

augmentation/noise.py

@@ -0,0 +1,94 @@
+
+import numpy as np
+import skimage as sk
+from PIL import Image
+
+'''
+    PIL resize (W,H)
+'''
+class GaussianNoise:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #c = np.random.uniform(.08, .38)
+        b = [.08, 0.1, 0.12]
+        if mag<0 or mag>=len(b):
+            index = 0
+        else:
+            index = mag
+        a = b[index]
+        c = np.random.uniform(a, a+0.03)
+        img = np.array(img) / 255.
+        img = np.clip(img + np.random.normal(size=img.shape, scale=c), 0, 1) * 255
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class ShotNoise:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #c = np.random.uniform(3, 60)
+        b = [13, 8, 3]
+        if mag<0 or mag>=len(b):
+            index = 2
+        else:
+            index = mag
+        a = b[index]
+        c = np.random.uniform(a, a+7)
+        img = np.array(img) / 255.
+        img = np.clip(np.random.poisson(img * c) / float(c), 0, 1) * 255
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class ImpulseNoise:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        #c = np.random.uniform(.03, .27)
+        b = [.03, .07, .11]
+        if mag<0 or mag>=len(b):
+            index = 0
+        else:
+            index = mag
+        a = b[index]
+        c = np.random.uniform(a, a+.04)
+        img = sk.util.random_noise(np.array(img) / 255., mode='s&p', amount=c) * 255
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class SpeckleNoise:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        # c = np.random.uniform(.15, .6)
+        b = [.15, .2, .25]
+        if mag<0 or mag>=len(b):
+            index = 0
+        else:
+            index = mag
+        a = b[index]
+        c = np.random.uniform(a, a+.05)
+        img = np.array(img) / 255.
+        img = np.clip(img + img * np.random.normal(size=img.shape, scale=c), 0, 1) * 255
+        return Image.fromarray(img.astype(np.uint8))
+

augmentation/ops.py

@@ -0,0 +1,87 @@
+
+import cv2
+import numpy as np
+from wand.image import Image as WandImage
+from scipy.ndimage import zoom as scizoom
+from wand.api import library as wandlibrary
+
+class MotionImage(WandImage):
+    def motion_blur(self, radius=0.0, sigma=0.0, angle=0.0):
+        wandlibrary.MagickMotionBlurImage(self.wand, radius, sigma, angle)
+
+def clipped_zoom(img, zoom_factor):
+    h = img.shape[1]
+    # ceil crop height(= crop width)
+    ch = int(np.ceil(h / float(zoom_factor)))
+
+    top = (h - ch) // 2
+    img = scizoom(img[top:top + ch, top:top + ch], (zoom_factor, zoom_factor, 1), order=1)
+    # trim off any extra pixels
+    trim_top = (img.shape[0] - h) // 2
+
+    return img[trim_top:trim_top + h, trim_top:trim_top + h]
+
+def disk(radius, alias_blur=0.1, dtype=np.float32):
+    if radius <= 8:
+        L = np.arange(-8, 8 + 1)
+        ksize = (3, 3)
+    else:
+        L = np.arange(-radius, radius + 1)
+        ksize = (5, 5)
+    X, Y = np.meshgrid(L, L)
+    aliased_disk = np.array((X ** 2 + Y ** 2) <= radius ** 2, dtype=dtype)
+    aliased_disk /= np.sum(aliased_disk)
+
+    # supersample disk to antialias
+    return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)
+
+# modification of https://github.com/FLHerne/mapgen/blob/master/diamondsquare.py
+def plasma_fractal(mapsize=256, wibbledecay=3):
+    """
+    Generate a heightmap using diamond-square algorithm.
+    Return square 2d array, side length 'mapsize', of floats in range 0-255.
+    'mapsize' must be a power of two.
+    """
+    assert (mapsize & (mapsize - 1) == 0)
+    maparray = np.empty((mapsize, mapsize), dtype=np.float_)
+    maparray[0, 0] = 0
+    stepsize = mapsize
+    wibble = 100
+
+    def wibbledmean(array):
+        return array / 4 + wibble * np.random.uniform(-wibble, wibble, array.shape)
+
+    def fillsquares():
+        """For each square of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        cornerref = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        squareaccum = cornerref + np.roll(cornerref, shift=-1, axis=0)
+        squareaccum += np.roll(squareaccum, shift=-1, axis=1)
+        maparray[stepsize // 2:mapsize:stepsize,
+        stepsize // 2:mapsize:stepsize] = wibbledmean(squareaccum)
+
+    def filldiamonds():
+        """For each diamond of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        mapsize = maparray.shape[0]
+        drgrid = maparray[stepsize // 2:mapsize:stepsize, stepsize // 2:mapsize:stepsize]
+        ulgrid = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        ldrsum = drgrid + np.roll(drgrid, 1, axis=0)
+        lulsum = ulgrid + np.roll(ulgrid, -1, axis=1)
+        ltsum = ldrsum + lulsum
+        maparray[0:mapsize:stepsize, stepsize // 2:mapsize:stepsize] = wibbledmean(ltsum)
+        tdrsum = drgrid + np.roll(drgrid, 1, axis=1)
+        tulsum = ulgrid + np.roll(ulgrid, -1, axis=0)
+        ttsum = tdrsum + tulsum
+        maparray[stepsize // 2:mapsize:stepsize, 0:mapsize:stepsize] = wibbledmean(ttsum)
+
+    while stepsize >= 2:
+        fillsquares()
+        filldiamonds()
+        stepsize //= 2
+        wibble /= wibbledecay
+
+    maparray -= maparray.min()
+    return maparray / maparray.max()
+
+

augmentation/pattern.py

@@ -0,0 +1,115 @@
+
+import cv2
+import numpy as np
+from PIL import Image, ImageOps, ImageDraw
+
+'''
+    PIL resize (W,H)
+    Torch resize is (H,W)
+'''
+class VGrid:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, copy=True, max_width=4, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        if copy:
+            img = img.copy()
+        W, H = img.size
+
+        if mag<0 or mag>max_width:
+            line_width = np.random.randint(1, max_width)
+            image_stripe = np.random.randint(1, max_width)
+        else:
+            line_width = 1
+            image_stripe = 3 - mag
+
+        n_lines = W // (line_width + image_stripe) + 1
+        draw = ImageDraw.Draw(img)
+        for i in range(1, n_lines):
+            x = image_stripe*i + line_width*(i-1)
+            draw.line([(x,0), (x,H)], width=line_width, fill='black')
+
+        return img
+
+class HGrid:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, copy=True, max_width=4, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        if copy:
+            img = img.copy()
+        W, H = img.size
+        if mag<0 or mag>max_width:
+            line_width = np.random.randint(1, max_width)
+            image_stripe = np.random.randint(1, max_width)
+        else:
+            line_width = 1
+            image_stripe = 3 - mag
+
+        n_lines = H // (line_width + image_stripe) + 1
+        draw = ImageDraw.Draw(img)
+        for i in range(1, n_lines):
+            y = image_stripe*i + line_width*(i-1)
+            draw.line([(0,y), (W, y)], width=line_width, fill='black')
+
+        return img
+
+class Grid:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        img = VGrid()(img, copy=True, mag=mag)
+        img = HGrid()(img, copy=False, mag=mag)
+        return img
+
+class RectGrid:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, isellipse=False, mag=-1,  prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        img = img.copy()
+        W, H = img.size
+        line_width = 1 
+        image_stripe = 3 - mag #np.random.randint(2, 6)
+        offset = 4 if isellipse else 1
+        n_lines = ((H//2) // (line_width + image_stripe)) + offset
+        draw = ImageDraw.Draw(img)
+        x_center = W // 2
+        y_center = H // 2
+        for i in range(1, n_lines):
+            dx = image_stripe*i + line_width*(i-1)
+            dy = image_stripe*i + line_width*(i-1)
+            x1 = x_center - (dx * W//H)
+            y1 = y_center - dy
+            x2 = x_center + (dx * W/H) 
+            y2 = y_center + dy
+            if isellipse:
+                draw.ellipse([(x1,y1), (x2, y2)], width=line_width, outline='black')
+            else:
+                draw.rectangle([(x1,y1), (x2, y2)], width=line_width, outline='black')
+
+        return img
+
+class EllipseGrid:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        img = RectGrid()(img, isellipse=True, mag=mag, prob=prob)
+        return img

augmentation/process.py

@@ -0,0 +1,123 @@
+
+from PIL import Image
+import PIL.ImageOps, PIL.ImageEnhance
+import numpy as np
+
+class Posterize:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        c = [1, 3, 6]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        bit = np.random.randint(c, c+2)
+        img = PIL.ImageOps.posterize(img, bit)
+
+        return img
+
+
+class Solarize:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        c = [64, 128, 192]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        thresh = np.random.randint(c, c+64)
+        img = PIL.ImageOps.solarize(img, thresh)
+
+        return img
+
+class Invert:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        img = PIL.ImageOps.invert(img)
+
+        return img
+
+    
+class Equalize:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        mg = PIL.ImageOps.equalize(img)
+
+        return img
+
+
+class AutoContrast:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        mg = PIL.ImageOps.autocontrast(img)
+
+        return img
+
+
+class Sharpness:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        c = [.1, .7, 1.3]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        magnitude = np.random.uniform(c, c+.6)
+        img = PIL.ImageEnhance.Sharpness(img).enhance(magnitude)
+
+        return img
+
+
+class Color:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        c = [.1, .7, 1.3]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+        magnitude = np.random.uniform(c, c+.6)
+        img = PIL.ImageEnhance.Color(img).enhance(magnitude)
+
+        return img
+
+

augmentation/test.py

@@ -0,0 +1,43 @@
+
+import os
+import cv2
+from warp import Curve, Distort, Stretch
+from geometry import Rotate, Perspective, Shrink, TranslateX, TranslateY
+from pattern import VGrid, HGrid, Grid, RectGrid, EllipseGrid
+from noise import GaussianNoise, ShotNoise, ImpulseNoise, SpeckleNoise
+from blur import GaussianBlur, DefocusBlur, MotionBlur, GlassBlur, ZoomBlur
+from camera import Contrast, Brightness, JpegCompression, Pixelate
+from weather import Fog, Snow, Frost, Rain, Shadow
+from process import Posterize, Solarize, Invert, Equalize, AutoContrast, Sharpness, Color
+
+from PIL import Image
+import PIL.ImageOps
+import numpy as np
+import argparse
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image', default="images/delivery.png", help='Load image file')
+    parser.add_argument('--results', default="results", help='Load image file')
+    parser.add_argument('--gray', action='store_true', help='Convert to grayscale 1st')
+    opt = parser.parse_args()
+    os.makedirs(opt.results, exist_ok=True)
+
+    img = Image.open(opt.image)
+    img = img.resize( (100,32) )
+    ops = [Curve(), Rotate(), Perspective(), Distort(), Stretch(), Shrink(), TranslateX(), TranslateY(), VGrid(), HGrid(), Grid(), RectGrid(), EllipseGrid()]
+    ops.extend([GaussianNoise(), ShotNoise(), ImpulseNoise(), SpeckleNoise()])
+    ops.extend([GaussianBlur(), DefocusBlur(), MotionBlur(), GlassBlur(), ZoomBlur()])
+    ops.extend([Contrast(), Brightness(), JpegCompression(), Pixelate()])
+    ops.extend([Fog(), Snow(), Frost(), Rain(), Shadow()])
+    ops.extend([Posterize(), Solarize(), Invert(), Equalize(), AutoContrast(), Sharpness(), Color()])
+    for op in ops:
+        for mag in range(3):
+            filename = type(op).__name__ + "-" + str(mag) + ".png"
+            out_img = op(img, mag=mag)
+            if opt.gray:
+                out_img = PIL.ImageOps.grayscale(out_img)
+            out_img.save(os.path.join(opt.results, filename))
+
+

augmentation/warp.py

@@ -0,0 +1,241 @@
+
+import cv2
+import numpy as np
+from PIL import Image, ImageOps
+
+'''
+    PIL resize (W,H)
+    Torch resize is (H,W)
+'''
+class Stretch:
+    def __init__(self):
+        self.tps = cv2.createThinPlateSplineShapeTransformer()
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        img = np.array(img)
+        srcpt = list()
+        dstpt = list()
+
+        W_33 = 0.33 * W
+        W_50 = 0.50 * W
+        W_66 = 0.66 * W
+
+        H_50 = 0.50 * H
+
+        P = 0
+        #frac = 0.4
+
+        b = [.2, .3, .4]
+        if mag<0 or mag>=len(b):
+            index = len(b)-1
+        else:
+            index = mag
+        frac = b[index]
+
+        # left-most
+        srcpt.append([P, P])
+        srcpt.append([P, H-P])
+        srcpt.append([P, H_50])
+        x = np.random.uniform(0, frac)*W_33 #if np.random.uniform(0,1) > 0.5 else 0
+        dstpt.append([P+x, P])
+        dstpt.append([P+x, H-P])
+        dstpt.append([P+x, H_50])
+        
+        # 2nd left-most 
+        srcpt.append([P+W_33, P])
+        srcpt.append([P+W_33, H-P])
+        x = np.random.uniform(-frac, frac)*W_33
+        dstpt.append([P+W_33+x, P])
+        dstpt.append([P+W_33+x, H-P])
+        
+        # 3rd left-most 
+        srcpt.append([P+W_66, P])
+        srcpt.append([P+W_66, H-P])
+        x = np.random.uniform(-frac, frac)*W_33
+        dstpt.append([P+W_66+x, P])
+        dstpt.append([P+W_66+x, H-P])
+        
+        # right-most 
+        srcpt.append([W-P, P])
+        srcpt.append([W-P, H-P])
+        srcpt.append([W-P, H_50])
+        x = np.random.uniform(-frac, 0)*W_33 #if np.random.uniform(0,1) > 0.5 else 0
+        dstpt.append([W-P+x, P])
+        dstpt.append([W-P+x, H-P])
+        dstpt.append([W-P+x, H_50])
+
+        N = len(dstpt)
+        matches = [cv2.DMatch(i, i, 0) for i in range(N)]
+        dst_shape = np.array(dstpt).reshape((-1, N, 2))
+        src_shape = np.array(srcpt).reshape((-1, N, 2))
+        self.tps.estimateTransformation(dst_shape, src_shape, matches)
+        img = self.tps.warpImage(img)
+        img = Image.fromarray(img)
+
+        return img
+
+
+class Distort:
+    def __init__(self):
+        self.tps = cv2.createThinPlateSplineShapeTransformer()
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        img = np.array(img)
+        srcpt = list()
+        dstpt = list()
+
+        W_33 = 0.33 * W
+        W_50 = 0.50 * W
+        W_66 = 0.66 * W
+
+        H_50 = 0.50 * H
+
+        P = 0
+        #frac = 0.4
+
+        b = [.2, .3, .4]
+        if mag<0 or mag>=len(b):
+            index = len(b)-1
+        else:
+            index = mag
+        frac = b[index]
+
+        # top pts
+        srcpt.append([P, P])
+        x = np.random.uniform(0, frac)*W_33
+        y = np.random.uniform(0, frac)*H_50
+        dstpt.append([P+x, P+y])
+        
+        srcpt.append([P+W_33, P])
+        x = np.random.uniform(-frac, frac)*W_33
+        y = np.random.uniform(0, frac)*H_50
+        dstpt.append([P+W_33+x, P+y])
+        
+        srcpt.append([P+W_66, P])
+        x = np.random.uniform(-frac, frac)*W_33
+        y = np.random.uniform(0, frac)*H_50
+        dstpt.append([P+W_66+x, P+y])
+        
+        srcpt.append([W-P, P])
+        x = np.random.uniform(-frac, 0)*W_33
+        y = np.random.uniform(0, frac)*H_50
+        dstpt.append([W-P+x, P+y])
+
+        # bottom pts
+        srcpt.append([P, H-P])
+        x = np.random.uniform(0, frac)*W_33
+        y = np.random.uniform(-frac, 0)*H_50
+        dstpt.append([P+x, H-P+y])
+        
+        srcpt.append([P+W_33, H-P])
+        x = np.random.uniform(-frac, frac)*W_33
+        y = np.random.uniform(-frac, 0)*H_50
+        dstpt.append([P+W_33+x, H-P+y])
+        
+        srcpt.append([P+W_66, H-P])
+        x = np.random.uniform(-frac, frac)*W_33
+        y = np.random.uniform(-frac, 0)*H_50
+        dstpt.append([P+W_66+x, H-P+y])
+        
+        srcpt.append([W-P, H-P])
+        x = np.random.uniform(-frac, 0)*W_33
+        y = np.random.uniform(-frac, 0)*H_50
+        dstpt.append([W-P+x, H-P+y])
+
+        N = len(dstpt)
+        matches = [cv2.DMatch(i, i, 0) for i in range(N)]
+        dst_shape = np.array(dstpt).reshape((-1, N, 2))
+        src_shape = np.array(srcpt).reshape((-1, N, 2))
+        self.tps.estimateTransformation(dst_shape, src_shape, matches)
+        img = self.tps.warpImage(img)
+        img = Image.fromarray(img)
+
+        return img
+
+
+class Curve:
+    def __init__(self, square_side=224):
+        self.tps = cv2.createThinPlateSplineShapeTransformer()
+        self.side = square_side
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+
+        if H!=self.side or W!=self.side:
+            img = img.resize((self.side, self.side), Image.BICUBIC)
+
+        isflip = np.random.uniform(0,1) > 0.5
+        if isflip:
+            img = ImageOps.flip(img)
+            #img = TF.vflip(img)
+
+        img = np.array(img)
+        w = self.side
+        h = self.side
+        w_25 = 0.25 * w
+        w_50 = 0.50 * w
+        w_75 = 0.75 * w
+
+        b = [1.1, .95, .8]
+        if mag<0 or mag>=len(b):
+            index = 0
+        else:
+            index = mag
+        rmin = b[index]
+
+        r = np.random.uniform(rmin, rmin+.1)*h
+        x1 = (r**2 - w_50**2)**0.5
+        h1 = r - x1
+
+        t = np.random.uniform(0.4, 0.5)*h
+
+        w2 = w_50*t/r
+        hi = x1*t/r
+        h2 = h1 + hi  
+
+        sinb_2 = ((1 - x1/r)/2)**0.5
+        cosb_2 = ((1 + x1/r)/2)**0.5
+        w3 = w_50 - r*sinb_2
+        h3 = r - r*cosb_2
+
+        w4 = w_50 - (r-t)*sinb_2
+        h4 = r - (r-t)*cosb_2
+
+        w5 = 0.5*w2
+        h5 = h1 + 0.5*hi
+        h_50 = 0.50*h
+
+        srcpt = [(0,0 ), (w,0 ), (w_50,0), (0,h  ), (w,h    ), (w_25,0), (w_75,0 ),  (w_50,h), (w_25,h), (w_75,h ), (0,h_50), (w,h_50 )]
+        dstpt = [(0,h1), (w,h1), (w_50,0), (w2,h2), (w-w2,h2), (w3, h3), (w-w3,h3),  (w_50,t), (w4,h4 ), (w-w4,h4), (w5,h5 ), (w-w5,h5)]
+
+        N = len(dstpt)
+        matches = [cv2.DMatch(i, i, 0) for i in range(N)]
+        dst_shape = np.array(dstpt).reshape((-1, N, 2))
+        src_shape = np.array(srcpt).reshape((-1, N, 2))
+        self.tps.estimateTransformation(dst_shape, src_shape, matches)
+        img = self.tps.warpImage(img)
+        img = Image.fromarray(img)
+
+        if isflip:
+            #img = TF.vflip(img)
+            img = ImageOps.flip(img)
+            rect = (0, self.side//2, self.side, self.side)
+        else:
+            rect = (0, 0, self.side, self.side//2)
+
+        img = img.crop(rect)
+        img = img.resize((W, H), Image.BICUBIC)
+        return img
+
+

augmentation/weather.py

@@ -0,0 +1,231 @@
+
+import cv2
+import numpy as np
+import math
+from PIL import Image, ImageOps, ImageDraw
+from skimage import color
+from pkg_resources import resource_filename
+from io import BytesIO
+from .ops import plasma_fractal, clipped_zoom, MotionImage
+
+'''
+    PIL resize (W,H)
+'''
+class Fog:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        c = [(1.5, 2), (2., 2), (2.5, 1.7)]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.array(img) / 255.
+        max_val = img.max()
+        fog = c[0] * plasma_fractal(wibbledecay=c[1])[:H, :W][..., np.newaxis]
+        #x += c[0] * plasma_fractal(wibbledecay=c[1])[:224, :224][..., np.newaxis]
+        #return np.clip(x * max_val / (max_val + c[0]), 0, 1) * 255
+        if isgray:
+            fog = np.squeeze(fog)
+        else:
+            fog = np.repeat(fog, 3, axis=2)
+
+        img += fog
+        img = np.clip(img * max_val / (max_val + c[0]), 0, 1) * 255
+        return Image.fromarray(img.astype(np.uint8))
+
+
+class Frost:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        c = [(1, 0.4), (0.8, 0.6), (0.7, 0.7)]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        filename = [resource_filename(__name__, 'frost/frost1.png'),
+                    resource_filename(__name__, 'frost/frost2.png'),
+                    resource_filename(__name__, 'frost/frost3.png'),
+                    resource_filename(__name__, 'frost/frost4.jpg'),
+                    resource_filename(__name__, 'frost/frost5.jpg'),
+                    resource_filename(__name__, 'frost/frost6.jpg')]
+        index = np.random.randint(0, len(filename))
+        filename = filename[index]
+        frost = cv2.imread(filename)
+        #randomly crop and convert to rgb
+        x_start, y_start = np.random.randint(0, frost.shape[0] - H), np.random.randint(0, frost.shape[1] - W)
+        frost = frost[x_start:x_start + H, y_start:y_start + W][..., [2, 1, 0]]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.array(img)
+        
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+
+        img = img * c[0]
+        frost = frost * c[1]
+        img = np.clip(c[0] * img + c[1] * frost, 0, 255)
+        img = Image.fromarray(img.astype(np.uint8))
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img
+
+class Snow:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        W, H = img.size
+        c = [(0.1, 0.3, 3, 0.5, 10, 4, 0.8),
+             (0.2, 0.3, 2, 0.5, 12, 4, 0.7),
+             (0.55, 0.3, 4, 0.9, 12, 8, 0.7)]
+        if mag<0 or mag>=len(c):
+            index = np.random.randint(0, len(c))
+        else:
+            index = mag
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.array(img, dtype=np.float32) / 255.
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+
+        snow_layer = np.random.normal(size=img.shape[:2], loc=c[0], scale=c[1])  # [:2] for monochrome
+
+        #snow_layer = clipped_zoom(snow_layer[..., np.newaxis], c[2])
+        snow_layer[snow_layer < c[3]] = 0
+
+        snow_layer = Image.fromarray((np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode='L')
+        output = BytesIO()
+        snow_layer.save(output, format='PNG')
+        snow_layer = MotionImage(blob=output.getvalue())
+
+        snow_layer.motion_blur(radius=c[4], sigma=c[5], angle=np.random.uniform(-135, -45))
+
+        snow_layer = cv2.imdecode(np.fromstring(snow_layer.make_blob(), np.uint8),
+                                  cv2.IMREAD_UNCHANGED) / 255.
+
+        #snow_layer = cv2.cvtColor(snow_layer, cv2.COLOR_BGR2RGB)
+
+        snow_layer = snow_layer[..., np.newaxis]
+
+        img = c[6] * img
+        gray_img = (1 - c[6]) * np.maximum(img, cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).reshape(H, W, 1) * 1.5 + 0.5)
+        img += gray_img
+        img = np.clip(img + snow_layer + np.rot90(snow_layer, k=2), 0, 1) * 255
+        img = Image.fromarray(img.astype(np.uint8))
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img
+
+class Rain:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        img = img.copy()
+        W, H = img.size
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+        line_width = np.random.randint(1, 2)
+
+        c =[50, 70, 90]
+        if mag<0 or mag>=len(c):
+            index = 0
+        else:
+            index = mag
+        c = c[index]
+
+        n_rains = np.random.randint(c, c+20)
+        slant = np.random.randint(-60, 60)
+        fillcolor = 200 if isgray else (200,200,200)
+
+        draw = ImageDraw.Draw(img)
+        for i in range(1, n_rains):
+            length = np.random.randint(5, 10)
+            x1 = np.random.randint(0, W-length)
+            y1 = np.random.randint(0, H-length)
+            x2 = x1 + length*math.sin(slant*math.pi/180.)
+            y2 = y1 + length*math.cos(slant*math.pi/180.)
+            x2 = int(x2)
+            y2 = int(y2)
+            draw.line([(x1,y1), (x2,y2)], width=line_width, fill=fillcolor)
+
+        return img
+
+class Shadow:
+    def __init__(self):
+        pass
+
+    def __call__(self, img, mag=-1, prob=1.):
+        if np.random.uniform(0,1) > prob:
+            return img
+
+        #img = img.copy()
+        W, H = img.size
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        c =[64, 96, 128]
+        if mag<0 or mag>=len(c):
+            index = 0
+        else:
+            index = mag
+        c = c[index]
+
+        img = img.convert('RGBA')
+        overlay = Image.new('RGBA', img.size, (255,255,255,0))
+        draw = ImageDraw.Draw(overlay) 
+        transparency = np.random.randint(c, c+32)
+        x1 = np.random.randint(0, W//2)
+        y1 = 0
+
+        x2 = np.random.randint(W//2, W)
+        y2 = 0
+
+        x3 = np.random.randint(W//2, W)
+        y3 = H - 1
+
+        x4 = np.random.randint(0, W//2)
+        y4 = H - 1
+
+        draw.polygon([(x1,y1), (x2,y2), (x3,y3), (x4,y4)], fill=(0,0,0,transparency))
+
+        img = Image.alpha_composite(img, overlay)
+        img = img.convert("RGB")
+        if isgray:
+            img = ImageOps.grayscale(img)
+
+        return img

callbacks.py

@@ -0,0 +1,360 @@
+import logging
+import shutil
+import time
+
+import editdistance as ed
+import torchvision.utils as vutils
+from fastai.callbacks.tensorboard import (LearnerTensorboardWriter,
+                                          SummaryWriter, TBWriteRequest,
+                                          asyncTBWriter)
+from fastai.vision import *
+from torch.nn.parallel import DistributedDataParallel
+from torchvision import transforms
+
+import dataset_abinet
+from utils_abinet import CharsetMapper, Timer, blend_mask
+
+
+class IterationCallback(LearnerTensorboardWriter):
+    "A `TrackerCallback` that monitor in each iteration."
+    def __init__(self, learn:Learner, name:str='model', checpoint_keep_num=5,
+                 show_iters:int=50, eval_iters:int=1000, save_iters:int=20000,
+                 start_iters:int=0, stats_iters=20000):
+        #if self.learn.rank is not None: time.sleep(self.learn.rank)  # keep all event files
+        super().__init__(learn, base_dir='.', name=learn.path, loss_iters=show_iters,
+                        stats_iters=stats_iters, hist_iters=stats_iters)
+        self.name, self.bestname = Path(name).name, f'best-{Path(name).name}'
+        self.show_iters = show_iters
+        self.eval_iters = eval_iters
+        self.save_iters = save_iters
+        self.start_iters = start_iters
+        self.checpoint_keep_num = checpoint_keep_num
+        self.metrics_root = 'metrics/'  # rewrite
+        self.timer = Timer()
+        self.host = self.learn.rank is None or self.learn.rank == 0
+
+    def _write_metrics(self, iteration:int, names:List[str], last_metrics:MetricsList)->None:
+        "Writes training metrics to Tensorboard."
+        for i, name in enumerate(names):
+            if last_metrics is None or len(last_metrics) < i+1: return
+            scalar_value = last_metrics[i]
+            self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)
+
+    def _write_sub_loss(self, iteration:int, last_losses:dict)->None:
+        "Writes sub loss to Tensorboard."
+        for name, loss in last_losses.items():
+            scalar_value = to_np(loss)
+            tag = self.metrics_root + name
+            self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    def _save(self, name):
+        if isinstance(self.learn.model, DistributedDataParallel):
+            tmp = self.learn.model
+            self.learn.model = self.learn.model.module
+            self.learn.save(name)
+            self.learn.model = tmp
+        else: self.learn.save(name)
+
+    def _validate(self, dl=None, callbacks=None, metrics=None, keeped_items=False):
+        "Validate on `dl` with potential `callbacks` and `metrics`."
+        dl = ifnone(dl, self.learn.data.valid_dl)
+        metrics = ifnone(metrics, self.learn.metrics)
+        cb_handler = CallbackHandler(ifnone(callbacks, []), metrics)
+        cb_handler.on_train_begin(1, None, metrics); cb_handler.on_epoch_begin()
+        if keeped_items: cb_handler.state_dict.update(dict(keeped_items=[]))
+        val_metrics = validate(self.learn.model, dl, self.loss_func, cb_handler)
+        cb_handler.on_epoch_end(val_metrics)
+        if keeped_items: return cb_handler.state_dict['keeped_items']
+        else: return cb_handler.state_dict['last_metrics']
+
+    def jump_to_epoch_iter(self, epoch:int, iteration:int)->None:
+        try:
+            self.learn.load(f'{self.name}_{epoch}_{iteration}', purge=False)
+            logging.info(f'Loaded {self.name}_{epoch}_{iteration}')
+        except: logging.info(f'Model {self.name}_{epoch}_{iteration} not found.')
+
+    def on_train_begin(self, n_epochs, **kwargs):
+        # TODO: can not write graph here
+        # super().on_train_begin(**kwargs)
+        self.best = -float('inf')
+        self.timer.tic()
+        if self.host:
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if checkpoint_path.exists():
+                os.remove(checkpoint_path)
+            open(checkpoint_path, 'w').close()
+        return {'skip_validate': True, 'iteration':self.start_iters}  # disable default validate
+
+    def on_batch_begin(self, **kwargs:Any)->None:
+        self.timer.toc_data()
+        super().on_batch_begin(**kwargs)
+
+    def on_batch_end(self, iteration, epoch, last_loss, smooth_loss, train, **kwargs):
+        super().on_batch_end(last_loss, iteration, train, **kwargs)
+        if iteration == 0: return
+
+        if iteration % self.loss_iters == 0:
+            last_losses = self.learn.loss_func.last_losses
+            self._write_sub_loss(iteration=iteration, last_losses=last_losses)
+            self.tbwriter.add_scalar(tag=self.metrics_root + 'lr',
+                scalar_value=self.opt.lr, global_step=iteration)
+
+        if iteration % self.show_iters == 0:
+            log_str = f'epoch {epoch} iter {iteration}: loss = {last_loss:6.4f},  ' \
+                      f'smooth loss = {smooth_loss:6.4f}'
+            logging.info(log_str)
+            # log_str = f'data time = {self.timer.data_diff:.4f}s, runing time = {self.timer.running_diff:.4f}s'
+            # logging.info(log_str)
+
+        if iteration % self.eval_iters == 0:
+            # TODO: or remove time to on_epoch_end
+            # 1. Record time
+            log_str = f'average data time = {self.timer.average_data_time():.4f}s, ' \
+                      f'average running time = {self.timer.average_running_time():.4f}s'
+            logging.info(log_str)
+
+            # 2. Call validate
+            last_metrics = self._validate()
+            self.learn.model.train()
+            log_str = f'epoch {epoch} iter {iteration}: eval loss = {last_metrics[0]:6.4f},  ' \
+                      f'ccr = {last_metrics[1]:6.4f},  cwr = {last_metrics[2]:6.4f},  ' \
+                      f'ted = {last_metrics[3]:6.4f},  ned = {last_metrics[4]:6.4f},  ' \
+                      f'ted/w = {last_metrics[5]:6.4f}, '
+            logging.info(log_str)
+            names = ['eval_loss', 'ccr', 'cwr', 'ted', 'ned', 'ted/w']
+            self._write_metrics(iteration, names, last_metrics)
+
+            # 3. Save best model
+            current = last_metrics[2]
+            if current is not None and current > self.best:
+                logging.info(f'Better model found at epoch {epoch}, '\
+                             f'iter {iteration} with accuracy value: {current:6.4f}.')
+                self.best = current
+                self._save(f'{self.bestname}')
+
+        if iteration % self.save_iters == 0 and self.host:
+            logging.info(f'Save model {self.name}_{epoch}_{iteration}')
+            filename = f'{self.name}_{epoch}_{iteration}'
+            self._save(filename)
+
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if not checkpoint_path.exists():
+                open(checkpoint_path, 'w').close()
+            with open(checkpoint_path, 'r') as file:
+                checkpoints = yaml.load(file, Loader=yaml.FullLoader) or dict()
+            checkpoints['all_checkpoints'] = (
+                checkpoints.get('all_checkpoints') or list())
+            checkpoints['all_checkpoints'].insert(0, filename)
+            if len(checkpoints['all_checkpoints']) > self.checpoint_keep_num:
+                removed_checkpoint = checkpoints['all_checkpoints'].pop()
+                removed_checkpoint =  self.learn.path/self.learn.model_dir/f'{removed_checkpoint}.pth'
+                os.remove(removed_checkpoint)
+            checkpoints['current_checkpoint'] = filename
+            with open(checkpoint_path, 'w') as file:
+                yaml.dump(checkpoints, file)
+
+
+        self.timer.toc_running()
+
+    def on_train_end(self, **kwargs):
+        #self.learn.load(f'{self.bestname}', purge=False)
+        pass
+
+    def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:
+        self._write_embedding(iteration=iteration)
+
+
+class TextAccuracy(Callback):
+    _names = ['ccr', 'cwr', 'ted', 'ned', 'ted/w']
+    def __init__(self, charset_path, max_length, case_sensitive, model_eval):
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+
+        self.model_eval = model_eval or 'alignment'
+        assert self.model_eval in ['vision', 'language', 'alignment']
+
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+        self.total_ed = 0.
+        self.total_ned = 0.
+
+    def _get_output(self, last_output):
+        if isinstance(last_output, (tuple, list)):
+            for res in last_output:
+                if res['name'] == self.model_eval: output = res
+        else: output = last_output
+        return output
+
+    def _update_output(self, last_output, items):
+        if isinstance(last_output, (tuple, list)):
+            for res in last_output:
+                if res['name'] == self.model_eval: res.update(items)
+        else: last_output.update(items)
+        return last_output
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        output = self._get_output(last_output)
+        logits, pt_lengths = output['logits'], output['pt_lengths']
+        pt_text, pt_scores, pt_lengths_ = self.decode(logits)
+        assert (pt_lengths == pt_lengths_).all(), f'{pt_lengths} != {pt_lengths_} for {pt_text}'
+        last_output = self._update_output(last_output, {'pt_text':pt_text, 'pt_scores':pt_scores})
+
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        label = last_target[0]
+        if label.dim() == 3: label = label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in label]
+
+        for i in range(len(gt_text)):
+            if not self.case_sensitive:
+                gt_text[i], pt_text[i] = gt_text[i].lower(), pt_text[i].lower()
+            distance = ed.eval(gt_text[i], pt_text[i])
+            self.total_ed += distance
+            self.total_ned += float(distance) / max(len(gt_text[i]), 1)
+
+            if gt_text[i] == pt_text[i]:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+
+            for j in range(min(len(gt_text[i]), len(pt_text[i]))):
+                if gt_text[i][j] == pt_text[i][j]:
+                    self.correct_num_char += 1
+            self.total_num_char += len(gt_text[i])
+
+        return {'last_output': last_output}
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                self.total_ed,
+                self.total_ned,
+                self.total_ed / self.total_num_word]
+        return add_metrics(last_metrics, mets)
+
+    def decode(self, logit):
+        """ Greed decode """
+        # TODO: test running time and decode on GPU
+        out = F.softmax(logit, dim=2)
+        pt_text, pt_scores, pt_lengths = [], [], []
+        for o in out:
+            text = self.charset.get_text(o.argmax(dim=1), padding=False, trim=False)
+            text = text.split(self.charset.null_char)[0]  # end at end-token
+            pt_text.append(text)
+            pt_scores.append(o.max(dim=1)[0])
+            pt_lengths.append(min(len(text) + 1, self.max_length))  # one for end-token
+        pt_scores = torch.stack(pt_scores)
+        pt_lengths = pt_scores.new_tensor(pt_lengths, dtype=torch.long)
+        return pt_text, pt_scores, pt_lengths
+
+
+class TopKTextAccuracy(TextAccuracy):
+    _names = ['ccr', 'cwr']
+    def __init__(self, k, charset_path, max_length, case_sensitive, model_eval):
+        self.k = k
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        logits, pt_lengths = last_output['logits'], last_output['pt_lengths']
+        gt_labels, gt_lengths = last_target[:]
+
+        for logit, pt_length, label, length in zip(logits, pt_lengths, gt_labels, gt_lengths):
+            word_flag = True
+            for i in range(length):
+                char_logit = logit[i].topk(self.k)[1]
+                char_label = label[i].argmax(-1)
+                if char_label in char_logit: self.correct_num_char += 1
+                else: word_flag = False
+                self.total_num_char += 1
+            if pt_length == length and word_flag:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                0., 0., 0.]
+        return add_metrics(last_metrics, mets)
+
+
+class DumpPrediction(LearnerCallback):
+
+    def __init__(self, learn, dataset, charset_path, model_eval, image_only=False, debug=False):
+        super().__init__(learn=learn)
+        self.debug = debug
+        self.model_eval = model_eval or 'alignment'
+        self.image_only = image_only
+        assert self.model_eval in ['vision', 'language', 'alignment']
+
+        self.dataset, self.root = dataset, Path(self.learn.path)/f'{dataset}-{self.model_eval}'
+        self.attn_root = self.root/'attn'
+        self.charset = CharsetMapper(charset_path)
+        if self.root.exists(): shutil.rmtree(self.root)
+        self.root.mkdir(), self.attn_root.mkdir()
+
+        self.pil = transforms.ToPILImage()
+        self.tensor = transforms.ToTensor()
+        size = self.learn.data.img_h, self.learn.data.img_w
+        self.resize = transforms.Resize(size=size, interpolation=0)
+        self.c = 0
+
+    def on_batch_end(self, last_input, last_output, last_target, **kwargs):
+        if isinstance(last_output, (tuple, list)):
+            for res in last_output:
+                if res['name'] == self.model_eval: pt_text = res['pt_text']
+                if res['name'] == 'vision': attn_scores = res['attn_scores'].detach().cpu()
+                if res['name'] == self.model_eval: logits = res['logits']
+        else:
+            pt_text = last_output['pt_text']
+            attn_scores = last_output['attn_scores'].detach().cpu()
+            logits = last_output['logits']
+
+        images = last_input[0] if isinstance(last_input, (tuple, list)) else last_input
+        images = images.detach().cpu()
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        gt_label = last_target[0]
+        if gt_label.dim() == 3: gt_label = gt_label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in gt_label]
+
+        prediction, false_prediction = [], []
+        for gt, pt, image, attn, logit in zip(gt_text, pt_text, images, attn_scores, logits):
+            prediction.append(f'{gt}\t{pt}\n')
+            if gt != pt:
+                if self.debug:
+                    scores = torch.softmax(logit, dim=-1)[:max(len(pt), len(gt)) + 1]
+                    logging.info(f'{self.c} gt {gt}, pt {pt}, logit {logit.shape}, scores {scores.topk(5, dim=-1)}')
+                false_prediction.append(f'{gt}\t{pt}\n')
+
+            image = self.learn.data.denorm(image)
+            if not self.image_only:
+                image_np = np.array(self.pil(image))
+                attn_pil = [self.pil(a) for a in attn[:, None, :, :]]
+                attn = [self.tensor(self.resize(a)).repeat(3, 1, 1) for a in attn_pil]
+                attn_sum = np.array([np.array(a) for a in attn_pil[:len(pt)]]).sum(axis=0)
+                blended_sum = self.tensor(blend_mask(image_np, attn_sum))
+                blended = [self.tensor(blend_mask(image_np, np.array(a))) for a in attn_pil]
+                save_image = torch.stack([image] + attn + [blended_sum] + blended)
+                save_image = save_image.view(2, -1, *save_image.shape[1:])
+                save_image = save_image.permute(1, 0, 2, 3, 4).flatten(0, 1)
+                vutils.save_image(save_image, self.attn_root/f'{self.c}_{gt}_{pt}.jpg',
+                                nrow=2, normalize=True, scale_each=True)
+            else:
+                self.pil(image).save(self.attn_root/f'{self.c}_{gt}_{pt}.jpg')
+            self.c += 1
+
+        with open(self.root/f'{self.model_eval}.txt', 'a') as f: f.writelines(prediction)
+        with open(self.root/f'{self.model_eval}-false.txt', 'a') as f: f.writelines(false_prediction)

captum/__init__.py

@@ -0,0 +1,3 @@
+#!/usr/bin/env python3
+
+__version__ = "0.5.0"

captum/_utils/av.py

@@ -0,0 +1,499 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+import re
+import warnings
+from typing import Any, List, Optional, Tuple, Union
+
+import captum._utils.common as common
+import torch
+from captum.attr import LayerActivation
+from torch import Tensor
+from torch.nn import Module
+from torch.utils.data import DataLoader, Dataset
+
+
+class AV:
+    r"""
+    This class provides functionality to store and load activation vectors
+    generated for pre-defined neural network layers.
+    It also provides functionality to check if activation vectors already
+    exist in the manifold and other auxiliary functions.
+
+    This class also defines a torch `Dataset`, representing Activation Vectors,
+    which enables lazy access to activation vectors and layer stored in the manifold.
+
+    """
+
+    r"""
+        The name of the subfolder in the manifold where the activation vectors
+        are stored.
+    """
+
+    class AVDataset(Dataset):
+        r"""
+        This dataset enables access to activation vectors for a given `model` stored
+        under a pre-defined path.
+        The iterator of this dataset returns a batch of data tensors.
+        Additionally, subsets of the model activations can be loaded based on layer
+        or identifier or num_id (representing batch number in source dataset).
+        """
+
+        def __init__(
+            self,
+            path: str,
+            model_id: str,
+            identifier: Optional[str] = None,
+            layer: Optional[str] = None,
+            num_id: Optional[str] = None,
+        ):
+            r"""
+            Loads into memory the list of all activation file paths associated
+            with the input `model_id`.
+
+            Args:
+                path (str): The path where the activation vectors
+                        for the `layer` are stored.
+                model_id (str): The name/version of the model for which layer
+                        activations are being computed and stored.
+                identifier (str or None): An optional identifier for the layer
+                        activations. Can be used to distinguish between activations for
+                        different training batches.
+                layer (str or None): The layer for which the activation vectors
+                        are computed.
+                num_id (str): An optional string representing the batch number for
+                    which the activation vectors are computed
+            """
+
+            self.av_filesearch = AV._construct_file_search(
+                path, model_id, identifier, layer, num_id
+            )
+
+            files = glob.glob(self.av_filesearch)
+
+            self.files = AV.sort_files(files)
+
+        def __getitem__(self, idx: int) -> Union[Tensor, Tuple[Tensor, ...]]:
+            assert idx < len(self.files), "Layer index is out of bounds!"
+            fl = self.files[idx]
+            av = torch.load(fl)
+            return av
+
+        def __len__(self):
+            return len(self.files)
+
+    AV_DIR_NAME: str = "av"
+
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    def _assemble_model_dir(path: str, model_id: str) -> str:
+        r"""
+        Returns a directory path for the given source path `path` and `model_id.`
+        This path is suffixed with the '/' delimiter.
+        """
+        return "/".join([path, AV.AV_DIR_NAME, model_id, ""])
+
+    @staticmethod
+    def _assemble_file_path(source_dir: str, identifier: str, layer: str) -> str:
+        r"""
+        Returns a full filepath given a source directory, layer, and required
+        identifier. The source dir is not required to end with a "/" delimiter.
+        """
+        if not source_dir.endswith("/"):
+            source_dir += "/"
+
+        filepath = os.path.join(source_dir, identifier)
+
+        filepath = os.path.join(filepath, layer)
+
+        return filepath
+
+    @staticmethod
+    def _construct_file_search(
+        source_dir: str,
+        model_id: str,
+        identifier: Optional[str] = None,
+        layer: Optional[str] = None,
+        num_id: Optional[str] = None,
+    ) -> str:
+        r"""
+        Returns a search string that can be used by glob to search `source_dir/model_id`
+        for the desired layer/identifier pair. Leaving `layer` as None will search ids
+        over all layers, and leaving `identifier` as none will search layers over all
+        ids.  Leaving both as none will return a path to glob for every activation.
+        Assumes identifier is always specified when saving activations, so that
+        activations live at source_dir/model_id/identifier/layer
+        (and never source_dir/model_id/layer)
+        """
+
+        av_filesearch = AV._assemble_model_dir(source_dir, model_id)
+
+        av_filesearch = os.path.join(
+            av_filesearch, "*" if identifier is None else identifier
+        )
+
+        av_filesearch = os.path.join(av_filesearch, "*" if layer is None else layer)
+
+        av_filesearch = os.path.join(
+            av_filesearch, "*.pt" if num_id is None else "%s.pt" % num_id
+        )
+
+        return av_filesearch
+
+    @staticmethod
+    def exists(
+        path: str,
+        model_id: str,
+        identifier: Optional[str] = None,
+        layer: Optional[str] = None,
+        num_id: Optional[str] = None,
+    ) -> bool:
+        r"""
+        Verifies whether the model + layer activations exist
+        under the path.
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `model_id` are stored.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            identifier (str or None): An optional identifier for the layer activations.
+                    Can be used to distinguish between activations for different
+                    training batches. For example, the id could be a suffix composed of
+                    a train/test label and numerical value, such as "-train-xxxxx".
+                    The numerical id is often a monotonic sequence taken from datetime.
+            layer (str or None): The layer for which the activation vectors are
+                    computed.
+            num_id (str): An optional string representing the batch number for which
+                    the activation vectors are computed
+
+        Returns:
+            exists (bool): Indicating whether the activation vectors for the `layer`
+                    and `identifier` (if provided) and num_id (if provided) were stored
+                    in the manifold. If no `identifier` is provided, will return `True`
+                    if any layer activation exists, whether it has an identifier or
+                    not, and vice-versa.
+        """
+        av_dir = AV._assemble_model_dir(path, model_id)
+        av_filesearch = AV._construct_file_search(
+            path, model_id, identifier, layer, num_id
+        )
+        return os.path.exists(av_dir) and len(glob.glob(av_filesearch)) > 0
+
+    @staticmethod
+    def save(
+        path: str,
+        model_id: str,
+        identifier: str,
+        layers: Union[str, List[str]],
+        act_tensors: Union[Tensor, List[Tensor]],
+        num_id: str,
+    ) -> None:
+        r"""
+        Saves the activation vectors `act_tensor` for the
+        `layer` under the manifold `path`.
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `layer` are stored.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            identifier (str or None): An optional identifier for the layer
+                    activations. Can be used to distinguish between activations for
+                    different training batches. For example, the identifier could be
+                    a suffix composed of a train/test label and numerical value, such
+                    as "-src-abc".
+                    Additionally, (abc) could be a unique identifying number. For
+                    example, it is automatically created in
+                    AV.generate_dataset_activations from batch index.
+                    It assumes identifier is same for all layers if a list of
+                    `layers` is provided.
+            layers (str or List of str): The layer(s) for which the activation vectors
+                    are computed.
+            act_tensors (Tensor or List of Tensor): A batch of activation vectors.
+                    This must match the dimension of `layers`.
+            num_id (str): string representing the batch number for which the activation
+                    vectors are computed
+        """
+        if isinstance(layers, str):
+            layers = [layers]
+        if isinstance(act_tensors, Tensor):
+            act_tensors = [act_tensors]
+
+        if len(layers) != len(act_tensors):
+            raise ValueError("The dimension of `layers` and `act_tensors` must match!")
+
+        av_dir = AV._assemble_model_dir(path, model_id)
+
+        for i, layer in enumerate(layers):
+            av_save_fl_path = os.path.join(
+                AV._assemble_file_path(av_dir, identifier, layer), "%s.pt" % num_id
+            )
+
+            layer_dir = os.path.dirname(av_save_fl_path)
+            if not os.path.exists(layer_dir):
+                os.makedirs(layer_dir)
+            torch.save(act_tensors[i], av_save_fl_path)
+
+    @staticmethod
+    def load(
+        path: str,
+        model_id: str,
+        identifier: Optional[str] = None,
+        layer: Optional[str] = None,
+        num_id: Optional[str] = None,
+    ) -> AVDataset:
+        r"""
+        Loads lazily the activation vectors for given `model_id` and
+        `layer` saved under the `path`.
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `layer` are stored.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            identifier (str or None): An optional identifier for the layer
+                    activations. Can be used to distinguish between activations for
+                    different training batches.
+            layer (str or None): The layer for which the activation vectors
+                are computed.
+            num_id (str): An optional string representing the batch number for which
+                    the activation vectors are computed
+
+        Returns:
+            dataset (AV.AVDataset): AV.AVDataset that allows to iterate
+                    over the activation vectors for given layer, identifier (if
+                    provided), num_id (if provided).  Returning an AV.AVDataset as
+                    opposed to a DataLoader constructed from it offers more
+                    flexibility.  Raises RuntimeError if activation vectors are not
+                    found.
+        """
+
+        av_save_dir = AV._assemble_model_dir(path, model_id)
+
+        if os.path.exists(av_save_dir):
+            avdataset = AV.AVDataset(path, model_id, identifier, layer, num_id)
+            return avdataset
+        else:
+            raise RuntimeError(
+                f"Activation vectors for model {model_id} was not found at path {path}"
+            )
+
+    @staticmethod
+    def _manage_loading_layers(
+        path: str,
+        model_id: str,
+        layers: Union[str, List[str]],
+        load_from_disk: bool = True,
+        identifier: Optional[str] = None,
+        num_id: Optional[str] = None,
+    ) -> List[str]:
+        r"""
+        Returns unsaved layers, and deletes saved layers if load_from_disk is False.
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `layer` are stored.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            layers (str or List of str): The layer(s) for which the activation vectors
+                    are computed.
+            identifier (str or None): An optional identifier for the layer
+                    activations. Can be used to distinguish between activations for
+                    different training batches.
+            num_id (str): An optional string representing the batch number for which the
+                    activation vectors are computed
+
+        Returns:
+            List of layer names for which activations should be generated
+        """
+
+        layers = [layers] if isinstance(layers, str) else layers
+        unsaved_layers = []
+
+        if load_from_disk:
+            for layer in layers:
+                if not AV.exists(path, model_id, identifier, layer, num_id):
+                    unsaved_layers.append(layer)
+        else:
+            unsaved_layers = layers
+            warnings.warn(
+                "Overwriting activations: load_from_disk is set to False. Removing all "
+                f"activations matching specified parameters {{path: {path}, "
+                f"model_id: {model_id}, layers: {layers}, identifier: {identifier}}} "
+                "before generating new activations."
+            )
+            for layer in layers:
+                files = glob.glob(
+                    AV._construct_file_search(path, model_id, identifier, layer)
+                )
+                for filename in files:
+                    os.remove(filename)
+
+        return unsaved_layers
+
+    @staticmethod
+    def _compute_and_save_activations(
+        path: str,
+        model: Module,
+        model_id: str,
+        layers: Union[str, List[str]],
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        identifier: str,
+        num_id: str,
+        additional_forward_args: Any = None,
+        load_from_disk: bool = True,
+    ) -> None:
+        r"""
+        Computes layer activations for the given inputs and specified `layers`
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `layer` are stored.
+            model (torch.nn.Module): An instance of pytorch model. This model should
+                    define all of its layers as attributes of the model.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            layers (str or List of str): The layer(s) for which the activation vectors
+                    are computed.
+            inputs (tensor or tuple of tensors): Batch of examples for
+                    which influential instances are computed. They are passed to the
+                    input `model`. The first dimension in `inputs` tensor or tuple of
+                    tensors corresponds to the batch size.
+            identifier (str or None): An optional identifier for the layer
+                    activations. Can be used to distinguish between activations for
+                    different training batches.
+            num_id (str): An required string representing the batch number for which the
+                    activation vectors are computed
+            additional_forward_args (optional):  Additional arguments that will be
+                    passed to `model` after inputs.
+                    Default: None
+            load_from_disk (bool): Forces function to regenerate activations if False.
+                    Default: True
+        """
+        unsaved_layers = AV._manage_loading_layers(
+            path,
+            model_id,
+            layers,
+            load_from_disk,
+            identifier,
+            num_id,
+        )
+        layer_modules = [
+            common._get_module_from_name(model, layer) for layer in unsaved_layers
+        ]
+        if len(unsaved_layers) > 0:
+            layer_act = LayerActivation(model, layer_modules)
+            new_activations = layer_act.attribute.__wrapped__(  # type: ignore
+                layer_act, inputs, additional_forward_args
+            )
+            AV.save(path, model_id, identifier, unsaved_layers, new_activations, num_id)
+
+    @staticmethod
+    def _unpack_data(data: Union[Any, Tuple[Any, Any]]) -> Any:
+        r"""
+        Helper to extract input from labels when getting items from a Dataset. Assumes
+        that data is either a single value, or a tuple containing two elements.
+        The input could itself be a Tuple containing multiple values. If your
+        dataset returns a Tuple with more than 2 elements, please reformat it such that
+        all inputs are formatted into a tuple stored at the first position.
+        """
+        if isinstance(data, tuple) or isinstance(data, list):
+            data = data[0]
+        return data
+
+    r"""TODO:
+    1. Can propagate saving labels along with activations.
+    2. Use of additional_forward_args when sourcing from dataset?
+    """
+
+    @staticmethod
+    def generate_dataset_activations(
+        path: str,
+        model: Module,
+        model_id: str,
+        layers: Union[str, List[str]],
+        dataloader: DataLoader,
+        identifier: str = "default",
+        load_from_disk: bool = True,
+        return_activations: bool = False,
+    ) -> Optional[Union[AVDataset, List[AVDataset]]]:
+        r"""
+        Computes layer activations for a source dataset and specified `layers`. Assumes
+        that the dataset returns a single value, or a tuple containing two elements
+        (see AV._unpack_data).
+
+        Args:
+            path (str): The path where the activation vectors
+                    for the `layer` are stored.
+            module (torch.nn.Module): An instance of pytorch model. This model should
+                    define all of its layers as attributes of the model.
+            model_id (str): The name/version of the model for which layer activations
+                    are being computed and stored.
+            layers (str or List of str): The layer(s) for which the activation vectors
+                    are computed.
+            dataloader (torch.utils.data.DataLoader): DataLoader that yields Dataset
+                    for which influential instances are computed. They are passed to
+                    input `model`.
+            identifier (str or None): An identifier for the layer
+                    activations. Can be used to distinguish between activations for
+                    different training batches.
+                    Default: "default"
+            load_from_disk (bool): Forces function to regenerate activations if False.
+                    Default: True
+            return_activations (bool, optional): Whether to return the activations.
+                    Default: False
+        Returns: If `return_activations == True`, returns a single `AVDataset` if
+                    `layers` is a str, otherwise, a list of `AVDataset`s of the length
+                    of `layers`, where each element corresponds to a layer.  In either
+                    case, `AVDataset`'s represent the activations for a single layer,
+                    over the entire `dataloader`.  If `return_activations == False`,
+                    does not return anything.
+
+        """
+
+        unsaved_layers = AV._manage_loading_layers(
+            path,
+            model_id,
+            layers,
+            load_from_disk,
+            identifier,
+        )
+        if len(unsaved_layers) > 0:
+            for i, data in enumerate(dataloader):
+                AV._compute_and_save_activations(
+                    path,
+                    model,
+                    model_id,
+                    layers,
+                    AV._unpack_data(data),
+                    identifier,
+                    str(i),
+                )
+
+        if not return_activations:
+            return None
+        if isinstance(layers, str):
+            return AV.load(path, model_id, identifier, layers)
+        else:
+            return [AV.load(path, model_id, identifier, layer) for layer in layers]
+
+    @staticmethod
+    def sort_files(files: List[str]) -> List[str]:
+        r"""
+        Utility for sorting files based on natural sorting instead of the default
+        lexigraphical sort.
+        """
+
+        def split_alphanum(s):
+            r"""
+            Splits string into a list of strings and numbers
+                "z23a" -> ["z", 23, "a"]
+            """
+
+            return [int(x) if x.isdigit() else x for x in re.split("([0-9]+)", s)]
+
+        return sorted(files, key=split_alphanum)

captum/_utils/common.py

@@ -0,0 +1,679 @@
+#!/usr/bin/env python3
+import typing
+from enum import Enum
+from functools import reduce
+from inspect import signature
+from typing import Any, Callable, cast, Dict, List, overload, Tuple, Union
+
+import numpy as np
+import torch
+from captum._utils.typing import (
+    BaselineType,
+    Literal,
+    TargetType,
+    TensorOrTupleOfTensorsGeneric,
+    TupleOrTensorOrBoolGeneric,
+)
+from torch import device, Tensor
+from torch.nn import Module
+
+
+class ExpansionTypes(Enum):
+    repeat = 1
+    repeat_interleave = 2
+
+
+def safe_div(
+    numerator: Tensor,
+    denom: Union[Tensor, int, float],
+    default_denom: Union[Tensor, int, float] = 1.0,
+) -> Tensor:
+    r"""
+    A simple utility function to perform `numerator / denom`
+    if the statement is undefined => result will be `numerator / default_denorm`
+    """
+    if isinstance(denom, (int, float)):
+        return numerator / (denom if denom != 0 else default_denom)
+
+    # convert default_denom to tensor if it is float
+    if not torch.is_tensor(default_denom):
+        default_denom = torch.tensor(
+            default_denom, dtype=denom.dtype, device=denom.device
+        )
+
+    return numerator / torch.where(denom != 0, denom, default_denom)
+
+
+@typing.overload
+def _is_tuple(inputs: Tensor) -> Literal[False]:
+    ...
+
+
+@typing.overload
+def _is_tuple(inputs: Tuple[Tensor, ...]) -> Literal[True]:
+    ...
+
+
+def _is_tuple(inputs: Union[Tensor, Tuple[Tensor, ...]]) -> bool:
+    return isinstance(inputs, tuple)
+
+
+def _validate_target(num_samples: int, target: TargetType) -> None:
+    if isinstance(target, list) or (
+        isinstance(target, torch.Tensor) and torch.numel(target) > 1
+    ):
+        assert num_samples == len(target), (
+            "The number of samples provied in the"
+            "input {} does not match with the number of targets. {}".format(
+                num_samples, len(target)
+            )
+        )
+
+
+def _validate_input(
+    inputs: Tuple[Tensor, ...],
+    baselines: Tuple[Union[Tensor, int, float], ...],
+    draw_baseline_from_distrib: bool = False,
+) -> None:
+    assert len(inputs) == len(baselines), (
+        "Input and baseline must have the same "
+        "dimensions, baseline has {} features whereas input has {}.".format(
+            len(baselines), len(inputs)
+        )
+    )
+
+    for input, baseline in zip(inputs, baselines):
+        if draw_baseline_from_distrib:
+            assert (
+                isinstance(baseline, (int, float))
+                or input.shape[1:] == baseline.shape[1:]
+            ), (
+                "The samples in input and baseline batches must have"
+                " the same shape or the baseline corresponding to the"
+                " input tensor must be a scalar."
+                " Found baseline: {} and input: {} ".format(baseline, input)
+            )
+        else:
+            assert (
+                isinstance(baseline, (int, float))
+                or input.shape == baseline.shape
+                or baseline.shape[0] == 1
+            ), (
+                "Baseline can be provided as a tensor for just one input and"
+                " broadcasted to the batch or input and baseline must have the"
+                " same shape or the baseline corresponding to each input tensor"
+                " must be a scalar. Found baseline: {} and input: {}".format(
+                    baseline, input
+                )
+            )
+
+
+def _zeros(inputs: Tuple[Tensor, ...]) -> Tuple[int, ...]:
+    r"""
+    Takes a tuple of tensors as input and returns a tuple that has the same
+    length as `inputs` with each element as the integer 0.
+    """
+    return tuple(0 if input.dtype is not torch.bool else False for input in inputs)
+
+
+def _format_baseline(
+    baselines: BaselineType, inputs: Tuple[Tensor, ...]
+) -> Tuple[Union[Tensor, int, float], ...]:
+    if baselines is None:
+        return _zeros(inputs)
+
+    if not isinstance(baselines, tuple):
+        baselines = (baselines,)
+
+    for baseline in baselines:
+        assert isinstance(
+            baseline, (torch.Tensor, int, float)
+        ), "baseline input argument must be either a torch.Tensor or a number \
+            however {} detected".format(
+            type(baseline)
+        )
+
+    return baselines
+
+
+@overload
+def _format_tensor_into_tuples(inputs: None) -> None:
+    ...
+
+
+@overload
+def _format_tensor_into_tuples(
+    inputs: Union[Tensor, Tuple[Tensor, ...]]
+) -> Tuple[Tensor, ...]:
+    ...
+
+
+def _format_tensor_into_tuples(
+    inputs: Union[None, Tensor, Tuple[Tensor, ...]]
+) -> Union[None, Tuple[Tensor, ...]]:
+    if inputs is None:
+        return None
+    if not isinstance(inputs, tuple):
+        assert isinstance(
+            inputs, torch.Tensor
+        ), "`inputs` must have type " "torch.Tensor but {} found: ".format(type(inputs))
+        inputs = (inputs,)
+    return inputs
+
+
+def _format_inputs(inputs: Any, unpack_inputs: bool = True) -> Any:
+    return (
+        inputs
+        if (isinstance(inputs, tuple) or isinstance(inputs, list)) and unpack_inputs
+        else (inputs,)
+    )
+
+
+def _format_float_or_tensor_into_tuples(
+    inputs: Union[float, Tensor, Tuple[Union[float, Tensor], ...]]
+) -> Tuple[Union[float, Tensor], ...]:
+    if not isinstance(inputs, tuple):
+        assert isinstance(
+            inputs, (torch.Tensor, float)
+        ), "`inputs` must have type float or torch.Tensor but {} found: ".format(
+            type(inputs)
+        )
+        inputs = (inputs,)
+    return inputs
+
+
+@overload
+def _format_additional_forward_args(additional_forward_args: None) -> None:
+    ...
+
+
+@overload
+def _format_additional_forward_args(
+    additional_forward_args: Union[Tensor, Tuple]
+) -> Tuple:
+    ...
+
+
+@overload
+def _format_additional_forward_args(additional_forward_args: Any) -> Union[None, Tuple]:
+    ...
+
+
+def _format_additional_forward_args(additional_forward_args: Any) -> Union[None, Tuple]:
+    if additional_forward_args is not None and not isinstance(
+        additional_forward_args, tuple
+    ):
+        additional_forward_args = (additional_forward_args,)
+    return additional_forward_args
+
+
+def _expand_additional_forward_args(
+    additional_forward_args: Any,
+    n_steps: int,
+    expansion_type: ExpansionTypes = ExpansionTypes.repeat,
+) -> Union[None, Tuple]:
+    def _expand_tensor_forward_arg(
+        additional_forward_arg: Tensor,
+        n_steps: int,
+        expansion_type: ExpansionTypes = ExpansionTypes.repeat,
+    ) -> Tensor:
+        if len(additional_forward_arg.size()) == 0:
+            return additional_forward_arg
+        if expansion_type == ExpansionTypes.repeat:
+            return torch.cat([additional_forward_arg] * n_steps, dim=0)
+        elif expansion_type == ExpansionTypes.repeat_interleave:
+            return additional_forward_arg.repeat_interleave(n_steps, dim=0)
+        else:
+            raise NotImplementedError(
+                "Currently only `repeat` and `repeat_interleave`"
+                " expansion_types are supported"
+            )
+
+    if additional_forward_args is None:
+        return None
+
+    return tuple(
+        _expand_tensor_forward_arg(additional_forward_arg, n_steps, expansion_type)
+        if isinstance(additional_forward_arg, torch.Tensor)
+        else additional_forward_arg
+        for additional_forward_arg in additional_forward_args
+    )
+
+
+def _expand_target(
+    target: TargetType,
+    n_steps: int,
+    expansion_type: ExpansionTypes = ExpansionTypes.repeat,
+) -> TargetType:
+    if isinstance(target, list):
+        if expansion_type == ExpansionTypes.repeat:
+            return target * n_steps
+        elif expansion_type == ExpansionTypes.repeat_interleave:
+            expanded_target = []
+            for i in target:
+                expanded_target.extend([i] * n_steps)
+            return cast(Union[List[Tuple[int, ...]], List[int]], expanded_target)
+        else:
+            raise NotImplementedError(
+                "Currently only `repeat` and `repeat_interleave`"
+                " expansion_types are supported"
+            )
+
+    elif isinstance(target, torch.Tensor) and torch.numel(target) > 1:
+        if expansion_type == ExpansionTypes.repeat:
+            return torch.cat([target] * n_steps, dim=0)
+        elif expansion_type == ExpansionTypes.repeat_interleave:
+            return target.repeat_interleave(n_steps, dim=0)
+        else:
+            raise NotImplementedError(
+                "Currently only `repeat` and `repeat_interleave`"
+                " expansion_types are supported"
+            )
+
+    return target
+
+
+def _expand_feature_mask(
+    feature_mask: Union[Tensor, Tuple[Tensor, ...]], n_samples: int
+):
+    is_feature_mask_tuple = _is_tuple(feature_mask)
+    feature_mask = _format_tensor_into_tuples(feature_mask)
+    feature_mask_new = tuple(
+        feature_mask_elem.repeat_interleave(n_samples, dim=0)
+        if feature_mask_elem.size(0) > 1
+        else feature_mask_elem
+        for feature_mask_elem in feature_mask
+    )
+    return _format_output(is_feature_mask_tuple, feature_mask_new)
+
+
+def _expand_and_update_baselines(
+    inputs: Tuple[Tensor, ...],
+    n_samples: int,
+    kwargs: dict,
+    draw_baseline_from_distrib: bool = False,
+):
+    def get_random_baseline_indices(bsz, baseline):
+        num_ref_samples = baseline.shape[0]
+        return np.random.choice(num_ref_samples, n_samples * bsz).tolist()
+
+    # expand baselines to match the sizes of input
+    if "baselines" not in kwargs:
+        return
+
+    baselines = kwargs["baselines"]
+    baselines = _format_baseline(baselines, inputs)
+    _validate_input(
+        inputs, baselines, draw_baseline_from_distrib=draw_baseline_from_distrib
+    )
+
+    if draw_baseline_from_distrib:
+        bsz = inputs[0].shape[0]
+        baselines = tuple(
+            baseline[get_random_baseline_indices(bsz, baseline)]
+            if isinstance(baseline, torch.Tensor)
+            else baseline
+            for baseline in baselines
+        )
+    else:
+        baselines = tuple(
+            baseline.repeat_interleave(n_samples, dim=0)
+            if isinstance(baseline, torch.Tensor)
+            and baseline.shape[0] == input.shape[0]
+            and baseline.shape[0] > 1
+            else baseline
+            for input, baseline in zip(inputs, baselines)
+        )
+    # update kwargs with expanded baseline
+    kwargs["baselines"] = baselines
+
+
+def _expand_and_update_additional_forward_args(n_samples: int, kwargs: dict):
+    if "additional_forward_args" not in kwargs:
+        return
+    additional_forward_args = kwargs["additional_forward_args"]
+    additional_forward_args = _format_additional_forward_args(additional_forward_args)
+    if additional_forward_args is None:
+        return
+    additional_forward_args = _expand_additional_forward_args(
+        additional_forward_args,
+        n_samples,
+        expansion_type=ExpansionTypes.repeat_interleave,
+    )
+    # update kwargs with expanded baseline
+    kwargs["additional_forward_args"] = additional_forward_args
+
+
+def _expand_and_update_target(n_samples: int, kwargs: dict):
+    if "target" not in kwargs:
+        return
+    target = kwargs["target"]
+    target = _expand_target(
+        target, n_samples, expansion_type=ExpansionTypes.repeat_interleave
+    )
+    # update kwargs with expanded baseline
+    kwargs["target"] = target
+
+
+def _expand_and_update_feature_mask(n_samples: int, kwargs: dict):
+    if "feature_mask" not in kwargs:
+        return
+
+    feature_mask = kwargs["feature_mask"]
+    if feature_mask is None:
+        return
+
+    feature_mask = _expand_feature_mask(feature_mask, n_samples)
+    kwargs["feature_mask"] = feature_mask
+
+
+@typing.overload
+def _format_output(
+    is_inputs_tuple: Literal[True], output: Tuple[Tensor, ...]
+) -> Tuple[Tensor, ...]:
+    ...
+
+
+@typing.overload
+def _format_output(
+    is_inputs_tuple: Literal[False], output: Tuple[Tensor, ...]
+) -> Tensor:
+    ...
+
+
+@typing.overload
+def _format_output(
+    is_inputs_tuple: bool, output: Tuple[Tensor, ...]
+) -> Union[Tensor, Tuple[Tensor, ...]]:
+    ...
+
+
+def _format_output(
+    is_inputs_tuple: bool, output: Tuple[Tensor, ...]
+) -> Union[Tensor, Tuple[Tensor, ...]]:
+    r"""
+    In case input is a tensor and the output is returned in form of a
+    tuple we take the first element of the output's tuple to match the
+    same shape signatues of the inputs
+    """
+    assert isinstance(output, tuple), "Output must be in shape of a tuple"
+    assert is_inputs_tuple or len(output) == 1, (
+        "The input is a single tensor however the output isn't."
+        "The number of output tensors is: {}".format(len(output))
+    )
+    return output if is_inputs_tuple else output[0]
+
+
+@typing.overload
+def _format_outputs(
+    is_multiple_inputs: Literal[False], outputs: List[Tuple[Tensor, ...]]
+) -> Union[Tensor, Tuple[Tensor, ...]]:
+    ...
+
+
+@typing.overload
+def _format_outputs(
+    is_multiple_inputs: Literal[True], outputs: List[Tuple[Tensor, ...]]
+) -> List[Union[Tensor, Tuple[Tensor, ...]]]:
+    ...
+
+
+@typing.overload
+def _format_outputs(
+    is_multiple_inputs: bool, outputs: List[Tuple[Tensor, ...]]
+) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
+    ...
+
+
+def _format_outputs(
+    is_multiple_inputs: bool, outputs: List[Tuple[Tensor, ...]]
+) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
+    assert isinstance(outputs, list), "Outputs must be a list"
+    assert is_multiple_inputs or len(outputs) == 1, (
+        "outputs should contain multiple inputs or have a single output"
+        f"however the number of outputs is: {len(outputs)}"
+    )
+
+    return (
+        [_format_output(len(output) > 1, output) for output in outputs]
+        if is_multiple_inputs
+        else _format_output(len(outputs[0]) > 1, outputs[0])
+    )
+
+
+def _run_forward(
+    forward_func: Callable,
+    inputs: Any,
+    target: TargetType = None,
+    additional_forward_args: Any = None,
+) -> Tensor:
+    forward_func_args = signature(forward_func).parameters
+    if len(forward_func_args) == 0:
+        output = forward_func()
+        return output if target is None else _select_targets(output, target)
+
+    # make everything a tuple so that it is easy to unpack without
+    # using if-statements
+    inputs = _format_inputs(inputs)
+    additional_forward_args = _format_additional_forward_args(additional_forward_args)
+
+    output = forward_func(
+        *(*inputs, *additional_forward_args)
+        if additional_forward_args is not None
+        else inputs
+    )
+    return _select_targets(output, target)
+
+
+def _select_targets(output: Tensor, target: TargetType) -> Tensor:
+    if target is None:
+        return output
+
+    num_examples = output.shape[0]
+    dims = len(output.shape)
+    device = output.device
+    if isinstance(target, (int, tuple)):
+        return _verify_select_column(output, target)
+    elif isinstance(target, torch.Tensor):
+        if torch.numel(target) == 1 and isinstance(target.item(), int):
+            return _verify_select_column(output, cast(int, target.item()))
+        elif len(target.shape) == 1 and torch.numel(target) == num_examples:
+            assert dims == 2, "Output must be 2D to select tensor of targets."
+            return torch.gather(output, 1, target.reshape(len(output), 1))
+        else:
+            raise AssertionError(
+                "Tensor target dimension %r is not valid. %r"
+                % (target.shape, output.shape)
+            )
+    elif isinstance(target, list):
+        assert len(target) == num_examples, "Target list length does not match output!"
+        if isinstance(target[0], int):
+            assert dims == 2, "Output must be 2D to select tensor of targets."
+            return torch.gather(
+                output, 1, torch.tensor(target, device=device).reshape(len(output), 1)
+            )
+        elif isinstance(target[0], tuple):
+            return torch.stack(
+                [
+                    output[(i,) + cast(Tuple, targ_elem)]
+                    for i, targ_elem in enumerate(target)
+                ]
+            )
+        else:
+            raise AssertionError("Target element type in list is not valid.")
+    else:
+        raise AssertionError("Target type %r is not valid." % target)
+
+
+def _contains_slice(target: Union[int, Tuple[Union[int, slice], ...]]) -> bool:
+    if isinstance(target, tuple):
+        for index in target:
+            if isinstance(index, slice):
+                return True
+        return False
+    return isinstance(target, slice)
+
+
+def _verify_select_column(
+    output: Tensor, target: Union[int, Tuple[Union[int, slice], ...]]
+) -> Tensor:
+    target = (target,) if isinstance(target, int) else target
+    assert (
+        len(target) <= len(output.shape) - 1
+    ), "Cannot choose target column with output shape %r." % (output.shape,)
+    return output[(slice(None), *target)]
+
+
+def _verify_select_neuron(
+    layer_output: Tuple[Tensor, ...],
+    selector: Union[int, Tuple[Union[int, slice], ...], Callable],
+) -> Tensor:
+    if callable(selector):
+        return selector(layer_output if len(layer_output) > 1 else layer_output[0])
+
+    assert len(layer_output) == 1, (
+        "Cannot select neuron index from layer with multiple tensors,"
+        "consider providing a neuron selector function instead."
+    )
+
+    selected_neurons = _verify_select_column(layer_output[0], selector)
+    if _contains_slice(selector):
+        return selected_neurons.reshape(selected_neurons.shape[0], -1).sum(1)
+    return selected_neurons
+
+
+def _extract_device(
+    module: Module,
+    hook_inputs: Union[None, Tensor, Tuple[Tensor, ...]],
+    hook_outputs: Union[None, Tensor, Tuple[Tensor, ...]],
+) -> device:
+    params = list(module.parameters())
+    if (
+        (hook_inputs is None or len(hook_inputs) == 0)
+        and (hook_outputs is None or len(hook_outputs) == 0)
+        and len(params) == 0
+    ):
+        raise RuntimeError(
+            """Unable to extract device information for the module
+            {}. Both inputs and outputs to the forward hook and
+            `module.parameters()` are empty.
+            The reason that the inputs to the forward hook are empty
+            could be due to the fact that the arguments to that
+            module {} are all named and are passed as named
+            variables to its forward function.
+            """.format(
+                module, module
+            )
+        )
+    if hook_inputs is not None and len(hook_inputs) > 0:
+        return hook_inputs[0].device
+    if hook_outputs is not None and len(hook_outputs) > 0:
+        return hook_outputs[0].device
+
+    return params[0].device
+
+
+def _reduce_list(
+    val_list: List[TupleOrTensorOrBoolGeneric],
+    red_func: Callable[[List], Any] = torch.cat,
+) -> TupleOrTensorOrBoolGeneric:
+    """
+    Applies reduction function to given list. If each element in the list is
+    a Tensor, applies reduction function to all elements of the list, and returns
+    the output Tensor / value. If each element is a boolean, apply any method (or).
+    If each element is a tuple, applies reduction
+    function to corresponding elements of each tuple in the list, and returns
+    tuple of reduction function outputs with length matching the length of tuple
+    val_list[0]. It is assumed that all tuples in the list have the same length
+    and red_func can be applied to all elements in each corresponding position.
+    """
+    assert len(val_list) > 0, "Cannot reduce empty list!"
+    if isinstance(val_list[0], torch.Tensor):
+        first_device = val_list[0].device
+        return red_func([elem.to(first_device) for elem in val_list])
+    elif isinstance(val_list[0], bool):
+        return any(val_list)
+    elif isinstance(val_list[0], tuple):
+        final_out = []
+        for i in range(len(val_list[0])):
+            final_out.append(
+                _reduce_list([val_elem[i] for val_elem in val_list], red_func)
+            )
+    else:
+        raise AssertionError(
+            "Elements to be reduced can only be"
+            "either Tensors or tuples containing Tensors."
+        )
+    return tuple(final_out)
+
+
+def _sort_key_list(
+    keys: List[device], device_ids: Union[None, List[int]] = None
+) -> List[device]:
+    """
+    Sorts list of torch devices (keys) by given index list, device_ids. If keys
+    contains only one device, then the list is returned unchanged. If keys
+    contains a device for which the id is not contained in device_ids, then
+    an error is returned. This method is used to identify the order of DataParallel
+    batched devices, given the device ID ordering.
+    """
+    if len(keys) == 1:
+        return keys
+    id_dict: Dict[int, device] = {}
+    assert device_ids is not None, "Device IDs must be provided with multiple devices."
+    for key in keys:
+        if key.index in id_dict:
+            raise AssertionError("Duplicate CUDA Device ID identified in device list.")
+        id_dict[key.index] = key
+
+    out_list = [
+        id_dict[device_id]
+        for device_id in filter(lambda device_id: device_id in id_dict, device_ids)
+    ]
+
+    assert len(out_list) == len(keys), "Given Device ID List does not match"
+    "devices with computed tensors."
+
+    return out_list
+
+
+def _flatten_tensor_or_tuple(inp: TensorOrTupleOfTensorsGeneric) -> Tensor:
+    if isinstance(inp, Tensor):
+        return inp.flatten()
+    return torch.cat([single_inp.flatten() for single_inp in inp])
+
+
+def _get_module_from_name(model: Module, layer_name: str) -> Any:
+    r"""
+    Returns the module (layer) object, given its (string) name
+    in the model.
+
+    Args:
+            name (str): Module or nested modules name string in self.model
+
+    Returns:
+            The module (layer) in self.model.
+    """
+
+    return reduce(getattr, layer_name.split("."), model)
+
+
+def _register_backward_hook(
+    module: Module, hook: Callable, attr_obj: Any
+) -> torch.utils.hooks.RemovableHandle:
+    # Special case for supporting output attributions for neuron methods
+    # This can be removed after deprecation of neuron output attributions
+    # for NeuronDeepLift, NeuronDeconvolution, and NeuronGuidedBackprop
+    # in v0.6.0
+    if (
+        hasattr(attr_obj, "skip_new_hook_layer")
+        and attr_obj.skip_new_hook_layer == module
+    ):
+        return module.register_backward_hook(hook)
+
+    if torch.__version__ >= "1.9":
+        # Only supported for torch >= 1.9
+        return module.register_full_backward_hook(hook)
+    else:
+        # Fallback for previous versions of PyTorch
+        return module.register_backward_hook(hook)

captum/_utils/gradient.py

@@ -0,0 +1,865 @@
+#!/usr/bin/env python3
+import threading
+import typing
+import warnings
+from collections import defaultdict
+from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _reduce_list,
+    _run_forward,
+    _sort_key_list,
+    _verify_select_neuron,
+)
+from captum._utils.sample_gradient import SampleGradientWrapper
+from captum._utils.typing import (
+    Literal,
+    ModuleOrModuleList,
+    TargetType,
+    TensorOrTupleOfTensorsGeneric,
+)
+from torch import device, Tensor
+from torch.nn import Module
+
+
+def apply_gradient_requirements(
+    inputs: Tuple[Tensor, ...], warn: bool = True
+) -> List[bool]:
+    """
+    Iterates through tuple on input tensors and sets requires_grad to be true on
+    each Tensor, and ensures all grads are set to zero. To ensure that the input
+    is returned to its initial state, a list of flags representing whether or not
+     a tensor originally required grad is returned.
+    """
+    assert isinstance(
+        inputs, tuple
+    ), "Inputs should be wrapped in a tuple prior to preparing for gradients"
+    grad_required = []
+    for index, input in enumerate(inputs):
+        assert isinstance(input, torch.Tensor), "Given input is not a torch.Tensor"
+        grad_required.append(input.requires_grad)
+        inputs_dtype = input.dtype
+        # Note: torch 1.2 doesn't support is_complex for dtype that's why we check
+        # on the existance of is_complex method.
+        if not inputs_dtype.is_floating_point and not (
+            hasattr(inputs_dtype, "is_complex") and inputs_dtype.is_complex
+        ):
+            if warn:
+                warnings.warn(
+                    """Input Tensor %d has a dtype of %s.
+                    Gradients cannot be activated
+                    for these data types."""
+                    % (index, str(inputs_dtype))
+                )
+        elif not input.requires_grad:
+            if warn:
+                warnings.warn(
+                    "Input Tensor %d did not already require gradients, "
+                    "required_grads has been set automatically." % index
+                )
+            input.requires_grad_()
+    return grad_required
+
+
+def undo_gradient_requirements(
+    inputs: Tuple[Tensor, ...], grad_required: List[bool]
+) -> None:
+    """
+    Iterates through list of tensors, zeros each gradient, and sets required
+    grad to false if the corresponding index in grad_required is False.
+    This method is used to undo the effects of prepare_gradient_inputs, making
+    grads not required for any input tensor that did not initially require
+    gradients.
+    """
+
+    assert isinstance(
+        inputs, tuple
+    ), "Inputs should be wrapped in a tuple prior to preparing for gradients."
+    assert len(inputs) == len(
+        grad_required
+    ), "Input tuple length should match gradient mask."
+    for index, input in enumerate(inputs):
+        assert isinstance(input, torch.Tensor), "Given input is not a torch.Tensor"
+        if not grad_required[index]:
+            input.requires_grad_(False)
+
+
+def compute_gradients(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+) -> Tuple[Tensor, ...]:
+    r"""
+    Computes gradients of the output with respect to inputs for an
+    arbitrary forward function.
+
+    Args:
+
+        forward_fn: forward function. This can be for example model's
+                    forward function.
+        input:      Input at which gradients are evaluated,
+                    will be passed to forward_fn.
+        target_ind: Index of the target class for which gradients
+                    must be computed (classification only).
+        additional_forward_args: Additional input arguments that forward
+                    function requires. It takes an empty tuple (no additional
+                    arguments) if no additional arguments are required
+    """
+    with torch.autograd.set_grad_enabled(True):
+        # runs forward pass
+        outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)
+        assert outputs[0].numel() == 1, (
+            "Target not provided when necessary, cannot"
+            " take gradient with respect to multiple outputs."
+        )
+        # torch.unbind(forward_out) is a list of scalar tensor tuples and
+        # contains batch_size * #steps elements
+        grads = torch.autograd.grad(torch.unbind(outputs), inputs)
+    return grads
+
+
+def _neuron_gradients(
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    saved_layer: Dict[device, Tuple[Tensor, ...]],
+    key_list: List[device],
+    gradient_neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
+) -> Tuple[Tensor, ...]:
+    with torch.autograd.set_grad_enabled(True):
+        gradient_tensors = []
+        for key in key_list:
+            current_out_tensor = _verify_select_neuron(
+                saved_layer[key], gradient_neuron_selector
+            )
+            gradient_tensors.append(
+                torch.autograd.grad(
+                    torch.unbind(current_out_tensor)
+                    if current_out_tensor.numel() > 1
+                    else current_out_tensor,
+                    inputs,
+                )
+            )
+        _total_gradients = _reduce_list(gradient_tensors, sum)
+    return _total_gradients
+
+
+@typing.overload
+def _forward_layer_eval(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: Module,
+    additional_forward_args: Any = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    grad_enabled: bool = False,
+) -> Tuple[Tensor, ...]:
+    ...
+
+
+@typing.overload
+def _forward_layer_eval(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: List[Module],
+    additional_forward_args: Any = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    grad_enabled: bool = False,
+) -> List[Tuple[Tensor, ...]]:
+    ...
+
+
+def _forward_layer_eval(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: ModuleOrModuleList,
+    additional_forward_args: Any = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    grad_enabled: bool = False,
+) -> Union[Tuple[Tensor, ...], List[Tuple[Tensor, ...]]]:
+    return _forward_layer_eval_with_neuron_grads(
+        forward_fn,
+        inputs,
+        layer,
+        additional_forward_args=additional_forward_args,
+        gradient_neuron_selector=None,
+        grad_enabled=grad_enabled,
+        device_ids=device_ids,
+        attribute_to_layer_input=attribute_to_layer_input,
+    )
+
+
+@typing.overload
+def _forward_layer_distributed_eval(
+    forward_fn: Callable,
+    inputs: Any,
+    layer: ModuleOrModuleList,
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    attribute_to_layer_input: bool = False,
+    forward_hook_with_return: Literal[False] = False,
+    require_layer_grads: bool = False,
+) -> Dict[Module, Dict[device, Tuple[Tensor, ...]]]:
+    ...
+
+
+@typing.overload
+def _forward_layer_distributed_eval(
+    forward_fn: Callable,
+    inputs: Any,
+    layer: ModuleOrModuleList,
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    attribute_to_layer_input: bool = False,
+    *,
+    forward_hook_with_return: Literal[True],
+    require_layer_grads: bool = False,
+) -> Tuple[Dict[Module, Dict[device, Tuple[Tensor, ...]]], Tensor]:
+    ...
+
+
+def _forward_layer_distributed_eval(
+    forward_fn: Callable,
+    inputs: Any,
+    layer: ModuleOrModuleList,
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    attribute_to_layer_input: bool = False,
+    forward_hook_with_return: bool = False,
+    require_layer_grads: bool = False,
+) -> Union[
+    Tuple[Dict[Module, Dict[device, Tuple[Tensor, ...]]], Tensor],
+    Dict[Module, Dict[device, Tuple[Tensor, ...]]],
+]:
+    r"""
+    A helper function that allows to set a hook on model's `layer`, run the forward
+    pass and returns intermediate layer results, stored in a dictionary,
+    and optionally also the output of the forward function. The keys in the
+    dictionary are the device ids and the values are corresponding intermediate layer
+    results, either the inputs or the outputs of the layer depending on whether we set
+    `attribute_to_layer_input` to True or False.
+    This is especially useful when we execute forward pass in a distributed setting,
+    using `DataParallel`s for example.
+    """
+    saved_layer: Dict[Module, Dict[device, Tuple[Tensor, ...]]] = defaultdict(dict)
+    lock = threading.Lock()
+    all_layers: List[Module] = [layer] if isinstance(layer, Module) else layer
+
+    # Set a forward hook on specified module and run forward pass to
+    # get layer output tensor(s).
+    # For DataParallel models, each partition adds entry to dictionary
+    # with key as device and value as corresponding Tensor.
+    def hook_wrapper(original_module):
+        def forward_hook(module, inp, out=None):
+            eval_tsrs = inp if attribute_to_layer_input else out
+            is_eval_tuple = isinstance(eval_tsrs, tuple)
+
+            if not is_eval_tuple:
+                eval_tsrs = (eval_tsrs,)
+            if require_layer_grads:
+                apply_gradient_requirements(eval_tsrs, warn=False)
+            with lock:
+                nonlocal saved_layer
+                # Note that cloning behaviour of `eval_tsr` is different
+                # when `forward_hook_with_return` is set to True. This is because
+                # otherwise `backward()` on the last output layer won't execute.
+                if forward_hook_with_return:
+                    saved_layer[original_module][eval_tsrs[0].device] = eval_tsrs
+                    eval_tsrs_to_return = tuple(
+                        eval_tsr.clone() for eval_tsr in eval_tsrs
+                    )
+                    if not is_eval_tuple:
+                        eval_tsrs_to_return = eval_tsrs_to_return[0]
+                    return eval_tsrs_to_return
+                else:
+                    saved_layer[original_module][eval_tsrs[0].device] = tuple(
+                        eval_tsr.clone() for eval_tsr in eval_tsrs
+                    )
+
+        return forward_hook
+
+    all_hooks = []
+    try:
+        for single_layer in all_layers:
+            if attribute_to_layer_input:
+                all_hooks.append(
+                    single_layer.register_forward_pre_hook(hook_wrapper(single_layer))
+                )
+            else:
+                all_hooks.append(
+                    single_layer.register_forward_hook(hook_wrapper(single_layer))
+                )
+        output = _run_forward(
+            forward_fn,
+            inputs,
+            target=target_ind,
+            additional_forward_args=additional_forward_args,
+        )
+    finally:
+        for hook in all_hooks:
+            hook.remove()
+
+    if len(saved_layer) == 0:
+        raise AssertionError("Forward hook did not obtain any outputs for given layer")
+
+    if forward_hook_with_return:
+        return saved_layer, output
+    return saved_layer
+
+
+def _gather_distributed_tensors(
+    saved_layer: Dict[device, Tuple[Tensor, ...]],
+    device_ids: Union[None, List[int]] = None,
+    key_list: Union[None, List[device]] = None,
+) -> Tuple[Tensor, ...]:
+    r"""
+    A helper function to concatenate intermediate layer results stored on
+    different devices in `saved_layer`. `saved_layer` is a dictionary that
+    contains `device_id` as a key and intermediate layer results (either
+    the input or the output of the layer) stored on the device corresponding to
+    the key.
+    `key_list` is a list of devices in appropriate ordering for concatenation
+    and if not provided, keys are sorted based on device ids.
+
+    If only one key exists (standard model), key list simply has one element.
+    """
+    if key_list is None:
+        key_list = _sort_key_list(list(saved_layer.keys()), device_ids)
+    return _reduce_list([saved_layer[device_id] for device_id in key_list])
+
+
+def _extract_device_ids(
+    forward_fn: Callable,
+    saved_layer: Dict[Module, Dict[device, Tuple[Tensor, ...]]],
+    device_ids: Union[None, List[int]],
+) -> Union[None, List[int]]:
+    r"""
+    A helper function to extract device_ids from `forward_function` in case it is
+    provided as part of a `DataParallel` model or if is accessible from
+    `forward_fn`.
+    In case input device_ids is not None, this function returns that value.
+    """
+    # Multiple devices / keys implies a DataParallel model, so we look for
+    # device IDs if given or available from forward function
+    # (DataParallel model object).
+    if (
+        max(len(saved_layer[single_layer]) for single_layer in saved_layer) > 1
+        and device_ids is None
+    ):
+        if (
+            hasattr(forward_fn, "device_ids")
+            and cast(Any, forward_fn).device_ids is not None
+        ):
+            device_ids = cast(Any, forward_fn).device_ids
+        else:
+            raise AssertionError(
+                "Layer tensors are saved on multiple devices, however unable to access"
+                " device ID list from the `forward_fn`. Device ID list must be"
+                " accessible from `forward_fn`. For example, they can be retrieved"
+                " if `forward_fn` is a model of type `DataParallel`. It is used"
+                " for identifying device batch ordering."
+            )
+    return device_ids
+
+
+@typing.overload
+def _forward_layer_eval_with_neuron_grads(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: Module,
+    additional_forward_args: Any = None,
+    *,
+    gradient_neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
+    grad_enabled: bool = False,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]]:
+    ...
+
+
+@typing.overload
+def _forward_layer_eval_with_neuron_grads(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: Module,
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: None = None,
+    grad_enabled: bool = False,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+) -> Tuple[Tensor, ...]:
+    ...
+
+
+@typing.overload
+def _forward_layer_eval_with_neuron_grads(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: List[Module],
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: None = None,
+    grad_enabled: bool = False,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+) -> List[Tuple[Tensor, ...]]:
+    ...
+
+
+def _forward_layer_eval_with_neuron_grads(
+    forward_fn: Callable,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    layer: ModuleOrModuleList,
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: Union[
+        None, int, Tuple[Union[int, slice], ...], Callable
+    ] = None,
+    grad_enabled: bool = False,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+) -> Union[
+    Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]],
+    Tuple[Tensor, ...],
+    List[Tuple[Tensor, ...]],
+]:
+    """
+    This method computes forward evaluation for a particular layer using a
+    forward hook. If a gradient_neuron_selector is provided, then gradients with
+    respect to that neuron in the layer output are also returned.
+
+    These functionalities are combined due to the behavior of DataParallel models
+    with hooks, in which hooks are executed once per device. We need to internally
+    combine the separated tensors from devices by concatenating based on device_ids.
+    Any necessary gradients must be taken with respect to each independent batched
+    tensor, so the gradients are computed and combined appropriately.
+
+    More information regarding the behavior of forward hooks with DataParallel models
+    can be found in the PyTorch data parallel documentation. We maintain the separate
+    evals in a dictionary protected by a lock, analogous to the gather implementation
+    for the core PyTorch DataParallel implementation.
+    """
+    grad_enabled = True if gradient_neuron_selector is not None else grad_enabled
+
+    with torch.autograd.set_grad_enabled(grad_enabled):
+        saved_layer = _forward_layer_distributed_eval(
+            forward_fn,
+            inputs,
+            layer,
+            additional_forward_args=additional_forward_args,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+    device_ids = _extract_device_ids(forward_fn, saved_layer, device_ids)
+    # Identifies correct device ordering based on device ids.
+    # key_list is a list of devices in appropriate ordering for concatenation.
+    # If only one key exists (standard model), key list simply has one element.
+    key_list = _sort_key_list(list(next(iter(saved_layer.values())).keys()), device_ids)
+    if gradient_neuron_selector is not None:
+        assert isinstance(
+            layer, Module
+        ), "Cannot compute neuron gradients for multiple layers simultaneously!"
+        inp_grads = _neuron_gradients(
+            inputs, saved_layer[layer], key_list, gradient_neuron_selector
+        )
+        return (
+            _gather_distributed_tensors(saved_layer[layer], key_list=key_list),
+            inp_grads,
+        )
+    else:
+        if isinstance(layer, Module):
+            return _gather_distributed_tensors(saved_layer[layer], key_list=key_list)
+        else:
+            return [
+                _gather_distributed_tensors(saved_layer[curr_layer], key_list=key_list)
+                for curr_layer in layer
+            ]
+
+
+@typing.overload
+def compute_layer_gradients_and_eval(
+    forward_fn: Callable,
+    layer: Module,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    *,
+    gradient_neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    output_fn: Union[None, Callable] = None,
+) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...]]:
+    ...
+
+
+@typing.overload
+def compute_layer_gradients_and_eval(
+    forward_fn: Callable,
+    layer: List[Module],
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: None = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    output_fn: Union[None, Callable] = None,
+) -> Tuple[List[Tuple[Tensor, ...]], List[Tuple[Tensor, ...]]]:
+    ...
+
+
+@typing.overload
+def compute_layer_gradients_and_eval(
+    forward_fn: Callable,
+    layer: Module,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: None = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    output_fn: Union[None, Callable] = None,
+) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]]:
+    ...
+
+
+def compute_layer_gradients_and_eval(
+    forward_fn: Callable,
+    layer: ModuleOrModuleList,
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
+    target_ind: TargetType = None,
+    additional_forward_args: Any = None,
+    gradient_neuron_selector: Union[
+        None, int, Tuple[Union[int, slice], ...], Callable
+    ] = None,
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_layer_input: bool = False,
+    output_fn: Union[None, Callable] = None,
+) -> Union[
+    Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]],
+    Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...]],
+    Tuple[List[Tuple[Tensor, ...]], List[Tuple[Tensor, ...]]],
+]:
+    r"""
+    Computes gradients of the output with respect to a given layer as well
+    as the output evaluation of the layer for an arbitrary forward function
+    and given input.
+
+    For data parallel models, hooks are executed once per device ,so we
+    need to internally combine the separated tensors from devices by
+    concatenating based on device_ids. Any necessary gradients must be taken
+    with respect to each independent batched tensor, so the gradients are
+    computed and combined appropriately.
+
+    More information regarding the behavior of forward hooks with DataParallel
+    models can be found in the PyTorch data parallel documentation. We maintain
+    the separate inputs in a dictionary protected by a lock, analogous to the
+    gather implementation for the core PyTorch DataParallel implementation.
+
+    NOTE: To properly handle inplace operations, a clone of the layer output
+    is stored. This structure inhibits execution of a backward hook on the last
+    module for the layer output when computing the gradient with respect to
+    the input, since we store an intermediate clone, as
+    opposed to the true module output. If backward module hooks are necessary
+    for the final module when computing input gradients, utilize
+    _forward_layer_eval_with_neuron_grads instead.
+
+    Args:
+
+        forward_fn: forward function. This can be for example model's
+                    forward function.
+        layer:      Layer for which gradients / output will be evaluated.
+        inputs:     Input at which gradients are evaluated,
+                    will be passed to forward_fn.
+        target_ind: Index of the target class for which gradients
+                    must be computed (classification only).
+        output_fn:  An optional function that is applied to the layer inputs or
+                    outputs depending whether the `attribute_to_layer_input` is
+                    set to `True` or `False`
+        args:       Additional input arguments that forward function requires.
+                    It takes an empty tuple (no additional arguments) if no
+                    additional arguments are required
+
+
+    Returns:
+        2-element tuple of **gradients**, **evals**:
+        - **gradients**:
+            Gradients of output with respect to target layer output.
+        - **evals**:
+            Target layer output for given input.
+    """
+    with torch.autograd.set_grad_enabled(True):
+        # saved_layer is a dictionary mapping device to a tuple of
+        # layer evaluations on that device.
+        saved_layer, output = _forward_layer_distributed_eval(
+            forward_fn,
+            inputs,
+            layer,
+            target_ind=target_ind,
+            additional_forward_args=additional_forward_args,
+            attribute_to_layer_input=attribute_to_layer_input,
+            forward_hook_with_return=True,
+            require_layer_grads=True,
+        )
+        assert output[0].numel() == 1, (
+            "Target not provided when necessary, cannot"
+            " take gradient with respect to multiple outputs."
+        )
+
+        device_ids = _extract_device_ids(forward_fn, saved_layer, device_ids)
+
+        # Identifies correct device ordering based on device ids.
+        # key_list is a list of devices in appropriate ordering for concatenation.
+        # If only one key exists (standard model), key list simply has one element.
+        key_list = _sort_key_list(
+            list(next(iter(saved_layer.values())).keys()), device_ids
+        )
+        all_outputs: Union[Tuple[Tensor, ...], List[Tuple[Tensor, ...]]]
+        if isinstance(layer, Module):
+            all_outputs = _reduce_list(
+                [
+                    saved_layer[layer][device_id]
+                    if output_fn is None
+                    else output_fn(saved_layer[layer][device_id])
+                    for device_id in key_list
+                ]
+            )
+        else:
+            all_outputs = [
+                _reduce_list(
+                    [
+                        saved_layer[single_layer][device_id]
+                        if output_fn is None
+                        else output_fn(saved_layer[single_layer][device_id])
+                        for device_id in key_list
+                    ]
+                )
+                for single_layer in layer
+            ]
+        all_layers: List[Module] = [layer] if isinstance(layer, Module) else layer
+        grad_inputs = tuple(
+            layer_tensor
+            for single_layer in all_layers
+            for device_id in key_list
+            for layer_tensor in saved_layer[single_layer][device_id]
+        )
+        saved_grads = torch.autograd.grad(torch.unbind(output), grad_inputs)
+
+        offset = 0
+        all_grads: List[Tuple[Tensor, ...]] = []
+        for single_layer in all_layers:
+            num_tensors = len(next(iter(saved_layer[single_layer].values())))
+            curr_saved_grads = [
+                saved_grads[i : i + num_tensors]
+                for i in range(
+                    offset, offset + len(key_list) * num_tensors, num_tensors
+                )
+            ]
+            offset += len(key_list) * num_tensors
+            if output_fn is not None:
+                curr_saved_grads = [
+                    output_fn(curr_saved_grad) for curr_saved_grad in curr_saved_grads
+                ]
+
+            all_grads.append(_reduce_list(curr_saved_grads))
+
+        layer_grads: Union[Tuple[Tensor, ...], List[Tuple[Tensor, ...]]]
+        layer_grads = all_grads
+        if isinstance(layer, Module):
+            layer_grads = all_grads[0]
+
+        if gradient_neuron_selector is not None:
+            assert isinstance(
+                layer, Module
+            ), "Cannot compute neuron gradients for multiple layers simultaneously!"
+            inp_grads = _neuron_gradients(
+                inputs, saved_layer[layer], key_list, gradient_neuron_selector
+            )
+            return (
+                cast(Tuple[Tensor, ...], layer_grads),
+                cast(Tuple[Tensor, ...], all_outputs),
+                inp_grads,
+            )
+    return layer_grads, all_outputs  # type: ignore
+
+
+def construct_neuron_grad_fn(
+    layer: Module,
+    neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
+    device_ids: Union[None, List[int]] = None,
+    attribute_to_neuron_input: bool = False,
+) -> Callable:
+    def grad_fn(
+        forward_fn: Callable,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target_ind: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> Tuple[Tensor, ...]:
+        _, grads = _forward_layer_eval_with_neuron_grads(
+            forward_fn,
+            inputs,
+            layer,
+            additional_forward_args,
+            gradient_neuron_selector=neuron_selector,
+            device_ids=device_ids,
+            attribute_to_layer_input=attribute_to_neuron_input,
+        )
+        return grads
+
+    return grad_fn
+
+
+def _compute_jacobian_wrt_params(
+    model: Module,
+    inputs: Tuple[Any, ...],
+    labels: Optional[Tensor] = None,
+    loss_fn: Optional[Union[Module, Callable]] = None,
+) -> Tuple[Tensor, ...]:
+    r"""
+    Computes the Jacobian of a batch of test examples given a model, and optional
+    loss function and target labels. This method is equivalent to calculating the
+    gradient for every individual example in the minibatch.
+
+    Args:
+        model (torch.nn.Module): The trainable model providing the forward pass
+        inputs (tuple of Any): The minibatch for which the forward pass is computed.
+                It is unpacked before passing to `model`, so it must be a tuple.  The
+                individual elements of `inputs` can be anything.
+        labels (Tensor or None): Labels for input if computing a loss function.
+        loss_fn (torch.nn.Module or Callable or None): The loss function. If a library
+                defined loss function is provided, it would be expected to be a
+                torch.nn.Module. If a custom loss is provided, it can be either type,
+                but must behave as a library loss function would if `reduction='none'`.
+
+    Returns:
+        grads (Tuple of Tensor): Returns the Jacobian for the minibatch as a
+                tuple of gradients corresponding to the tuple of trainable parameters
+                returned by `model.parameters()`. Each object grads[i] references to the
+                gradients for the parameters in the i-th trainable layer of the model.
+                Each grads[i] object is a tensor with the gradients for the `inputs`
+                batch. For example, grads[i][j] would reference the gradients for the
+                parameters of the i-th layer, for the j-th member of the minibatch.
+    """
+    with torch.autograd.set_grad_enabled(True):
+        out = model(*inputs)
+        assert out.dim() != 0, "Please ensure model output has at least one dimension."
+
+        if labels is not None and loss_fn is not None:
+            loss = loss_fn(out, labels)
+            if hasattr(loss_fn, "reduction"):
+                msg0 = "Please ensure loss_fn.reduction is set to `none`"
+                assert loss_fn.reduction == "none", msg0  # type: ignore
+            else:
+                msg1 = (
+                    "Loss function is applying a reduction. Please ensure "
+                    f"Output shape: {out.shape} and Loss shape: {loss.shape} "
+                    "are matching."
+                )
+                assert loss.dim() != 0, msg1
+                assert out.shape[0] == loss.shape[0], msg1
+            out = loss
+
+        grads_list = [
+            torch.autograd.grad(
+                outputs=out[i],
+                inputs=model.parameters(),  # type: ignore
+                grad_outputs=torch.ones_like(out[i]),
+                retain_graph=True,
+            )
+            for i in range(out.shape[0])
+        ]
+
+        grads = tuple([torch.stack(x) for x in zip(*grads_list)])
+
+        return tuple(grads)
+
+
+def _compute_jacobian_wrt_params_with_sample_wise_trick(
+    model: Module,
+    inputs: Tuple[Any, ...],
+    labels: Optional[Tensor] = None,
+    loss_fn: Optional[Union[Module, Callable]] = None,
+    reduction_type: Optional[str] = "sum",
+) -> Tuple[Any, ...]:
+    r"""
+    Computes the Jacobian of a batch of test examples given a model, and optional
+    loss function and target labels. This method uses sample-wise gradients per
+    batch trick to fully vectorize the Jacobian calculation. Currently, only
+    linear and conv2d layers are supported.
+
+    User must `add_hooks(model)` before calling this function.
+
+    Args:
+        model (torch.nn.Module): The trainable model providing the forward pass
+        inputs (tuple of Any): The minibatch for which the forward pass is computed.
+                It is unpacked before passing to `model`, so it must be a tuple.  The
+                individual elements of `inputs` can be anything.
+        labels (Tensor or None): Labels for input if computing a loss function.
+        loss_fn (torch.nn.Module or Callable or None): The loss function. If a library
+                defined loss function is provided, it would be expected to be a
+                torch.nn.Module. If a custom loss is provided, it can be either type,
+                but must behave as a library loss function would if `reduction='sum'` or
+                `reduction='mean'`.
+        reduction_type (str): The type of reduction applied. If a loss_fn is passed,
+                this should match `loss_fn.reduction`. Else if gradients are being
+                computed on direct model outputs (scores), then 'sum' should be used.
+                Defaults to 'sum'.
+
+    Returns:
+        grads (Tuple of Tensor): Returns the Jacobian for the minibatch as a
+                tuple of gradients corresponding to the tuple of trainable parameters
+                returned by `model.parameters()`. Each object grads[i] references to the
+                gradients for the parameters in the i-th trainable layer of the model.
+                Each grads[i] object is a tensor with the gradients for the `inputs`
+                batch. For example, grads[i][j] would reference the gradients for the
+                parameters of the i-th layer, for the j-th member of the minibatch.
+    """
+    with torch.autograd.set_grad_enabled(True):
+        sample_grad_wrapper = SampleGradientWrapper(model)
+        try:
+            sample_grad_wrapper.add_hooks()
+
+            out = model(*inputs)
+            assert (
+                out.dim() != 0
+            ), "Please ensure model output has at least one dimension."
+
+            if labels is not None and loss_fn is not None:
+                loss = loss_fn(out, labels)
+                # TODO: allow loss_fn to be Callable
+                if isinstance(loss_fn, Module) and hasattr(loss_fn, "reduction"):
+                    msg0 = (
+                        "Please ensure that loss_fn.reduction is set to `sum` or `mean`"
+                    )
+
+                    assert loss_fn.reduction != "none", msg0
+                    msg1 = (
+                        f"loss_fn.reduction ({loss_fn.reduction}) does not match"
+                        f"reduction type ({reduction_type}). Please ensure they are"
+                        " matching."
+                    )
+                    assert loss_fn.reduction == reduction_type, msg1
+                msg2 = (
+                    "Please ensure custom loss function is applying either a "
+                    "sum or mean reduction."
+                )
+                assert out.shape != loss.shape, msg2
+
+                if reduction_type != "sum" and reduction_type != "mean":
+                    raise ValueError(
+                        f"{reduction_type} is not a valid value for reduction_type. "
+                        "Must be either 'sum' or 'mean'."
+                    )
+                out = loss
+
+            sample_grad_wrapper.compute_param_sample_gradients(
+                out, loss_mode=reduction_type
+            )
+
+            grads = tuple(
+                param.sample_grad  # type: ignore
+                for param in model.parameters()
+                if hasattr(param, "sample_grad")
+            )
+        finally:
+            sample_grad_wrapper.remove_hooks()
+
+        return grads

captum/_utils/models/__init__.py

@@ -0,0 +1,25 @@
+from captum._utils.models.linear_model import (
+    LinearModel,
+    SGDLasso,
+    SGDLinearModel,
+    SGDLinearRegression,
+    SGDRidge,
+    SkLearnLasso,
+    SkLearnLinearModel,
+    SkLearnLinearRegression,
+    SkLearnRidge,
+)
+from captum._utils.models.model import Model
+
+__all__ = [
+    "Model",
+    "LinearModel",
+    "SGDLinearModel",
+    "SGDLasso",
+    "SGDRidge",
+    "SGDLinearRegression",
+    "SkLearnLinearModel",
+    "SkLearnLasso",
+    "SkLearnRidge",
+    "SkLearnLinearRegression",
+]

captum/_utils/models/linear_model/__init__.py

@@ -0,0 +1,23 @@
+from captum._utils.models.linear_model.model import (
+    LinearModel,
+    SGDLasso,
+    SGDLinearModel,
+    SGDLinearRegression,
+    SGDRidge,
+    SkLearnLasso,
+    SkLearnLinearModel,
+    SkLearnLinearRegression,
+    SkLearnRidge,
+)
+
+__all__ = [
+    "LinearModel",
+    "SGDLinearModel",
+    "SGDLasso",
+    "SGDRidge",
+    "SGDLinearRegression",
+    "SkLearnLinearModel",
+    "SkLearnLasso",
+    "SkLearnRidge",
+    "SkLearnLinearRegression",
+]

captum/_utils/models/linear_model/model.py

@@ -0,0 +1,341 @@
+from typing import Callable, cast, List, Optional
+
+import torch.nn as nn
+from captum._utils.models.model import Model
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+
+class LinearModel(nn.Module, Model):
+    SUPPORTED_NORMS: List[Optional[str]] = [None, "batch_norm", "layer_norm"]
+
+    def __init__(self, train_fn: Callable, **kwargs) -> None:
+        r"""
+        Constructs a linear model with a training function and additional
+        construction arguments that will be sent to
+        `self._construct_model_params` after a `self.fit` is called. Please note
+        that this assumes the `self.train_fn` will call
+        `self._construct_model_params`.
+
+        Please note that this is an experimental feature.
+
+        Args:
+            train_fn (callable)
+                The function to train with. See
+                `captum._utils.models.linear_model.train.sgd_train_linear_model`
+                and
+                `captum._utils.models.linear_model.train.sklearn_train_linear_model`
+                for examples
+            kwargs
+                Any additional keyword arguments to send to
+                `self._construct_model_params` once a `self.fit` is called.
+        """
+        super().__init__()
+
+        self.norm: Optional[nn.Module] = None
+        self.linear: Optional[nn.Linear] = None
+        self.train_fn = train_fn
+        self.construct_kwargs = kwargs
+
+    def _construct_model_params(
+        self,
+        in_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        norm_type: Optional[str] = None,
+        affine_norm: bool = False,
+        bias: bool = True,
+        weight_values: Optional[Tensor] = None,
+        bias_value: Optional[Tensor] = None,
+        classes: Optional[Tensor] = None,
+    ):
+        r"""
+        Lazily initializes a linear model. This will be called for you in a
+        train method.
+
+        Args:
+            in_features (int):
+                The number of input features
+            output_features (int):
+                The number of output features.
+            norm_type (str, optional):
+                The type of normalization that can occur. Please assign this
+                to one of `PyTorchLinearModel.SUPPORTED_NORMS`.
+            affine_norm (bool):
+                Whether or not to learn an affine transformation of the
+                normalization parameters used.
+            bias (bool):
+                Whether to add a bias term. Not needed if normalized input.
+            weight_values (tensor, optional):
+                The values to initialize the linear model with. This must be a
+                1D or 2D tensor, and of the form `(num_outputs, num_features)` or
+                `(num_features,)`. Additionally, if this is provided you need not
+                to provide `in_features` or `out_features`.
+            bias_value (tensor, optional):
+                The bias value to initialize the model with.
+            classes (tensor, optional):
+                The list of prediction classes supported by the model in case it
+                performs classificaton. In case of regression it is set to None.
+                Default: None
+        """
+        if norm_type not in LinearModel.SUPPORTED_NORMS:
+            raise ValueError(
+                f"{norm_type} not supported. Please use {LinearModel.SUPPORTED_NORMS}"
+            )
+
+        if weight_values is not None:
+            in_features = weight_values.shape[-1]
+            out_features = (
+                1 if len(weight_values.shape) == 1 else weight_values.shape[0]
+            )
+
+        if in_features is None or out_features is None:
+            raise ValueError(
+                "Please provide `in_features` and `out_features` or `weight_values`"
+            )
+
+        if norm_type == "batch_norm":
+            self.norm = nn.BatchNorm1d(in_features, eps=1e-8, affine=affine_norm)
+        elif norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(
+                in_features, eps=1e-8, elementwise_affine=affine_norm
+            )
+        else:
+            self.norm = None
+
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+        if weight_values is not None:
+            self.linear.weight.data = weight_values
+
+        if bias_value is not None:
+            if not bias:
+                raise ValueError("`bias_value` is not None and bias is False")
+
+            self.linear.bias.data = bias_value
+
+        if classes is not None:
+            self.linear.classes = classes
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        r"""
+        Calls `self.train_fn`
+        """
+        return self.train_fn(
+            self,
+            dataloader=train_data,
+            construct_kwargs=self.construct_kwargs,
+            **kwargs,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert self.linear is not None
+        if self.norm is not None:
+            x = self.norm(x)
+        return self.linear(x)
+
+    def representation(self) -> Tensor:
+        r"""
+        Returns a tensor which describes the hyper-plane input space. This does
+        not include the bias. For bias/intercept, please use `self.bias`
+        """
+        assert self.linear is not None
+        return self.linear.weight.detach()
+
+    def bias(self) -> Optional[Tensor]:
+        r"""
+        Returns the bias of the linear model
+        """
+        if self.linear is None or self.linear.bias is None:
+            return None
+        return self.linear.bias.detach()
+
+    def classes(self) -> Optional[Tensor]:
+        if self.linear is None or self.linear.classes is None:
+            return None
+        return cast(Tensor, self.linear.classes).detach()
+
+
+class SGDLinearModel(LinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Construct a a `LinearModel` with the
+        `sgd_train_linear_model` as the train method
+
+        Args:
+            kwargs
+                Arguments send to `self._construct_model_params` after
+                `self.fit` is called. Please refer to that method for parameter
+                documentation.
+        """
+        # avoid cycles
+        from captum._utils.models.linear_model.train import sgd_train_linear_model
+
+        super().__init__(train_fn=sgd_train_linear_model, **kwargs)
+
+
+class SGDLasso(SGDLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class to train a `LinearModel` with SGD
+        (`sgd_train_linear_model`) whilst setting appropriate parameters to
+        optimize for ridge regression loss. This optimizes L2 loss + alpha * L1
+        regularization.
+
+        Please note that with SGD it is not guaranteed that weights will
+        converge to 0.
+        """
+        super().__init__(**kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        # avoid cycles
+        from captum._utils.models.linear_model.train import l2_loss
+
+        return super().fit(train_data=train_data, loss_fn=l2_loss, reg_term=1, **kwargs)
+
+
+class SGDRidge(SGDLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class to train a `LinearModel` with SGD
+        (`sgd_train_linear_model`) whilst setting appropriate parameters to
+        optimize for ridge regression loss. This optimizes L2 loss + alpha *
+        L2 regularization.
+        """
+        super().__init__(**kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        # avoid cycles
+        from captum._utils.models.linear_model.train import l2_loss
+
+        return super().fit(train_data=train_data, loss_fn=l2_loss, reg_term=2, **kwargs)
+
+
+class SGDLinearRegression(SGDLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class to train a `LinearModel` with SGD
+        (`sgd_train_linear_model`). For linear regression this assigns the loss
+        to L2 and no regularization.
+        """
+        super().__init__(**kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        # avoid cycles
+        from captum._utils.models.linear_model.train import l2_loss
+
+        return super().fit(
+            train_data=train_data, loss_fn=l2_loss, reg_term=None, **kwargs
+        )
+
+
+class SkLearnLinearModel(LinearModel):
+    def __init__(self, sklearn_module: str, **kwargs) -> None:
+        r"""
+        Factory class to construct a `LinearModel` with sklearn training method.
+
+        Please note that this assumes:
+
+        0. You have sklearn and numpy installed
+        1. The dataset can fit into memory
+
+        SkLearn support does introduce some slight overhead as we convert the
+        tensors to numpy and then convert the resulting trained model to a
+        `LinearModel` object. However, this conversion should be negligible.
+
+        Args:
+            sklearn_module
+                The module under sklearn to construct and use for training, e.g.
+                use "svm.LinearSVC" for an SVM or "linear_model.Lasso" for Lasso.
+
+                There are factory classes defined for you for common use cases,
+                such as `SkLearnLasso`.
+            kwargs
+                The kwargs to pass to the construction of the sklearn model
+        """
+        # avoid cycles
+        from captum._utils.models.linear_model.train import sklearn_train_linear_model
+
+        super().__init__(train_fn=sklearn_train_linear_model, **kwargs)
+
+        self.sklearn_module = sklearn_module
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        r"""
+        Args:
+            train_data
+                Train data to use
+            kwargs
+                Arguments to feed to `.fit` method for sklearn
+        """
+        return super().fit(
+            train_data=train_data, sklearn_trainer=self.sklearn_module, **kwargs
+        )
+
+
+class SkLearnLasso(SkLearnLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Trains a `LinearModel` model with
+        `sklearn.linear_model.Lasso`. You will need sklearn version >= 0.23 to
+        support sample weights.
+        """
+        super().__init__(sklearn_module="linear_model.Lasso", **kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        return super().fit(train_data=train_data, **kwargs)
+
+
+class SkLearnRidge(SkLearnLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Trains a model with `sklearn.linear_model.Ridge`.
+
+        Any arguments provided to the sklearn constructor can be provided
+        as kwargs here.
+        """
+        super().__init__(sklearn_module="linear_model.Ridge", **kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        return super().fit(train_data=train_data, **kwargs)
+
+
+class SkLearnLinearRegression(SkLearnLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Trains a model with `sklearn.linear_model.LinearRegression`.
+
+        Any arguments provided to the sklearn constructor can be provided
+        as kwargs here.
+        """
+        super().__init__(sklearn_module="linear_model.LinearRegression", **kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        return super().fit(train_data=train_data, **kwargs)
+
+
+class SkLearnLogisticRegression(SkLearnLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Trains a model with `sklearn.linear_model.LogisticRegression`.
+
+        Any arguments provided to the sklearn constructor can be provided
+        as kwargs here.
+        """
+        super().__init__(sklearn_module="linear_model.LogisticRegression", **kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        return super().fit(train_data=train_data, **kwargs)
+
+
+class SkLearnSGDClassifier(SkLearnLinearModel):
+    def __init__(self, **kwargs) -> None:
+        r"""
+        Factory class. Trains a model with `sklearn.linear_model.SGDClassifier(`.
+
+        Any arguments provided to the sklearn constructor can be provided
+        as kwargs here.
+        """
+        super().__init__(sklearn_module="linear_model.SGDClassifier", **kwargs)
+
+    def fit(self, train_data: DataLoader, **kwargs):
+        return super().fit(train_data=train_data, **kwargs)

captum/_utils/models/linear_model/train.py

@@ -0,0 +1,364 @@
+import time
+import warnings
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from captum._utils.models.linear_model.model import LinearModel
+from torch.utils.data import DataLoader
+
+
+def l2_loss(x1, x2, weights=None):
+    if weights is None:
+        return torch.mean((x1 - x2) ** 2) / 2.0
+    else:
+        return torch.sum((weights / weights.norm(p=1)) * ((x1 - x2) ** 2)) / 2.0
+
+
+def sgd_train_linear_model(
+    model: LinearModel,
+    dataloader: DataLoader,
+    construct_kwargs: Dict[str, Any],
+    max_epoch: int = 100,
+    reduce_lr: bool = True,
+    initial_lr: float = 0.01,
+    alpha: float = 1.0,
+    loss_fn: Callable = l2_loss,
+    reg_term: Optional[int] = 1,
+    patience: int = 10,
+    threshold: float = 1e-4,
+    running_loss_window: Optional[int] = None,
+    device: Optional[str] = None,
+    init_scheme: str = "zeros",
+    debug: bool = False,
+) -> Dict[str, float]:
+    r"""
+    Trains a linear model with SGD. This will continue to iterate your
+    dataloader until we converged to a solution or alternatively until we have
+    exhausted `max_epoch`.
+
+    Convergence is defined by the loss not changing by `threshold` amount for
+    `patience` number of iterations.
+
+    Args:
+        model
+            The model to train
+        dataloader
+            The data to train it with. We will assume the dataloader produces
+            either pairs or triples of the form (x, y) or (x, y, w). Where x and
+            y are typical pairs for supervised learning and w is a weight
+            vector.
+
+            We will call `model._construct_model_params` with construct_kwargs
+            and the input features set to `x.shape[1]` (`x.shape[0]` corresponds
+            to the batch size). We assume that `len(x.shape) == 2`, i.e. the
+            tensor is flat. The number of output features will be set to
+            y.shape[1] or 1 (if `len(y.shape) == 1`); we require `len(y.shape)
+            <= 2`.
+        max_epoch
+            The maximum number of epochs to exhaust
+        reduce_lr
+            Whether or not to reduce the learning rate as iterations progress.
+            Halves the learning rate when the training loss does not move. This
+            uses torch.optim.lr_scheduler.ReduceLROnPlateau and uses the
+            parameters `patience` and `threshold`
+        initial_lr
+            The initial learning rate to use.
+        alpha
+            A constant for the regularization term.
+        loss_fn
+            The loss to optimise for. This must accept three parameters:
+            x1 (predicted), x2 (labels) and a weight vector
+        reg_term
+            Regularization is defined by the `reg_term` norm of the weights.
+            Please use `None` if you do not wish to use regularization.
+        patience
+            Defines the number of iterations in a row the loss must remain
+            within `threshold` in order to be classified as converged.
+        threshold
+            Threshold for convergence detection.
+        running_loss_window
+            Used to report the training loss once we have finished training and
+            to determine when we have converged (along with reducing the
+            learning rate).
+
+            The reported training loss will take the last `running_loss_window`
+            iterations and average them.
+
+            If `None` we will approximate this to be the number of examples in
+            an epoch.
+        init_scheme
+            Initialization to use prior to training the linear model.
+        device
+            The device to send the model and data to. If None then no `.to` call
+            will be used.
+        debug
+            Whether to print the loss, learning rate per iteration
+
+    Returns
+        This will return the final training loss (averaged with
+        `running_loss_window`)
+    """
+
+    loss_window: List[torch.Tensor] = []
+    min_avg_loss = None
+    convergence_counter = 0
+    converged = False
+
+    def get_point(datapoint):
+        if len(datapoint) == 2:
+            x, y = datapoint
+            w = None
+        else:
+            x, y, w = datapoint
+
+        if device is not None:
+            x = x.to(device)
+            y = y.to(device)
+            if w is not None:
+                w = w.to(device)
+
+        return x, y, w
+
+    # get a point and construct the model
+    data_iter = iter(dataloader)
+    x, y, w = get_point(next(data_iter))
+
+    model._construct_model_params(
+        in_features=x.shape[1],
+        out_features=y.shape[1] if len(y.shape) == 2 else 1,
+        **construct_kwargs,
+    )
+    model.train()
+
+    assert model.linear is not None
+
+    if init_scheme is not None:
+        assert init_scheme in ["xavier", "zeros"]
+
+        with torch.no_grad():
+            if init_scheme == "xavier":
+                torch.nn.init.xavier_uniform_(model.linear.weight)
+            else:
+                model.linear.weight.zero_()
+
+            if model.linear.bias is not None:
+                model.linear.bias.zero_()
+
+    optim = torch.optim.SGD(model.parameters(), lr=initial_lr)
+    if reduce_lr:
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optim, factor=0.5, patience=patience, threshold=threshold
+        )
+
+    t1 = time.time()
+    epoch = 0
+    i = 0
+    while epoch < max_epoch:
+        while True:  # for x, y, w in dataloader
+            if running_loss_window is None:
+                running_loss_window = x.shape[0] * len(dataloader)
+
+            y = y.view(x.shape[0], -1)
+            if w is not None:
+                w = w.view(x.shape[0], -1)
+
+            i += 1
+
+            out = model(x)
+
+            loss = loss_fn(y, out, w)
+            if reg_term is not None:
+                reg = torch.norm(model.linear.weight, p=reg_term)
+                loss += reg.sum() * alpha
+
+            if len(loss_window) >= running_loss_window:
+                loss_window = loss_window[1:]
+            loss_window.append(loss.clone().detach())
+            assert len(loss_window) <= running_loss_window
+
+            average_loss = torch.mean(torch.stack(loss_window))
+            if min_avg_loss is not None:
+                # if we haven't improved by at least `threshold`
+                if average_loss > min_avg_loss or torch.isclose(
+                    min_avg_loss, average_loss, atol=threshold
+                ):
+                    convergence_counter += 1
+                    if convergence_counter >= patience:
+                        converged = True
+                        break
+                else:
+                    convergence_counter = 0
+            if min_avg_loss is None or min_avg_loss >= average_loss:
+                min_avg_loss = average_loss.clone()
+
+            if debug:
+                print(
+                    f"lr={optim.param_groups[0]['lr']}, Loss={loss},"
+                    + "Aloss={average_loss}, min_avg_loss={min_avg_loss}"
+                )
+
+            loss.backward()
+
+            optim.step()
+            model.zero_grad()
+            if scheduler:
+                scheduler.step(average_loss)
+
+            temp = next(data_iter, None)
+            if temp is None:
+                break
+            x, y, w = get_point(temp)
+
+        if converged:
+            break
+
+        epoch += 1
+        data_iter = iter(dataloader)
+        x, y, w = get_point(next(data_iter))
+
+    t2 = time.time()
+    return {
+        "train_time": t2 - t1,
+        "train_loss": torch.mean(torch.stack(loss_window)).item(),
+        "train_iter": i,
+        "train_epoch": epoch,
+    }
+
+
+class NormLayer(nn.Module):
+    def __init__(self, mean, std, n=None, eps=1e-8) -> None:
+        super().__init__()
+        self.mean = mean
+        self.std = std
+        self.eps = eps
+
+    def forward(self, x):
+        return (x - self.mean) / (self.std + self.eps)
+
+
+def sklearn_train_linear_model(
+    model: LinearModel,
+    dataloader: DataLoader,
+    construct_kwargs: Dict[str, Any],
+    sklearn_trainer: str = "Lasso",
+    norm_input: bool = False,
+    **fit_kwargs,
+):
+    r"""
+    Alternative method to train with sklearn. This does introduce some slight
+    overhead as we convert the tensors to numpy and then convert the resulting
+    trained model to a `LinearModel` object. However, this conversion
+    should be negligible.
+
+    Please note that this assumes:
+
+    0. You have sklearn and numpy installed
+    1. The dataset can fit into memory
+
+    Args
+        model
+            The model to train.
+        dataloader
+            The data to use. This will be exhausted and converted to numpy
+            arrays. Therefore please do not feed an infinite dataloader.
+        norm_input
+            Whether or not to normalize the input
+        sklearn_trainer
+            The sklearn model to use to train the model. Please refer to
+            sklearn.linear_model for a list of modules to use.
+        construct_kwargs
+            Additional arguments provided to the `sklearn_trainer` constructor
+        fit_kwargs
+            Other arguments to send to `sklearn_trainer`'s `.fit` method
+    """
+    from functools import reduce
+
+    try:
+        import numpy as np
+    except ImportError:
+        raise ValueError("numpy is not available. Please install numpy.")
+
+    try:
+        import sklearn
+        import sklearn.linear_model
+        import sklearn.svm
+    except ImportError:
+        raise ValueError("sklearn is not available. Please install sklearn >= 0.23")
+
+    if not sklearn.__version__ >= "0.23.0":
+        warnings.warn(
+            "Must have sklearn version 0.23.0 or higher to use "
+            "sample_weight in Lasso regression."
+        )
+
+    num_batches = 0
+    xs, ys, ws = [], [], []
+    for data in dataloader:
+        if len(data) == 3:
+            x, y, w = data
+        else:
+            assert len(data) == 2
+            x, y = data
+            w = None
+
+        xs.append(x.cpu().numpy())
+        ys.append(y.cpu().numpy())
+        if w is not None:
+            ws.append(w.cpu().numpy())
+        num_batches += 1
+
+    x = np.concatenate(xs, axis=0)
+    y = np.concatenate(ys, axis=0)
+    if len(ws) > 0:
+        w = np.concatenate(ws, axis=0)
+    else:
+        w = None
+
+    if norm_input:
+        mean, std = x.mean(0), x.std(0)
+        x -= mean
+        x /= std
+
+    t1 = time.time()
+    sklearn_model = reduce(
+        lambda val, el: getattr(val, el), [sklearn] + sklearn_trainer.split(".")
+    )(**construct_kwargs)
+    try:
+        sklearn_model.fit(x, y, sample_weight=w, **fit_kwargs)
+    except TypeError:
+        sklearn_model.fit(x, y, **fit_kwargs)
+        warnings.warn(
+            "Sample weight is not supported for the provided linear model!"
+            " Trained model without weighting inputs. For Lasso, please"
+            " upgrade sklearn to a version >= 0.23.0."
+        )
+
+    t2 = time.time()
+
+    # Convert weights to pytorch
+    classes = (
+        torch.IntTensor(sklearn_model.classes_)
+        if hasattr(sklearn_model, "classes_")
+        else None
+    )
+
+    # extract model device
+    device = model.device if hasattr(model, "device") else "cpu"
+
+    num_outputs = sklearn_model.coef_.shape[0] if sklearn_model.coef_.ndim > 1 else 1
+    weight_values = torch.FloatTensor(sklearn_model.coef_).to(device)  # type: ignore
+    bias_values = torch.FloatTensor([sklearn_model.intercept_]).to(  # type: ignore
+        device  # type: ignore
+    )  # type: ignore
+    model._construct_model_params(
+        norm_type=None,
+        weight_values=weight_values.view(num_outputs, -1),
+        bias_value=bias_values.squeeze().unsqueeze(0),
+        classes=classes,
+    )
+
+    if norm_input:
+        model.norm = NormLayer(mean, std)
+
+    return {"train_time": t2 - t1}

captum/_utils/models/model.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Union
+
+from captum._utils.typing import TensorOrTupleOfTensorsGeneric
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+
+class Model(ABC):
+    r"""
+    Abstract Class to describe the interface of a trainable model to be used
+    within the algorithms of captum.
+
+    Please note that this is an experimental feature.
+    """
+
+    @abstractmethod
+    def fit(
+        self, train_data: DataLoader, **kwargs
+    ) -> Optional[Dict[str, Union[int, float, Tensor]]]:
+        r"""
+        Override this method to actually train your model.
+
+        The specification of the dataloader will be supplied by the algorithm
+        you are using within captum. This will likely be a supervised learning
+        task, thus you should expect batched (x, y) pairs or (x, y, w) triples.
+
+        Args:
+            train_data (DataLoader):
+                The data to train on
+
+        Returns:
+            Optional statistics about training, e.g.  iterations it took to
+            train, training loss, etc.
+        """
+        pass
+
+    @abstractmethod
+    def representation(self) -> Tensor:
+        r"""
+        Returns the underlying representation of the interpretable model. For a
+        linear model this is simply a tensor (the concatenation of weights
+        and bias). For something slightly more complicated, such as a decision
+        tree, this could be the nodes of a decision tree.
+
+        Returns:
+            A Tensor describing the representation of the model.
+        """
+        pass
+
+    @abstractmethod
+    def __call__(
+        self, x: TensorOrTupleOfTensorsGeneric
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Predicts with the interpretable model.
+
+        Args:
+            x (TensorOrTupleOfTensorsGeneric)
+                A batched input of tensor(s) to the model to predict
+        Returns:
+            The prediction of the input as a TensorOrTupleOfTensorsGeneric.
+        """
+        pass

captum/_utils/progress.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+
+import sys
+import warnings
+from time import time
+from typing import cast, Iterable, Sized, TextIO
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    tqdm = None
+
+
+class DisableErrorIOWrapper(object):
+    def __init__(self, wrapped: TextIO):
+        """
+        The wrapper around a TextIO object to ignore write errors like tqdm
+        https://github.com/tqdm/tqdm/blob/bcce20f771a16cb8e4ac5cc5b2307374a2c0e535/tqdm/utils.py#L131
+        """
+        self._wrapped = wrapped
+
+    def __getattr__(self, name):
+        return getattr(self._wrapped, name)
+
+    @staticmethod
+    def _wrapped_run(func, *args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except OSError as e:
+            if e.errno != 5:
+                raise
+        except ValueError as e:
+            if "closed" not in str(e):
+                raise
+
+    def write(self, *args, **kwargs):
+        return self._wrapped_run(self._wrapped.write, *args, **kwargs)
+
+    def flush(self, *args, **kwargs):
+        return self._wrapped_run(self._wrapped.flush, *args, **kwargs)
+
+
+class SimpleProgress:
+    def __init__(
+        self,
+        iterable: Iterable = None,
+        desc: str = None,
+        total: int = None,
+        file: TextIO = None,
+        mininterval: float = 0.5,
+    ):
+        """
+        Simple progress output used when tqdm is unavailable.
+        Same as tqdm, output to stderr channel
+        """
+        self.cur = 0
+
+        self.iterable = iterable
+        self.total = total
+        if total is None and hasattr(iterable, "__len__"):
+            self.total = len(cast(Sized, iterable))
+
+        self.desc = desc
+
+        file = DisableErrorIOWrapper(file if file else sys.stderr)
+        cast(TextIO, file)
+        self.file = file
+
+        self.mininterval = mininterval
+        self.last_print_t = 0.0
+        self.closed = False
+
+    def __iter__(self):
+        if self.closed or not self.iterable:
+            return
+        self._refresh()
+        for it in self.iterable:
+            yield it
+            self.update()
+        self.close()
+
+    def _refresh(self):
+        progress_str = self.desc + ": " if self.desc else ""
+        if self.total:
+            # e.g., progress: 60% 3/5
+            progress_str += f"{100 * self.cur // self.total}% {self.cur}/{self.total}"
+        else:
+            # e.g., progress: .....
+            progress_str += "." * self.cur
+
+        print("\r" + progress_str, end="", file=self.file)
+
+    def update(self, amount: int = 1):
+        if self.closed:
+            return
+        self.cur += amount
+
+        cur_t = time()
+        if cur_t - self.last_print_t >= self.mininterval:
+            self._refresh()
+            self.last_print_t = cur_t
+
+    def close(self):
+        if not self.closed:
+            self._refresh()
+            print(file=self.file)  # end with new line
+            self.closed = True
+
+
+def progress(
+    iterable: Iterable = None,
+    desc: str = None,
+    total: int = None,
+    use_tqdm=True,
+    file: TextIO = None,
+    mininterval: float = 0.5,
+    **kwargs,
+):
+    # Try to use tqdm is possible. Fall back to simple progress print
+    if tqdm and use_tqdm:
+        return tqdm(
+            iterable,
+            desc=desc,
+            total=total,
+            file=file,
+            mininterval=mininterval,
+            **kwargs,
+        )
+    else:
+        if not tqdm and use_tqdm:
+            warnings.warn(
+                "Tried to show progress with tqdm "
+                "but tqdm is not installed. "
+                "Fall back to simply print out the progress."
+            )
+        return SimpleProgress(
+            iterable, desc=desc, total=total, file=file, mininterval=mininterval
+        )

captum/_utils/sample_gradient.py

@@ -0,0 +1,184 @@
+from collections import defaultdict
+from enum import Enum
+from typing import cast, Iterable, Tuple, Union
+
+import torch
+from captum._utils.common import _format_tensor_into_tuples, _register_backward_hook
+from torch import Tensor
+from torch.nn import Module
+
+
+def _reset_sample_grads(module: Module):
+    module.weight.sample_grad = 0  # type: ignore
+    if module.bias is not None:
+        module.bias.sample_grad = 0  # type: ignore
+
+
+def linear_param_grads(
+    module: Module, activation: Tensor, gradient_out: Tensor, reset: bool = False
+) -> None:
+    r"""
+    Computes parameter gradients per sample for nn.Linear module, given module
+    input activations and output gradients.
+
+    Gradients are accumulated in the sample_grad attribute of each parameter
+    (weight and bias). If reset = True, any current sample_grad values are reset,
+    otherwise computed gradients are accumulated and added to the existing
+    stored gradients.
+
+    Inputs with more than 2 dimensions are only supported with torch 1.8 or later
+    """
+    if reset:
+        _reset_sample_grads(module)
+
+    module.weight.sample_grad += torch.einsum(  # type: ignore
+        "n...i,n...j->nij", gradient_out, activation
+    )
+    if module.bias is not None:
+        module.bias.sample_grad += torch.einsum(  # type: ignore
+            "n...i->ni", gradient_out
+        )
+
+
+def conv2d_param_grads(
+    module: Module, activation: Tensor, gradient_out: Tensor, reset: bool = False
+) -> None:
+    r"""
+    Computes parameter gradients per sample for nn.Conv2d module, given module
+    input activations and output gradients.
+
+    nn.Conv2d modules with padding set to a string option ('same' or 'valid') are
+    currently unsupported.
+
+    Gradients are accumulated in the sample_grad attribute of each parameter
+    (weight and bias). If reset = True, any current sample_grad values are reset,
+    otherwise computed gradients are accumulated and added to the existing
+    stored gradients.
+    """
+    if reset:
+        _reset_sample_grads(module)
+
+    batch_size = cast(int, activation.shape[0])
+    unfolded_act = torch.nn.functional.unfold(
+        activation,
+        cast(Union[int, Tuple[int, ...]], module.kernel_size),
+        dilation=cast(Union[int, Tuple[int, ...]], module.dilation),
+        padding=cast(Union[int, Tuple[int, ...]], module.padding),
+        stride=cast(Union[int, Tuple[int, ...]], module.stride),
+    )
+    reshaped_grad = gradient_out.reshape(batch_size, -1, unfolded_act.shape[-1])
+    grad1 = torch.einsum("ijk,ilk->ijl", reshaped_grad, unfolded_act)
+    shape = [batch_size] + list(cast(Iterable[int], module.weight.shape))
+    module.weight.sample_grad += grad1.reshape(shape)  # type: ignore
+    if module.bias is not None:
+        module.bias.sample_grad += torch.sum(reshaped_grad, dim=2)  # type: ignore
+
+
+SUPPORTED_MODULES = {
+    torch.nn.Conv2d: conv2d_param_grads,
+    torch.nn.Linear: linear_param_grads,
+}
+
+
+class LossMode(Enum):
+    SUM = 0
+    MEAN = 1
+
+
+class SampleGradientWrapper:
+    r"""
+    Wrapper which allows computing sample-wise gradients in a single backward pass.
+
+    This is accomplished by adding hooks to capture activations and output
+    gradients for supported modules, and using these activations and gradients
+    to compute the parameter gradients per-sample.
+
+    Currently, only nn.Linear and nn.Conv2d modules are supported.
+
+    Similar reference implementations of sample-based gradients include:
+    - https://github.com/cybertronai/autograd-hacks
+    - https://github.com/pytorch/opacus/tree/main/opacus/grad_sample
+    """
+
+    def __init__(self, model):
+        self.model = model
+        self.hooks_added = False
+        self.activation_dict = defaultdict(list)
+        self.gradient_dict = defaultdict(list)
+        self.forward_hooks = []
+        self.backward_hooks = []
+
+    def add_hooks(self):
+        self.hooks_added = True
+        self.model.apply(self._register_module_hooks)
+
+    def _register_module_hooks(self, module: torch.nn.Module):
+        if isinstance(module, tuple(SUPPORTED_MODULES.keys())):
+            self.forward_hooks.append(
+                module.register_forward_hook(self._forward_hook_fn)
+            )
+            self.backward_hooks.append(
+                _register_backward_hook(module, self._backward_hook_fn, None)
+            )
+
+    def _forward_hook_fn(
+        self,
+        module: Module,
+        module_input: Union[Tensor, Tuple[Tensor, ...]],
+        module_output: Union[Tensor, Tuple[Tensor, ...]],
+    ):
+        inp_tuple = _format_tensor_into_tuples(module_input)
+        self.activation_dict[module].append(inp_tuple[0].clone().detach())
+
+    def _backward_hook_fn(
+        self,
+        module: Module,
+        grad_input: Union[Tensor, Tuple[Tensor, ...]],
+        grad_output: Union[Tensor, Tuple[Tensor, ...]],
+    ):
+        grad_output_tuple = _format_tensor_into_tuples(grad_output)
+        self.gradient_dict[module].append(grad_output_tuple[0].clone().detach())
+
+    def remove_hooks(self):
+        self.hooks_added = False
+
+        for hook in self.forward_hooks:
+            hook.remove()
+
+        for hook in self.backward_hooks:
+            hook.remove()
+
+        self.forward_hooks = []
+        self.backward_hooks = []
+
+    def _reset(self):
+        self.activation_dict = defaultdict(list)
+        self.gradient_dict = defaultdict(list)
+
+    def compute_param_sample_gradients(self, loss_blob, loss_mode="mean"):
+        assert (
+            loss_mode.upper() in LossMode.__members__
+        ), f"Provided loss mode {loss_mode} is not valid"
+        mode = LossMode[loss_mode.upper()]
+
+        self.model.zero_grad()
+        loss_blob.backward(gradient=torch.ones_like(loss_blob))
+
+        for module in self.gradient_dict:
+            sample_grad_fn = SUPPORTED_MODULES[type(module)]
+            activations = self.activation_dict[module]
+            gradients = self.gradient_dict[module]
+            assert len(activations) == len(gradients), (
+                "Number of saved activations do not match number of saved gradients."
+                " This may occur if multiple forward passes are run without calling"
+                " reset or computing param gradients."
+            )
+            # Reversing grads since when a module is used multiple times,
+            # the activations will be aligned with the reverse order of the gradients,
+            # since the order is reversed in backprop.
+            for i, (act, grad) in enumerate(
+                zip(activations, list(reversed(gradients)))
+            ):
+                mult = 1 if mode is LossMode.SUM else act.shape[0]
+                sample_grad_fn(module, act, grad * mult, reset=(i == 0))
+        self._reset()

captum/_utils/typing.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+
+from typing import List, Tuple, TYPE_CHECKING, TypeVar, Union
+
+from torch import Tensor
+from torch.nn import Module
+
+if TYPE_CHECKING:
+    import sys
+
+    if sys.version_info >= (3, 8):
+        from typing import Literal  # noqa: F401
+    else:
+        from typing_extensions import Literal  # noqa: F401
+else:
+    Literal = {True: bool, False: bool, (True, False): bool}
+
+TensorOrTupleOfTensorsGeneric = TypeVar(
+    "TensorOrTupleOfTensorsGeneric", Tensor, Tuple[Tensor, ...]
+)
+TupleOrTensorOrBoolGeneric = TypeVar("TupleOrTensorOrBoolGeneric", Tuple, Tensor, bool)
+ModuleOrModuleList = TypeVar("ModuleOrModuleList", Module, List[Module])
+TargetType = Union[None, int, Tuple[int, ...], Tensor, List[Tuple[int, ...]], List[int]]
+BaselineType = Union[None, Tensor, int, float, Tuple[Union[Tensor, int, float], ...]]
+
+TensorLikeList1D = List[float]
+TensorLikeList2D = List[TensorLikeList1D]
+TensorLikeList3D = List[TensorLikeList2D]
+TensorLikeList4D = List[TensorLikeList3D]
+TensorLikeList5D = List[TensorLikeList4D]
+TensorLikeList = Union[
+    TensorLikeList1D,
+    TensorLikeList2D,
+    TensorLikeList3D,
+    TensorLikeList4D,
+    TensorLikeList5D,
+]

captum/attr/__init__.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+from captum.attr._core.deep_lift import DeepLift, DeepLiftShap  # noqa
+from captum.attr._core.feature_ablation import FeatureAblation  # noqa
+from captum.attr._core.feature_permutation import FeaturePermutation  # noqa
+from captum.attr._core.gradient_shap import GradientShap  # noqa
+from captum.attr._core.guided_backprop_deconvnet import (  # noqa
+    Deconvolution,
+    GuidedBackprop,
+)
+from captum.attr._core.guided_grad_cam import GuidedGradCam  # noqa
+from captum.attr._core.input_x_gradient import InputXGradient  # noqa
+from captum.attr._core.integrated_gradients import IntegratedGradients  # noqa
+from captum.attr._core.kernel_shap import KernelShap  # noqa
+from captum.attr._core.layer.grad_cam import LayerGradCam  # noqa
+from captum.attr._core.layer.internal_influence import InternalInfluence  # noqa
+from captum.attr._core.layer.layer_activation import LayerActivation  # noqa
+from captum.attr._core.layer.layer_conductance import LayerConductance  # noqa
+from captum.attr._core.layer.layer_deep_lift import (  # noqa
+    LayerDeepLift,
+    LayerDeepLiftShap,
+)
+from captum.attr._core.layer.layer_feature_ablation import LayerFeatureAblation  # noqa
+from captum.attr._core.layer.layer_gradient_shap import LayerGradientShap  # noqa
+from captum.attr._core.layer.layer_gradient_x_activation import (  # noqa
+    LayerGradientXActivation,
+)
+from captum.attr._core.layer.layer_integrated_gradients import (  # noqa
+    LayerIntegratedGradients,
+)
+from captum.attr._core.layer.layer_lrp import LayerLRP  # noqa
+from captum.attr._core.lime import Lime, LimeBase  # noqa
+from captum.attr._core.lrp import LRP  # noqa
+from captum.attr._core.neuron.neuron_conductance import NeuronConductance  # noqa
+from captum.attr._core.neuron.neuron_deep_lift import (  # noqa
+    NeuronDeepLift,
+    NeuronDeepLiftShap,
+)
+from captum.attr._core.neuron.neuron_feature_ablation import (  # noqa
+    NeuronFeatureAblation,
+)
+from captum.attr._core.neuron.neuron_gradient import NeuronGradient  # noqa
+from captum.attr._core.neuron.neuron_gradient_shap import NeuronGradientShap  # noqa
+from captum.attr._core.neuron.neuron_guided_backprop_deconvnet import (  # noqa
+    NeuronDeconvolution,
+    NeuronGuidedBackprop,
+)
+from captum.attr._core.neuron.neuron_integrated_gradients import (  # noqa
+    NeuronIntegratedGradients,
+)
+from captum.attr._core.noise_tunnel import NoiseTunnel  # noqa
+from captum.attr._core.occlusion import Occlusion  # noqa
+from captum.attr._core.saliency import Saliency  # noqa
+from captum.attr._core.shapley_value import ShapleyValues, ShapleyValueSampling  # noqa
+from captum.attr._models.base import (  # noqa
+    configure_interpretable_embedding_layer,
+    InterpretableEmbeddingBase,
+    remove_interpretable_embedding_layer,
+    TokenReferenceBase,
+)
+from captum.attr._utils import visualization  # noqa
+from captum.attr._utils.attribution import (  # noqa  # noqa  # noqa  # noqa  # noqa
+    Attribution,
+    GradientAttribution,
+    LayerAttribution,
+    NeuronAttribution,
+    PerturbationAttribution,
+)
+from captum.attr._utils.class_summarizer import ClassSummarizer
+from captum.attr._utils.stat import (
+    CommonStats,
+    Count,
+    Max,
+    Mean,
+    Min,
+    MSE,
+    StdDev,
+    Sum,
+    Var,
+)
+from captum.attr._utils.summarizer import Summarizer
+
+__all__ = [
+    "Attribution",
+    "GradientAttribution",
+    "PerturbationAttribution",
+    "NeuronAttribution",
+    "LayerAttribution",
+    "IntegratedGradients",
+    "DeepLift",
+    "DeepLiftShap",
+    "InputXGradient",
+    "Saliency",
+    "GuidedBackprop",
+    "Deconvolution",
+    "GuidedGradCam",
+    "FeatureAblation",
+    "FeaturePermutation",
+    "Occlusion",
+    "ShapleyValueSampling",
+    "ShapleyValues",
+    "LimeBase",
+    "Lime",
+    "LRP",
+    "KernelShap",
+    "LayerConductance",
+    "LayerGradientXActivation",
+    "LayerActivation",
+    "LayerFeatureAblation",
+    "InternalInfluence",
+    "LayerGradCam",
+    "LayerDeepLift",
+    "LayerDeepLiftShap",
+    "LayerGradientShap",
+    "LayerIntegratedGradients",
+    "LayerLRP",
+    "NeuronConductance",
+    "NeuronFeatureAblation",
+    "NeuronGradient",
+    "NeuronIntegratedGradients",
+    "NeuronDeepLift",
+    "NeuronDeepLiftShap",
+    "NeuronGradientShap",
+    "NeuronDeconvolution",
+    "NeuronGuidedBackprop",
+    "NoiseTunnel",
+    "GradientShap",
+    "InterpretableEmbeddingBase",
+    "TokenReferenceBase",
+    "visualization",
+    "configure_interpretable_embedding_layer",
+    "remove_interpretable_embedding_layer",
+    "Summarizer",
+    "CommonStats",
+    "ClassSummarizer",
+    "Mean",
+    "StdDev",
+    "MSE",
+    "Var",
+    "Min",
+    "Max",
+    "Sum",
+    "Count",
+]

captum/attr/_core/deep_lift.py

@@ -0,0 +1,1151 @@
+#!/usr/bin/env python3
+import typing
+import warnings
+from typing import Any, Callable, cast, List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from captum._utils.common import (
+    _expand_additional_forward_args,
+    _expand_target,
+    _format_additional_forward_args,
+    _format_baseline,
+    _format_output,
+    _format_tensor_into_tuples,
+    _is_tuple,
+    _register_backward_hook,
+    _run_forward,
+    _select_targets,
+    ExpansionTypes,
+)
+from captum._utils.gradient import (
+    apply_gradient_requirements,
+    undo_gradient_requirements,
+)
+from captum._utils.typing import (
+    BaselineType,
+    Literal,
+    TargetType,
+    TensorOrTupleOfTensorsGeneric,
+)
+from captum.attr._utils.attribution import GradientAttribution
+from captum.attr._utils.common import (
+    _call_custom_attribution_func,
+    _compute_conv_delta_and_format_attrs,
+    _format_callable_baseline,
+    _tensorize_baseline,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+from torch.utils.hooks import RemovableHandle
+
+
+# Check if module backward hook can safely be used for the module that produced
+# this inputs / outputs mapping
+def _check_valid_module(inputs_grad_fn, outputs) -> bool:
+    def is_output_cloned(output_fn, input_grad_fn) -> bool:
+        """
+        Checks if the output has been cloned. This happens especially in case of
+        layer deeplift.
+        """
+        return (
+            output_fn[0].next_functions is not None
+            and output_fn[0].next_functions[0][0] == input_grad_fn
+        )
+
+    curr_fn = outputs.grad_fn
+    first_next = curr_fn.next_functions[0]
+    try:
+        # if `inputs` in the input to the network then the grad_fn is None and
+        # for that input backward_hook isn't computed. That's the reason why we
+        # need to check on `inputs_grad_fns[first_next[1]]` being None.
+        return (
+            inputs_grad_fn is None
+            or first_next[0] == inputs_grad_fn
+            or is_output_cloned(first_next, inputs_grad_fn)
+        )
+    except IndexError:
+        return False
+
+
+class DeepLift(GradientAttribution):
+    r"""
+    Implements DeepLIFT algorithm based on the following paper:
+    Learning Important Features Through Propagating Activation Differences,
+    Avanti Shrikumar, et. al.
+    https://arxiv.org/abs/1704.02685
+
+    and the gradient formulation proposed in:
+    Towards better understanding of gradient-based attribution methods for
+    deep neural networks,  Marco Ancona, et.al.
+    https://openreview.net/pdf?id=Sy21R9JAW
+
+    This implementation supports only Rescale rule. RevealCancel rule will
+    be supported in later releases.
+    In addition to that, in order to keep the implementation cleaner, DeepLIFT
+    for internal neurons and layers extends current implementation and is
+    implemented separately in LayerDeepLift and NeuronDeepLift.
+    Although DeepLIFT's(Rescale Rule) attribution quality is comparable with
+    Integrated Gradients, it runs significantly faster than Integrated
+    Gradients and is preferred for large datasets.
+
+    Currently we only support a limited number of non-linear activations
+    but the plan is to expand the list in the future.
+
+    Note: As we know, currently we cannot access the building blocks,
+    of PyTorch's built-in LSTM, RNNs and GRUs such as Tanh and Sigmoid.
+    Nonetheless, it is possible to build custom LSTMs, RNNS and GRUs
+    with performance similar to built-in ones using TorchScript.
+    More details on how to build custom RNNs can be found here:
+    https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/
+    """
+
+    def __init__(
+        self,
+        model: Module,
+        multiply_by_inputs: bool = True,
+        eps: float = 1e-10,
+    ) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place nonlinear submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API
+                        starting from PyTorch v1.9.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in
+                        then that type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of DeepLift, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores
+                        are being multiplied by (inputs - baselines).
+                        This flag applies only if `custom_attribution_func` is
+                        set to None.
+
+            eps (float, optional): A value at which to consider output/input change
+                        significant when computing the gradients for non-linear layers.
+                        This is useful to adjust, depending on your model's bit depth,
+                        to avoid numerical issues during the gradient computation.
+                        Default: 1e-10
+        """
+        GradientAttribution.__init__(self, model)
+        self.model = model
+        self.eps = eps
+        self.forward_handles: List[RemovableHandle] = []
+        self.backward_handles: List[RemovableHandle] = []
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define reference samples that are compared with
+                        the inputs. In order to assign attribution scores DeepLift
+                        computes the differences between the inputs/outputs and
+                        corresponding references.
+                        Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            custom_attribution_func (callable, optional): A custom function for
+                        computing final attribution scores. This function can take
+                        at least one and at most three arguments with the
+                        following signature:
+
+                        - custom_attribution_func(multipliers)
+                        - custom_attribution_func(multipliers, inputs)
+                        - custom_attribution_func(multipliers, inputs, baselines)
+
+                        In case this function is not provided, we use the default
+                        logic defined as: multipliers * (inputs - baselines)
+                        It is assumed that all input arguments, `multipliers`,
+                        `inputs` and `baselines` are provided in tuples of same
+                        length. `custom_attribution_func` returns a tuple of
+                        attribution tensors that have the same length as the
+                        `inputs`.
+
+                        Default: None
+
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                Attribution score computed based on DeepLift rescale rule with respect
+                to each input feature. Attributions will always be
+                the same size as the provided inputs, with each value
+                providing the attribution of the corresponding input index.
+                If a single tensor is provided as inputs, a single tensor is
+                returned. If a tuple is provided for inputs, a tuple of
+                corresponding sized tensors is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                This is computed using the property that
+                the total sum of forward_func(inputs) - forward_func(baselines)
+                must equal the total sum of the attributions computed
+                based on DeepLift's rescale rule.
+                Delta is calculated per example, meaning that the number of
+                elements in returned delta tensor is equal to the number of
+                of examples in input.
+                Note that the logic described for deltas is guaranteed when the
+                default logic for attribution computations is used, meaning that the
+                `custom_attribution_func=None`, otherwise it is not guaranteed and
+                depends on the specifics of the `custom_attribution_func`.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> dl = DeepLift(net)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes deeplift attribution scores for class 3.
+            >>> attribution = dl.attribute(input, target=3)
+        """
+
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+
+        inputs = _format_tensor_into_tuples(inputs)
+        baselines = _format_baseline(baselines, inputs)
+
+        gradient_mask = apply_gradient_requirements(inputs)
+
+        _validate_input(inputs, baselines)
+
+        # set hooks for baselines
+        warnings.warn(
+            """Setting forward, backward hooks and attributes on non-linear
+               activations. The hooks and attributes will be removed
+            after the attribution is finished"""
+        )
+        baselines = _tensorize_baseline(inputs, baselines)
+        main_model_hooks = []
+        try:
+            main_model_hooks = self._hook_main_model()
+
+            self.model.apply(self._register_hooks)
+
+            additional_forward_args = _format_additional_forward_args(
+                additional_forward_args
+            )
+
+            expanded_target = _expand_target(
+                target, 2, expansion_type=ExpansionTypes.repeat
+            )
+
+            wrapped_forward_func = self._construct_forward_func(
+                self.model,
+                (inputs, baselines),
+                expanded_target,
+                additional_forward_args,
+            )
+            gradients = self.gradient_func(wrapped_forward_func, inputs)
+            if custom_attribution_func is None:
+                if self.multiplies_by_inputs:
+                    attributions = tuple(
+                        (input - baseline) * gradient
+                        for input, baseline, gradient in zip(
+                            inputs, baselines, gradients
+                        )
+                    )
+                else:
+                    attributions = gradients
+            else:
+                attributions = _call_custom_attribution_func(
+                    custom_attribution_func, gradients, inputs, baselines
+                )
+        finally:
+            # Even if any error is raised, remove all hooks before raising
+            self._remove_hooks(main_model_hooks)
+
+        undo_gradient_requirements(inputs, gradient_mask)
+        return _compute_conv_delta_and_format_attrs(
+            self,
+            return_convergence_delta,
+            attributions,
+            baselines,
+            inputs,
+            additional_forward_args,
+            target,
+            is_inputs_tuple,
+        )
+
+    def _construct_forward_func(
+        self,
+        forward_func: Callable,
+        inputs: Tuple,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> Callable:
+        def forward_fn():
+            model_out = _run_forward(
+                forward_func, inputs, None, additional_forward_args
+            )
+            return _select_targets(
+                torch.cat((model_out[:, 0], model_out[:, 1])), target
+            )
+
+        if hasattr(forward_func, "device_ids"):
+            forward_fn.device_ids = forward_func.device_ids  # type: ignore
+        return forward_fn
+
+    def _is_non_linear(self, module: Module) -> bool:
+        return type(module) in SUPPORTED_NON_LINEAR.keys()
+
+    def _forward_pre_hook_ref(
+        self, module: Module, inputs: Union[Tensor, Tuple[Tensor, ...]]
+    ) -> None:
+        inputs = _format_tensor_into_tuples(inputs)
+        module.input_ref = tuple(  # type: ignore
+            input.clone().detach() for input in inputs
+        )
+
+    def _forward_pre_hook(
+        self, module: Module, inputs: Union[Tensor, Tuple[Tensor, ...]]
+    ) -> None:
+        """
+        For the modules that perform in-place operations such as ReLUs, we cannot
+        use inputs from forward hooks. This is because in that case inputs
+        and outputs are the same. We need access the inputs in pre-hooks and
+        set necessary hooks on inputs there.
+        """
+        inputs = _format_tensor_into_tuples(inputs)
+        module.input = inputs[0].clone().detach()
+        module.input_grad_fns = inputs[0].grad_fn  # type: ignore
+
+        def tensor_backward_hook(grad):
+            if module.saved_grad is None:
+                raise RuntimeError(
+                    """Module {} was detected as not supporting correctly module
+                        backward hook. You should modify your hook to ignore the given
+                        grad_inputs (recompute them by hand if needed) and save the
+                        newly computed grad_inputs in module.saved_grad. See MaxPool1d
+                        as an example.""".format(
+                        module
+                    )
+                )
+            return module.saved_grad
+
+        # the hook is set by default but it will be used only for
+        # failure cases and will be removed otherwise
+        handle = inputs[0].register_hook(tensor_backward_hook)
+        module.input_hook = handle
+
+    def _forward_hook(
+        self,
+        module: Module,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        outputs: Union[Tensor, Tuple[Tensor, ...]],
+    ) -> None:
+        r"""
+        we need forward hook to access and detach the inputs and
+        outputs of a neuron
+        """
+        outputs = _format_tensor_into_tuples(outputs)
+        module.output = outputs[0].clone().detach()
+        if not _check_valid_module(module.input_grad_fns, outputs[0]):
+            warnings.warn(
+                """An invalid module {} is detected. Saved gradients will
+                be used as the gradients of the module's input tensor.
+                See MaxPool1d as an example.""".format(
+                    module
+                )
+            )
+            module.is_invalid = True  # type: ignore
+            module.saved_grad = None  # type: ignore
+            self.forward_handles.append(cast(RemovableHandle, module.input_hook))
+        else:
+            module.is_invalid = False  # type: ignore
+            # removing the hook if there is no failure case
+            cast(RemovableHandle, module.input_hook).remove()
+        del module.input_hook
+        del module.input_grad_fns
+
+    def _backward_hook(
+        self,
+        module: Module,
+        grad_input: Union[Tensor, Tuple[Tensor, ...]],
+        grad_output: Union[Tensor, Tuple[Tensor, ...]],
+    ):
+        r"""
+        `grad_input` is the gradient of the neuron with respect to its input
+        `grad_output` is the gradient of the neuron with respect to its output
+         we can override `grad_input` according to chain rule with.
+        `grad_output` * delta_out / delta_in.
+
+        """
+        # before accessing the attributes from the module we want
+        # to ensure that the properties exist, if not, then it is
+        # likely that the module is being reused.
+        attr_criteria = self.satisfies_attribute_criteria(module)
+        if not attr_criteria:
+            raise RuntimeError(
+                "A Module {} was detected that does not contain some of "
+                "the input/output attributes that are required for DeepLift "
+                "computations. This can occur, for example, if "
+                "your module is being used more than once in the network."
+                "Please, ensure that module is being used only once in the "
+                "network.".format(module)
+            )
+        multipliers = tuple(
+            SUPPORTED_NON_LINEAR[type(module)](
+                module,
+                module.input,
+                module.output,
+                grad_input,
+                grad_output,
+                eps=self.eps,
+            )
+        )
+        # remove all the properies that we set for the inputs and output
+        del module.input
+        del module.output
+
+        return multipliers
+
+    def satisfies_attribute_criteria(self, module: Module) -> bool:
+        return hasattr(module, "input") and hasattr(module, "output")
+
+    def _can_register_hook(self, module: Module) -> bool:
+        # TODO find a better way of checking if a module is a container or not
+        module_fullname = str(type(module))
+        has_already_hooks = len(module._backward_hooks) > 0  # type: ignore
+        return not (
+            "nn.modules.container" in module_fullname
+            or has_already_hooks
+            or not self._is_non_linear(module)
+        )
+
+    def _register_hooks(
+        self, module: Module, attribute_to_layer_input: bool = True
+    ) -> None:
+        if not self._can_register_hook(module) or (
+            not attribute_to_layer_input and module is self.layer  # type: ignore
+        ):
+            return
+        # adds forward hook to leaf nodes that are non-linear
+        forward_handle = module.register_forward_hook(self._forward_hook)
+        pre_forward_handle = module.register_forward_pre_hook(self._forward_pre_hook)
+        backward_handle = _register_backward_hook(module, self._backward_hook, self)
+        self.forward_handles.append(forward_handle)
+        self.forward_handles.append(pre_forward_handle)
+        self.backward_handles.append(backward_handle)
+
+    def _remove_hooks(self, extra_hooks_to_remove: List[RemovableHandle]) -> None:
+        for handle in extra_hooks_to_remove:
+            handle.remove()
+        for forward_handle in self.forward_handles:
+            forward_handle.remove()
+        for backward_handle in self.backward_handles:
+            backward_handle.remove()
+
+    def _hook_main_model(self) -> List[RemovableHandle]:
+        def pre_hook(module: Module, baseline_inputs_add_args: Tuple) -> Tuple:
+            inputs = baseline_inputs_add_args[0]
+            baselines = baseline_inputs_add_args[1]
+            additional_args = None
+            if len(baseline_inputs_add_args) > 2:
+                additional_args = baseline_inputs_add_args[2:]
+
+            baseline_input_tsr = tuple(
+                torch.cat([input, baseline])
+                for input, baseline in zip(inputs, baselines)
+            )
+            if additional_args is not None:
+                expanded_additional_args = cast(
+                    Tuple,
+                    _expand_additional_forward_args(
+                        additional_args, 2, ExpansionTypes.repeat
+                    ),
+                )
+                return (*baseline_input_tsr, *expanded_additional_args)
+            return baseline_input_tsr
+
+        def forward_hook(module: Module, inputs: Tuple, outputs: Tensor):
+            return torch.stack(torch.chunk(outputs, 2), dim=1)
+
+        if isinstance(
+            self.model, (nn.DataParallel, nn.parallel.DistributedDataParallel)
+        ):
+            return [
+                self.model.module.register_forward_pre_hook(pre_hook),  # type: ignore
+                self.model.module.register_forward_hook(forward_hook),
+            ]  # type: ignore
+        else:
+            return [
+                self.model.register_forward_pre_hook(pre_hook),  # type: ignore
+                self.model.register_forward_hook(forward_hook),
+            ]  # type: ignore
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+
+class DeepLiftShap(DeepLift):
+    r"""
+    Extends DeepLift algorithm and approximates SHAP values using Deeplift.
+    For each input sample it computes DeepLift attribution with respect to
+    each baseline and averages resulting attributions.
+    More details about the algorithm can be found here:
+
+    http://papers.nips.cc/paper/7062-a-unified-approach-to-interpreting-model-predictions.pdf
+
+    Note that the explanation model:
+        1. Assumes that input features are independent of one another
+        2. Is linear, meaning that the explanations are modeled through
+            the additive composition of feature effects.
+    Although, it assumes a linear model for each explanation, the overall
+    model across multiple explanations can be complex and non-linear.
+    """
+
+    def __init__(self, model: Module, multiply_by_inputs: bool = True) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place nonlinear submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in
+                        then that type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of DeepLiftShap, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores
+                        are being multiplied by (inputs - baselines).
+                        This flag applies only if `custom_attribution_func` is
+                        set to None.
+        """
+        DeepLift.__init__(self, model, multiply_by_inputs=multiply_by_inputs)
+
+    # There's a mismatch between the signatures of DeepLift.attribute and
+    # DeepLiftShap.attribute, so we ignore typing here
+    @typing.overload  # type: ignore
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            baselines (tensor, tuple of tensors, callable):
+                        Baselines define reference samples that are compared with
+                        the inputs. In order to assign attribution scores DeepLift
+                        computes the differences between the inputs/outputs and
+                        corresponding references. Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          the first dimension equal to the number of examples
+                          in the baselines' distribution. The remaining dimensions
+                          must match with input tensor's dimension starting from
+                          the second dimension.
+
+                        - a tuple of tensors, if inputs is a tuple of tensors,
+                          with the first dimension of any tensor inside the tuple
+                          equal to the number of examples in the baseline's
+                          distribution. The remaining dimensions must match
+                          the dimensions of the corresponding input tensor
+                          starting from the second dimension.
+
+                        - callable function, optionally takes `inputs` as an
+                          argument and either returns a single tensor
+                          or a tuple of those.
+
+                        It is recommended that the number of samples in the baselines'
+                        tensors is larger than one.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            custom_attribution_func (callable, optional): A custom function for
+                        computing final attribution scores. This function can take
+                        at least one and at most three arguments with the
+                        following signature:
+
+                        - custom_attribution_func(multipliers)
+                        - custom_attribution_func(multipliers, inputs)
+                        - custom_attribution_func(multipliers, inputs, baselines)
+
+                        In case this function is not provided we use the default
+                        logic defined as: multipliers * (inputs - baselines)
+                        It is assumed that all input arguments, `multipliers`,
+                        `inputs` and `baselines` are provided in tuples of same
+                        length. `custom_attribution_func` returns a tuple of
+                        attribution tensors that have the same length as the
+                        `inputs`.
+                        Default: None
+
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attribution score computed based on DeepLift rescale rule with
+                        respect to each input feature. Attributions will always be
+                        the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple is provided for inputs, a tuple of
+                        corresponding sized tensors is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        This is computed using the property that the
+                        total sum of forward_func(inputs) - forward_func(baselines)
+                        must be very close to the total sum of attributions
+                        computed based on approximated SHAP values using
+                        Deeplift's rescale rule.
+                        Delta is calculated for each example input and baseline pair,
+                        meaning that the number of elements in returned delta tensor
+                        is equal to the
+                        `number of examples in input` * `number of examples
+                        in baseline`. The deltas are ordered in the first place by
+                        input example, followed by the baseline.
+                        Note that the logic described for deltas is guaranteed
+                        when the default logic for attribution computations is used,
+                        meaning that the `custom_attribution_func=None`, otherwise
+                        it is not guaranteed and depends on the specifics of the
+                        `custom_attribution_func`.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> dl = DeepLiftShap(net)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes shap values using deeplift for class 3.
+            >>> attribution = dl.attribute(input, target=3)
+        """
+        baselines = _format_callable_baseline(baselines, inputs)
+
+        assert isinstance(baselines[0], torch.Tensor) and baselines[0].shape[0] > 1, (
+            "Baselines distribution has to be provided in form of a torch.Tensor"
+            " with more than one example but found: {}."
+            " If baselines are provided in shape of scalars or with a single"
+            " baseline example, `DeepLift`"
+            " approach can be used instead.".format(baselines[0])
+        )
+
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+
+        inputs = _format_tensor_into_tuples(inputs)
+
+        # batch sizes
+        inp_bsz = inputs[0].shape[0]
+        base_bsz = baselines[0].shape[0]
+
+        (
+            exp_inp,
+            exp_base,
+            exp_tgt,
+            exp_addit_args,
+        ) = self._expand_inputs_baselines_targets(
+            baselines, inputs, target, additional_forward_args
+        )
+        attributions = super().attribute.__wrapped__(  # type: ignore
+            self,
+            exp_inp,
+            exp_base,
+            target=exp_tgt,
+            additional_forward_args=exp_addit_args,
+            return_convergence_delta=cast(
+                Literal[True, False], return_convergence_delta
+            ),
+            custom_attribution_func=custom_attribution_func,
+        )
+        if return_convergence_delta:
+            attributions, delta = cast(Tuple[Tuple[Tensor, ...], Tensor], attributions)
+
+        attributions = tuple(
+            self._compute_mean_across_baselines(
+                inp_bsz, base_bsz, cast(Tensor, attribution)
+            )
+            for attribution in attributions
+        )
+
+        if return_convergence_delta:
+            return _format_output(is_inputs_tuple, attributions), delta
+        else:
+            return _format_output(is_inputs_tuple, attributions)
+
+    def _expand_inputs_baselines_targets(
+        self,
+        baselines: Tuple[Tensor, ...],
+        inputs: Tuple[Tensor, ...],
+        target: TargetType,
+        additional_forward_args: Any,
+    ) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], TargetType, Any]:
+        inp_bsz = inputs[0].shape[0]
+        base_bsz = baselines[0].shape[0]
+
+        expanded_inputs = tuple(
+            [
+                input.repeat_interleave(base_bsz, dim=0).requires_grad_()
+                for input in inputs
+            ]
+        )
+        expanded_baselines = tuple(
+            [
+                baseline.repeat(
+                    (inp_bsz,) + tuple([1] * (len(baseline.shape) - 1))
+                ).requires_grad_()
+                for baseline in baselines
+            ]
+        )
+        expanded_target = _expand_target(
+            target, base_bsz, expansion_type=ExpansionTypes.repeat_interleave
+        )
+        input_additional_args = (
+            _expand_additional_forward_args(
+                additional_forward_args,
+                base_bsz,
+                expansion_type=ExpansionTypes.repeat_interleave,
+            )
+            if additional_forward_args is not None
+            else None
+        )
+        return (
+            expanded_inputs,
+            expanded_baselines,
+            expanded_target,
+            input_additional_args,
+        )
+
+    def _compute_mean_across_baselines(
+        self, inp_bsz: int, base_bsz: int, attribution: Tensor
+    ) -> Tensor:
+        # Average for multiple references
+        attr_shape: Tuple = (inp_bsz, base_bsz)
+        if len(attribution.shape) > 1:
+            attr_shape += attribution.shape[1:]
+        return torch.mean(attribution.view(attr_shape), dim=1, keepdim=False)
+
+
+def nonlinear(
+    module: Module,
+    inputs: Tensor,
+    outputs: Tensor,
+    grad_input: Tensor,
+    grad_output: Tensor,
+    eps: float = 1e-10,
+):
+    r"""
+    grad_input: (dLoss / dprev_layer_out, dLoss / wij, dLoss / bij)
+    grad_output: (dLoss / dlayer_out)
+    https://github.com/pytorch/pytorch/issues/12331
+    """
+    delta_in, delta_out = _compute_diffs(inputs, outputs)
+
+    new_grad_inp = list(grad_input)
+
+    # supported non-linear modules take only single tensor as input hence accessing
+    # only the first element in `grad_input` and `grad_output`
+    new_grad_inp[0] = torch.where(
+        abs(delta_in) < eps, new_grad_inp[0], grad_output[0] * delta_out / delta_in
+    )
+
+    # If the module is invalid, save the newly computed gradients
+    # The original_grad_input will be overridden later in the Tensor hook
+    if module.is_invalid:
+        module.saved_grad = new_grad_inp[0]
+    return new_grad_inp
+
+
+def softmax(
+    module: Module,
+    inputs: Tensor,
+    outputs: Tensor,
+    grad_input: Tensor,
+    grad_output: Tensor,
+    eps: float = 1e-10,
+):
+    delta_in, delta_out = _compute_diffs(inputs, outputs)
+
+    new_grad_inp = list(grad_input)
+    grad_input_unnorm = torch.where(
+        abs(delta_in) < eps, new_grad_inp[0], grad_output[0] * delta_out / delta_in
+    )
+    # normalizing
+    n = grad_input[0].numel()
+
+    # updating only the first half
+    new_grad_inp[0] = grad_input_unnorm - grad_input_unnorm.sum() * 1 / n
+    return new_grad_inp
+
+
+def maxpool1d(
+    module: Module,
+    inputs: Tensor,
+    outputs: Tensor,
+    grad_input: Tensor,
+    grad_output: Tensor,
+    eps: float = 1e-10,
+):
+    return maxpool(
+        module,
+        F.max_pool1d,
+        F.max_unpool1d,
+        inputs,
+        outputs,
+        grad_input,
+        grad_output,
+        eps=eps,
+    )
+
+
+def maxpool2d(
+    module: Module,
+    inputs: Tensor,
+    outputs: Tensor,
+    grad_input: Tensor,
+    grad_output: Tensor,
+    eps: float = 1e-10,
+):
+    return maxpool(
+        module,
+        F.max_pool2d,
+        F.max_unpool2d,
+        inputs,
+        outputs,
+        grad_input,
+        grad_output,
+        eps=eps,
+    )
+
+
+def maxpool3d(
+    module: Module, inputs, outputs, grad_input, grad_output, eps: float = 1e-10
+):
+    return maxpool(
+        module,
+        F.max_pool3d,
+        F.max_unpool3d,
+        inputs,
+        outputs,
+        grad_input,
+        grad_output,
+        eps=eps,
+    )
+
+
+def maxpool(
+    module: Module,
+    pool_func: Callable,
+    unpool_func: Callable,
+    inputs,
+    outputs,
+    grad_input,
+    grad_output,
+    eps: float = 1e-10,
+):
+    with torch.no_grad():
+        input, input_ref = inputs.chunk(2)
+        output, output_ref = outputs.chunk(2)
+
+        delta_in = input - input_ref
+        delta_in = torch.cat(2 * [delta_in])
+        # Extracts cross maximum between the outputs of maxpool for the
+        # actual inputs and its corresponding references. In case the delta outputs
+        # for the references are larger the method relies on the references and
+        # corresponding gradients to compute the multiplies and contributions.
+        delta_out_xmax = torch.max(output, output_ref)
+        delta_out = torch.cat([delta_out_xmax - output_ref, output - delta_out_xmax])
+
+        _, indices = pool_func(
+            module.input,
+            module.kernel_size,
+            module.stride,
+            module.padding,
+            module.dilation,
+            module.ceil_mode,
+            True,
+        )
+        grad_output_updated = grad_output[0]
+        unpool_grad_out_delta, unpool_grad_out_ref_delta = torch.chunk(
+            unpool_func(
+                grad_output_updated * delta_out,
+                indices,
+                module.kernel_size,
+                module.stride,
+                module.padding,
+                list(cast(torch.Size, module.input.shape)),
+            ),
+            2,
+        )
+
+    unpool_grad_out_delta = unpool_grad_out_delta + unpool_grad_out_ref_delta
+    unpool_grad_out_delta = torch.cat(2 * [unpool_grad_out_delta])
+
+    # If the module is invalid, we need to recompute the grad_input
+    if module.is_invalid:
+        original_grad_input = grad_input
+        grad_input = (
+            unpool_func(
+                grad_output_updated,
+                indices,
+                module.kernel_size,
+                module.stride,
+                module.padding,
+                list(cast(torch.Size, module.input.shape)),
+            ),
+        )
+    if grad_input[0].shape != inputs.shape:
+        raise AssertionError(
+            "A problem occurred during maxpool modul's backward pass. "
+            "The gradients with respect to inputs include only a "
+            "subset of inputs. More details about this issue can "
+            "be found here: "
+            "https://pytorch.org/docs/stable/"
+            "nn.html#torch.nn.Module.register_backward_hook "
+            "This can happen for example if you attribute to the outputs of a "
+            "MaxPool. As a workaround, please, attribute to the inputs of "
+            "the following layer."
+        )
+
+    new_grad_inp = torch.where(
+        abs(delta_in) < eps, grad_input[0], unpool_grad_out_delta / delta_in
+    )
+    # If the module is invalid, save the newly computed gradients
+    # The original_grad_input will be overridden later in the Tensor hook
+    if module.is_invalid:
+        module.saved_grad = new_grad_inp
+        return original_grad_input
+    else:
+        return (new_grad_inp,)
+
+
+def _compute_diffs(inputs: Tensor, outputs: Tensor) -> Tuple[Tensor, Tensor]:
+    input, input_ref = inputs.chunk(2)
+    # if the model is a single non-linear module and we apply Rescale rule on it
+    # we might not be able to perform chunk-ing because the output of the module is
+    # usually being replaced by model output.
+    output, output_ref = outputs.chunk(2)
+    delta_in = input - input_ref
+    delta_out = output - output_ref
+
+    return torch.cat(2 * [delta_in]), torch.cat(2 * [delta_out])
+
+
+SUPPORTED_NON_LINEAR = {
+    nn.ReLU: nonlinear,
+    nn.ELU: nonlinear,
+    nn.LeakyReLU: nonlinear,
+    nn.Sigmoid: nonlinear,
+    nn.Tanh: nonlinear,
+    nn.Softplus: nonlinear,
+    nn.MaxPool1d: maxpool1d,
+    nn.MaxPool2d: maxpool2d,
+    nn.MaxPool3d: maxpool3d,
+    nn.Softmax: softmax,
+}

captum/attr/_core/feature_ablation.py

@@ -0,0 +1,591 @@
+#!/usr/bin/env python3
+
+import math
+from typing import Any, Callable, cast, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _expand_additional_forward_args,
+    _expand_target,
+    _format_additional_forward_args,
+    _format_output,
+    _format_tensor_into_tuples,
+    _is_tuple,
+    _run_forward,
+)
+from captum._utils.progress import progress
+from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._utils.attribution import PerturbationAttribution
+from captum.attr._utils.common import _format_input_baseline
+from captum.log import log_usage
+from torch import dtype, Tensor
+
+
+class FeatureAblation(PerturbationAttribution):
+    r"""
+    A perturbation based approach to computing attribution, involving
+    replacing each input feature with a given baseline / reference, and
+    computing the difference in output. By default, each scalar value within
+    each input tensor is taken as a feature and replaced independently. Passing
+    a feature mask, allows grouping features to be ablated together. This can
+    be used in cases such as images, where an entire segment or region
+    can be ablated, measuring the importance of the segment (feature group).
+    Each input scalar in the group will be given the same attribution value
+    equal to the change in target as a result of ablating the entire feature
+    group.
+
+    The forward function can either return a scalar per example or a tensor
+    of a fixed sized tensor (or scalar value) for the full batch, i.e. the
+    output does not grow as the batch size increase. If the output is fixed
+    we consider this model to be an "aggregation" of the inputs. In the fixed
+    sized output mode we require `perturbations_per_eval == 1` and the
+    `feature_mask` to be either `None` or for all of them to have 1 as their
+    first dimension (i.e. a feature mask requires to be applied to all inputs).
+    """
+
+    def __init__(self, forward_func: Callable) -> None:
+        r"""
+        Args:
+
+            forward_func (callable): The forward function of the model or
+                        any modification of it
+        """
+        PerturbationAttribution.__init__(self, forward_func)
+        self.use_weights = False
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        feature_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
+        perturbations_per_eval: int = 1,
+        show_progress: bool = False,
+        **kwargs: Any,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which ablation
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define reference value which replaces each
+                        feature when ablated.
+                        Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or
+                          broadcastable to match the dimensions of inputs
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. For all other types,
+                        the given argument is used for all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            feature_mask (tensor or tuple of tensors, optional):
+                        feature_mask defines a mask for the input, grouping
+                        features which should be ablated together. feature_mask
+                        should contain the same number of tensors as inputs.
+                        Each tensor should
+                        be the same size as the corresponding input or
+                        broadcastable to match the input tensor. Each tensor
+                        should contain integers in the range 0 to num_features
+                        - 1, and indices corresponding to the same feature should
+                        have the same value.
+                        Note that features within each input tensor are ablated
+                        independently (not across tensors).
+                        If the forward function returns a single scalar per batch,
+                        we enforce that the first dimension of each mask must be 1,
+                        since attributions are returned batch-wise rather than per
+                        example, so the attributions must correspond to the
+                        same features (indices) in each input example.
+                        If None, then a feature mask is constructed which assigns
+                        each scalar within a tensor as a separate feature, which
+                        is ablated independently.
+                        Default: None
+            perturbations_per_eval (int, optional): Allows ablation of multiple
+                        features to be processed simultaneously in one call to
+                        forward_fn.
+                        Each forward pass will contain a maximum of
+                        perturbations_per_eval * #examples samples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain at most
+                        (perturbations_per_eval * #examples) / num_devices
+                        samples.
+                        If the forward function's number of outputs does not
+                        change as the batch size grows (e.g. if it outputs a
+                        scalar value), you must set perturbations_per_eval to 1
+                        and use a single feature mask to describe the features
+                        for all examples in the batch.
+                        Default: 1
+            show_progress (bool, optional): Displays the progress of computation.
+                        It will try to use tqdm if available for advanced features
+                        (e.g. time estimation). Otherwise, it will fallback to
+                        a simple output of progress.
+                        Default: False
+            **kwargs (Any, optional): Any additional arguments used by child
+                        classes of FeatureAblation (such as Occlusion) to construct
+                        ablations. These arguments are ignored when using
+                        FeatureAblation directly.
+                        Default: None
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        The attributions with respect to each input feature.
+                        If the forward function returns
+                        a scalar value per example, attributions will be
+                        the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If the forward function returns a scalar per batch, then
+                        attribution tensor(s) will have first dimension 1 and
+                        the remaining dimensions will match the input.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple of tensors is provided for inputs, a
+                        tuple of corresponding sized tensors is returned.
+
+
+        Examples::
+
+            >>> # SimpleClassifier takes a single input tensor of size Nx4x4,
+            >>> # and returns an Nx3 tensor of class probabilities.
+            >>> net = SimpleClassifier()
+            >>> # Generating random input with size 2 x 4 x 4
+            >>> input = torch.randn(2, 4, 4)
+            >>> # Defining FeatureAblation interpreter
+            >>> ablator = FeatureAblation(net)
+            >>> # Computes ablation attribution, ablating each of the 16
+            >>> # scalar input independently.
+            >>> attr = ablator.attribute(input, target=1)
+
+            >>> # Alternatively, we may want to ablate features in groups, e.g.
+            >>> # grouping each 2x2 square of the inputs and ablating them together.
+            >>> # This can be done by creating a feature mask as follows, which
+            >>> # defines the feature groups, e.g.:
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # With this mask, all inputs with the same value are ablated
+            >>> # simultaneously, and the attribution for each input in the same
+            >>> # group (0, 1, 2, and 3) per example are the same.
+            >>> # The attributions can be calculated as follows:
+            >>> # feature mask has dimensions 1 x 4 x 4
+            >>> feature_mask = torch.tensor([[[0,0,1,1],[0,0,1,1],
+            >>>                             [2,2,3,3],[2,2,3,3]]])
+            >>> attr = ablator.attribute(input, target=1, feature_mask=feature_mask)
+        """
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        num_examples = inputs[0].shape[0]
+        feature_mask = (
+            _format_tensor_into_tuples(feature_mask)
+            if feature_mask is not None
+            else None
+        )
+        assert (
+            isinstance(perturbations_per_eval, int) and perturbations_per_eval >= 1
+        ), "Perturbations per evaluation must be an integer and at least 1."
+        with torch.no_grad():
+            if show_progress:
+                feature_counts = self._get_feature_counts(
+                    inputs, feature_mask, **kwargs
+                )
+                total_forwards = (
+                    sum(
+                        math.ceil(count / perturbations_per_eval)
+                        for count in feature_counts
+                    )
+                    + 1
+                )  # add 1 for the initial eval
+                attr_progress = progress(
+                    desc=f"{self.get_name()} attribution", total=total_forwards
+                )
+                attr_progress.update(0)
+
+            # Computes initial evaluation with all features, which is compared
+            # to each ablated result.
+            initial_eval = _run_forward(
+                self.forward_func, inputs, target, additional_forward_args
+            )
+
+            if show_progress:
+                attr_progress.update()
+
+            agg_output_mode = FeatureAblation._find_output_mode(
+                perturbations_per_eval, feature_mask
+            )
+
+            # get as a 2D tensor (if it is not a scalar)
+            if isinstance(initial_eval, torch.Tensor):
+                initial_eval = initial_eval.reshape(1, -1)
+                num_outputs = initial_eval.shape[1]
+            else:
+                num_outputs = 1
+
+            if not agg_output_mode:
+                assert (
+                    isinstance(initial_eval, torch.Tensor)
+                    and num_outputs == num_examples
+                ), (
+                    "expected output of `forward_func` to have "
+                    + "`batch_size` elements for perturbations_per_eval > 1 "
+                    + "and all feature_mask.shape[0] > 1"
+                )
+
+            # Initialize attribution totals and counts
+            attrib_type = cast(
+                dtype,
+                initial_eval.dtype
+                if isinstance(initial_eval, Tensor)
+                else type(initial_eval),
+            )
+
+            total_attrib = [
+                torch.zeros(
+                    (num_outputs,) + input.shape[1:],
+                    dtype=attrib_type,
+                    device=input.device,
+                )
+                for input in inputs
+            ]
+
+            # Weights are used in cases where ablations may be overlapping.
+            if self.use_weights:
+                weights = [
+                    torch.zeros(
+                        (num_outputs,) + input.shape[1:], device=input.device
+                    ).float()
+                    for input in inputs
+                ]
+
+            # Iterate through each feature tensor for ablation
+            for i in range(len(inputs)):
+                # Skip any empty input tensors
+                if torch.numel(inputs[i]) == 0:
+                    continue
+
+                for (
+                    current_inputs,
+                    current_add_args,
+                    current_target,
+                    current_mask,
+                ) in self._ith_input_ablation_generator(
+                    i,
+                    inputs,
+                    additional_forward_args,
+                    target,
+                    baselines,
+                    feature_mask,
+                    perturbations_per_eval,
+                    **kwargs,
+                ):
+                    # modified_eval dimensions: 1D tensor with length
+                    # equal to #num_examples * #features in batch
+                    modified_eval = _run_forward(
+                        self.forward_func,
+                        current_inputs,
+                        current_target,
+                        current_add_args,
+                    )
+
+                    if show_progress:
+                        attr_progress.update()
+
+                    # (contains 1 more dimension than inputs). This adds extra
+                    # dimensions of 1 to make the tensor broadcastable with the inputs
+                    # tensor.
+                    if not isinstance(modified_eval, torch.Tensor):
+                        eval_diff = initial_eval - modified_eval
+                    else:
+                        if not agg_output_mode:
+                            assert (
+                                modified_eval.numel() == current_inputs[0].shape[0]
+                            ), """expected output of forward_func to grow with
+                            batch_size. If this is not the case for your model
+                            please set perturbations_per_eval = 1"""
+
+                        eval_diff = (
+                            initial_eval - modified_eval.reshape((-1, num_outputs))
+                        ).reshape((-1, num_outputs) + (len(inputs[i].shape) - 1) * (1,))
+                        eval_diff = eval_diff.to(total_attrib[i].device)
+                    if self.use_weights:
+                        weights[i] += current_mask.float().sum(dim=0)
+                    total_attrib[i] += (eval_diff * current_mask.to(attrib_type)).sum(
+                        dim=0
+                    )
+
+            if show_progress:
+                attr_progress.close()
+
+            # Divide total attributions by counts and return formatted attributions
+            if self.use_weights:
+                attrib = tuple(
+                    single_attrib.float() / weight
+                    for single_attrib, weight in zip(total_attrib, weights)
+                )
+            else:
+                attrib = tuple(total_attrib)
+            _result = _format_output(is_inputs_tuple, attrib)
+        return _result
+
+    def _ith_input_ablation_generator(
+        self,
+        i,
+        inputs,
+        additional_args,
+        target,
+        baselines,
+        input_mask,
+        perturbations_per_eval,
+        **kwargs,
+    ):
+        """
+        This method return an generator of ablation perturbations of the i-th input
+
+        Returns:
+            ablation_iter (generator): yields each perturbation to be evaluated
+                        as a tuple (inputs, additional_forward_args, targets, mask).
+        """
+        extra_args = {}
+        for key, value in kwargs.items():
+            # For any tuple argument in kwargs, we choose index i of the tuple.
+            if isinstance(value, tuple):
+                extra_args[key] = value[i]
+            else:
+                extra_args[key] = value
+
+        input_mask = input_mask[i] if input_mask is not None else None
+        min_feature, num_features, input_mask = self._get_feature_range_and_mask(
+            inputs[i], input_mask, **extra_args
+        )
+        num_examples = inputs[0].shape[0]
+        perturbations_per_eval = min(perturbations_per_eval, num_features)
+        baseline = baselines[i] if isinstance(baselines, tuple) else baselines
+        if isinstance(baseline, torch.Tensor):
+            baseline = baseline.reshape((1,) + baseline.shape)
+
+        if perturbations_per_eval > 1:
+            # Repeat features and additional args for batch size.
+            all_features_repeated = [
+                torch.cat([inputs[j]] * perturbations_per_eval, dim=0)
+                for j in range(len(inputs))
+            ]
+            additional_args_repeated = (
+                _expand_additional_forward_args(additional_args, perturbations_per_eval)
+                if additional_args is not None
+                else None
+            )
+            target_repeated = _expand_target(target, perturbations_per_eval)
+        else:
+            all_features_repeated = list(inputs)
+            additional_args_repeated = additional_args
+            target_repeated = target
+
+        num_features_processed = min_feature
+        while num_features_processed < num_features:
+            current_num_ablated_features = min(
+                perturbations_per_eval, num_features - num_features_processed
+            )
+
+            # Store appropriate inputs and additional args based on batch size.
+            if current_num_ablated_features != perturbations_per_eval:
+                current_features = [
+                    feature_repeated[0 : current_num_ablated_features * num_examples]
+                    for feature_repeated in all_features_repeated
+                ]
+                current_additional_args = (
+                    _expand_additional_forward_args(
+                        additional_args, current_num_ablated_features
+                    )
+                    if additional_args is not None
+                    else None
+                )
+                current_target = _expand_target(target, current_num_ablated_features)
+            else:
+                current_features = all_features_repeated
+                current_additional_args = additional_args_repeated
+                current_target = target_repeated
+
+            # Store existing tensor before modifying
+            original_tensor = current_features[i]
+            # Construct ablated batch for features in range num_features_processed
+            # to num_features_processed + current_num_ablated_features and return
+            # mask with same size as ablated batch. ablated_features has dimension
+            # (current_num_ablated_features, num_examples, inputs[i].shape[1:])
+            # Note that in the case of sparse tensors, the second dimension
+            # may not necessarilly be num_examples and will match the first
+            # dimension of this tensor.
+            current_reshaped = current_features[i].reshape(
+                (current_num_ablated_features, -1) + current_features[i].shape[1:]
+            )
+
+            ablated_features, current_mask = self._construct_ablated_input(
+                current_reshaped,
+                input_mask,
+                baseline,
+                num_features_processed,
+                num_features_processed + current_num_ablated_features,
+                **extra_args,
+            )
+
+            # current_features[i] has dimension
+            # (current_num_ablated_features * num_examples, inputs[i].shape[1:]),
+            # which can be provided to the model as input.
+            current_features[i] = ablated_features.reshape(
+                (-1,) + ablated_features.shape[2:]
+            )
+            yield tuple(
+                current_features
+            ), current_additional_args, current_target, current_mask
+            # Replace existing tensor at index i.
+            current_features[i] = original_tensor
+            num_features_processed += current_num_ablated_features
+
+    def _construct_ablated_input(
+        self, expanded_input, input_mask, baseline, start_feature, end_feature, **kwargs
+    ):
+        r"""
+        Ablates given expanded_input tensor with given feature mask, feature range,
+        and baselines. expanded_input shape is (`num_features`, `num_examples`, ...)
+        with remaining dimensions corresponding to remaining original tensor
+        dimensions and `num_features` = `end_feature` - `start_feature`.
+        input_mask has same number of dimensions as original input tensor (one less
+        than `expanded_input`), and can have first dimension either 1, applying same
+        feature mask to all examples, or `num_examples`. baseline is expected to
+        be broadcastable to match `expanded_input`.
+
+        This method returns the ablated input tensor, which has the same
+        dimensionality as `expanded_input` as well as the corresponding mask with
+        either the same dimensionality as `expanded_input` or second dimension
+        being 1. This mask contains 1s in locations which have been ablated (and
+        thus counted towards ablations for that feature) and 0s otherwise.
+        """
+        current_mask = torch.stack(
+            [input_mask == j for j in range(start_feature, end_feature)], dim=0
+        ).long()
+        ablated_tensor = (
+            expanded_input * (1 - current_mask).to(expanded_input.dtype)
+        ) + (baseline * current_mask.to(expanded_input.dtype))
+        return ablated_tensor, current_mask
+
+    def _get_feature_range_and_mask(self, input, input_mask, **kwargs):
+        if input_mask is None:
+            # Obtain feature mask for selected input tensor, matches size of
+            # 1 input example, (1 x inputs[i].shape[1:])
+            input_mask = torch.reshape(
+                torch.arange(torch.numel(input[0]), device=input.device),
+                input[0:1].shape,
+            ).long()
+        return (
+            torch.min(input_mask).item(),
+            torch.max(input_mask).item() + 1,
+            input_mask,
+        )
+
+    def _get_feature_counts(self, inputs, feature_mask, **kwargs):
+        """return the numbers of input features"""
+        if not feature_mask:
+            return tuple(inp[0].numel() if inp.numel() else 0 for inp in inputs)
+
+        return tuple(
+            (mask.max() - mask.min()).item() + 1
+            if mask is not None
+            else (inp[0].numel() if inp.numel() else 0)
+            for inp, mask in zip(inputs, feature_mask)
+        )
+
+    @staticmethod
+    def _find_output_mode(
+        perturbations_per_eval: int,
+        feature_mask: Union[None, TensorOrTupleOfTensorsGeneric],
+    ) -> bool:
+        """
+        Returns True if the output mode is "aggregation output mode"
+
+        Aggregation output mode is defined as: when there is no 1:1 correspondence
+        with the `num_examples` (`batch_size`) and the amount of outputs your model
+        produces, i.e. the model output does not grow in size as the input becomes
+        larger.
+
+        We assume this is the case if `perturbations_per_eval == 1`
+        and your feature mask is None or is associated to all
+        examples in a batch (fm.shape[0] == 1 for all fm in feature_mask).
+        """
+        return perturbations_per_eval == 1 and (
+            feature_mask is None
+            or all(len(sm.shape) == 0 or sm.shape[0] == 1 for sm in feature_mask)
+        )

captum/attr/_core/feature_permutation.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, Tuple, Union
+
+import torch
+from captum._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._core.feature_ablation import FeatureAblation
+from captum.log import log_usage
+from torch import Tensor
+
+
+def _permute_feature(x: Tensor, feature_mask: Tensor) -> Tensor:
+    n = x.size(0)
+    assert n > 1, "cannot permute features with batch_size = 1"
+
+    perm = torch.randperm(n)
+    no_perm = torch.arange(n)
+    while (perm == no_perm).all():
+        perm = torch.randperm(n)
+
+    return (x[perm] * feature_mask.to(dtype=x.dtype)) + (
+        x * feature_mask.bitwise_not().to(dtype=x.dtype)
+    )
+
+
+class FeaturePermutation(FeatureAblation):
+    r"""
+    A perturbation based approach to compute attribution, which
+    takes each input feature, permutes the feature values within a batch,
+    and computes the difference between original and shuffled outputs for
+    the given batch. This difference signifies the feature importance
+    for the permuted feature.
+
+    Example pseudocode for the algorithm is as follows::
+
+        perm_feature_importance(batch):
+            importance = dict()
+            baseline_error = error_metric(model(batch), batch_labels)
+            for each feature:
+                permute this feature across the batch
+                error = error_metric(model(permuted_batch), batch_labels)
+                importance[feature] = baseline_error - error
+                "un-permute" the feature across the batch
+
+            return importance
+
+    It should be noted that the `error_metric` must be called in the
+    `forward_func`. You do not need to have an error metric, e.g. you
+    could simply return the logits (the model output), but this may or may
+    not provide a meaningful attribution.
+
+    This method, unlike other attribution methods, requires a batch
+    of examples to compute attributions and cannot be performed on a single example.
+
+    By default, each scalar value within
+    each input tensor is taken as a feature and shuffled independently. Passing
+    a feature mask, allows grouping features to be shuffled together.
+    Each input scalar in the group will be given the same attribution value
+    equal to the change in target as a result of shuffling the entire feature
+    group.
+
+    The forward function can either return a scalar per example, or a single
+    scalar for the full batch. If a single scalar is returned for the batch,
+    `perturbations_per_eval` must be 1, and the returned attributions will have
+    first dimension 1, corresponding to feature importance across all
+    examples in the batch.
+
+    More information can be found in the permutation feature
+    importance algorithm description here:
+    https://christophm.github.io/interpretable-ml-book/feature-importance.html
+    """
+
+    def __init__(
+        self, forward_func: Callable, perm_func: Callable = _permute_feature
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable): The forward function of the model or
+                any modification of it
+            perm_func (callable, optional): A function that accepts a batch of
+                inputs and a feature mask, and "permutes" the feature using
+                feature mask across the batch. This defaults to a function
+                which applies a random permutation, this argument only needs
+                to be provided if a custom permutation behavior is desired.
+                Default: `_permute_feature`
+        """
+        FeatureAblation.__init__(self, forward_func=forward_func)
+        self.perm_func = perm_func
+
+    # suppressing error caused by the child class not having a matching
+    # signature to the parent
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None,
+        perturbations_per_eval: int = 1,
+        show_progress: bool = False,
+        **kwargs: Any,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        This function is almost equivalent to `FeatureAblation.attribute`. The
+        main difference is the way ablated examples are generated. Specifically
+        they are generated through the `perm_func`, as we set the baselines for
+        `FeatureAblation.attribute` to None.
+
+
+        Args:
+                inputs (tensor or tuple of tensors):  Input for which
+                            permutation attributions are computed. If
+                            forward_func takes a single tensor as input, a
+                            single input tensor should be provided.  If
+                            forward_func takes multiple tensors as input, a
+                            tuple of the input tensors should be provided. It is
+                            assumed that for all given input tensors, dimension
+                            0 corresponds to the number of examples (aka batch
+                            size), and if multiple input tensors are provided,
+                            the examples must be aligned appropriately.
+                target (int, tuple, tensor or list, optional):  Output indices for
+                            which difference is computed (for classification cases,
+                            this is usually the target class).
+                            If the network returns a scalar value per example,
+                            no target index is necessary.
+                            For general 2D outputs, targets can be either:
+
+                            - a single integer or a tensor containing a single
+                              integer, which is applied to all input examples
+
+                            - a list of integers or a 1D tensor, with length matching
+                              the number of examples in inputs (dim 0). Each integer
+                              is applied as the target for the corresponding example.
+
+                            For outputs with > 2 dimensions, targets can be either:
+
+                            - A single tuple, which contains #output_dims - 1
+                              elements. This target index is applied to all examples.
+
+                            - A list of tuples with length equal to the number of
+                              examples in inputs (dim 0), and each tuple containing
+                              #output_dims - 1 elements. Each tuple is applied as the
+                              target for the corresponding example.
+
+                            Default: None
+                additional_forward_args (any, optional): If the forward function
+                            requires additional arguments other than the inputs for
+                            which attributions should not be computed, this argument
+                            can be provided. It must be either a single additional
+                            argument of a Tensor or arbitrary (non-tuple) type or a
+                            tuple containing multiple additional arguments including
+                            tensors or any arbitrary python types. These arguments
+                            are provided to forward_func in order following the
+                            arguments in inputs.
+                            For a tensor, the first dimension of the tensor must
+                            correspond to the number of examples. For all other types,
+                            the given argument is used for all forward evaluations.
+                            Note that attributions are not computed with respect
+                            to these arguments.
+                            Default: None
+                feature_mask (tensor or tuple of tensors, optional):
+                            feature_mask defines a mask for the input, grouping
+                            features which should be ablated together. feature_mask
+                            should contain the same number of tensors as inputs.
+                            Each tensor should be the same size as the
+                            corresponding input or broadcastable to match the
+                            input tensor. Each tensor should contain integers in
+                            the range 0 to num_features - 1, and indices
+                            corresponding to the same feature should have the
+                            same value.  Note that features within each input
+                            tensor are ablated independently (not across
+                            tensors).
+
+                            The first dimension of each mask must be 1, as we require
+                            to have the same group of features for each input sample.
+
+                            If None, then a feature mask is constructed which assigns
+                            each scalar within a tensor as a separate feature, which
+                            is permuted independently.
+                            Default: None
+                perturbations_per_eval (int, optional): Allows permutations
+                            of multiple features to be processed simultaneously
+                            in one call to forward_fn.  Each forward pass will
+                            contain a maximum of perturbations_per_eval * #examples
+                            samples.  For DataParallel models, each batch is
+                            split among the available devices, so evaluations on
+                            each available device contain at most
+                            (perturbations_per_eval * #examples) / num_devices
+                            samples.
+                            If the forward function returns a single scalar per batch,
+                            perturbations_per_eval must be set to 1.
+                            Default: 1
+                show_progress (bool, optional): Displays the progress of computation.
+                            It will try to use tqdm if available for advanced features
+                            (e.g. time estimation). Otherwise, it will fallback to
+                            a simple output of progress.
+                            Default: False
+                **kwargs (Any, optional): Any additional arguments used by child
+                            classes of FeatureAblation (such as Occlusion) to construct
+                            ablations. These arguments are ignored when using
+                            FeatureAblation directly.
+                            Default: None
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        The attributions with respect to each input feature.
+                        If the forward function returns
+                        a scalar value per example, attributions will be
+                        the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If the forward function returns a scalar per batch, then
+                        attribution tensor(s) will have first dimension 1 and
+                        the remaining dimensions will match the input.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple of tensors is provided for inputs,
+                        a tuple of corresponding sized tensors is returned.
+
+
+        Examples::
+
+            >>> # SimpleClassifier takes a single input tensor of size Nx4x4,
+            >>> # and returns an Nx3 tensor of class probabilities.
+            >>> net = SimpleClassifier()
+            >>> # Generating random input with size 10 x 4 x 4
+            >>> input = torch.randn(10, 4, 4)
+            >>> # Defining FeaturePermutation interpreter
+            >>> feature_perm = FeaturePermutation(net)
+            >>> # Computes permutation attribution, shuffling each of the 16
+            >>> # scalar input independently.
+            >>> attr = feature_perm.attribute(input, target=1)
+
+            >>> # Alternatively, we may want to permute features in groups, e.g.
+            >>> # grouping each 2x2 square of the inputs and shuffling them together.
+            >>> # This can be done by creating a feature mask as follows, which
+            >>> # defines the feature groups, e.g.:
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # With this mask, all inputs with the same value are shuffled
+            >>> # simultaneously, and the attribution for each input in the same
+            >>> # group (0, 1, 2, and 3) per example are the same.
+            >>> # The attributions can be calculated as follows:
+            >>> # feature mask has dimensions 1 x 4 x 4
+            >>> feature_mask = torch.tensor([[[0,0,1,1],[0,0,1,1],
+            >>>                             [2,2,3,3],[2,2,3,3]]])
+            >>> attr = feature_perm.attribute(input, target=1,
+            >>>                               feature_mask=feature_mask)
+        """
+        return FeatureAblation.attribute.__wrapped__(
+            self,
+            inputs,
+            baselines=None,
+            target=target,
+            additional_forward_args=additional_forward_args,
+            feature_mask=feature_mask,
+            perturbations_per_eval=perturbations_per_eval,
+            show_progress=show_progress,
+            **kwargs,
+        )
+
+    def _construct_ablated_input(
+        self,
+        expanded_input: Tensor,
+        input_mask: Tensor,
+        baseline: Union[int, float, Tensor],
+        start_feature: int,
+        end_feature: int,
+        **kwargs: Any,
+    ) -> Tuple[Tensor, Tensor]:
+        r"""
+        This function permutes the features of `expanded_input` with a given
+        feature mask and feature range. Permutation occurs via calling
+        `self.perm_func` across each batch within `expanded_input`. As with
+        `FeatureAblation._construct_ablated_input`:
+        - `expanded_input.shape = (num_features, num_examples, ...)`
+        - `num_features = end_feature - start_feature` (i.e. start and end is a
+          half-closed interval)
+        - `input_mask` is a tensor of the same shape as one input, which
+          describes the locations of each feature via their "index"
+
+        Since `baselines` is set to None for `FeatureAblation.attribute, this
+        will be the zero tensor, however, it is not used.
+        """
+        assert input_mask.shape[0] == 1, (
+            "input_mask.shape[0] != 1: pass in one mask in order to permute"
+            "the same features for each input"
+        )
+        current_mask = torch.stack(
+            [input_mask == j for j in range(start_feature, end_feature)], dim=0
+        ).bool()
+
+        output = torch.stack(
+            [
+                self.perm_func(x, mask.squeeze(0))
+                for x, mask in zip(expanded_input, current_mask)
+            ]
+        )
+        return output, current_mask

captum/attr/_core/gradient_shap.py

@@ -0,0 +1,414 @@
+#!/usr/bin/env python3
+import typing
+from typing import Any, Callable, Tuple, Union
+
+import numpy as np
+import torch
+from captum._utils.common import _is_tuple
+from captum._utils.typing import (
+    BaselineType,
+    Literal,
+    TargetType,
+    Tensor,
+    TensorOrTupleOfTensorsGeneric,
+)
+from captum.attr._core.noise_tunnel import NoiseTunnel
+from captum.attr._utils.attribution import GradientAttribution
+from captum.attr._utils.common import (
+    _compute_conv_delta_and_format_attrs,
+    _format_callable_baseline,
+    _format_input_baseline,
+)
+from captum.log import log_usage
+
+
+class GradientShap(GradientAttribution):
+    r"""
+    Implements gradient SHAP based on the implementation from SHAP's primary
+    author. For reference, please view the original
+    `implementation
+    <https://github.com/slundberg/shap#deep-learning-example-with-gradientexplainer-tensorflowkeraspytorch-models>`_
+    and the paper: `A Unified Approach to Interpreting Model Predictions
+    <https://papers.nips.cc/paper/7062-a-unified-approach-to-interpreting-model-predictions>`_
+
+    GradientShap approximates SHAP values by computing the expectations of
+    gradients by randomly sampling from the distribution of baselines/references.
+    It adds white noise to each input sample `n_samples` times, selects a
+    random baseline from baselines' distribution and a random point along the
+    path between the baseline and the input, and computes the gradient of outputs
+    with respect to those selected random points. The final SHAP values represent
+    the expected values of gradients * (inputs - baselines).
+
+    GradientShap makes an assumption that the input features are independent
+    and that the explanation model is linear, meaning that the explanations
+    are modeled through the additive composition of feature effects.
+    Under those assumptions, SHAP value can be approximated as the expectation
+    of gradients that are computed for randomly generated `n_samples` input
+    samples after adding gaussian noise `n_samples` times to each input for
+    different baselines/references.
+
+    In some sense it can be viewed as an approximation of integrated gradients
+    by computing the expectations of gradients for different baselines.
+
+    Current implementation uses Smoothgrad from `NoiseTunnel` in order to
+    randomly draw samples from the distribution of baselines, add noise to input
+    samples and compute the expectation (smoothgrad).
+    """
+
+    def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> None:
+        r"""
+        Args:
+
+            forward_func (function): The forward function of the model or
+                       any modification of it.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                    model inputs' multiplier in the final attribution scores.
+                    In the literature this is also known as local vs global
+                    attribution. If inputs' multiplier isn't factored in
+                    then this type of attribution method is also called local
+                    attribution. If it is, then that type of attribution
+                    method is called global.
+                    More detailed can be found here:
+                    https://arxiv.org/abs/1711.06104
+
+                    In case of gradient shap, if `multiply_by_inputs`
+                    is set to True, the sensitivity scores of scaled inputs
+                    are being multiplied by (inputs - baselines).
+        """
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which SHAP attribution
+                        values are computed. If `forward_func` takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If `forward_func` takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (tensor, tuple of tensors, callable):
+                        Baselines define the starting point from which expectation
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          the first dimension equal to the number of examples
+                          in the baselines' distribution. The remaining dimensions
+                          must match with input tensor's dimension starting from
+                          the second dimension.
+
+                        - a tuple of tensors, if inputs is a tuple of tensors,
+                          with the first dimension of any tensor inside the tuple
+                          equal to the number of examples in the baseline's
+                          distribution. The remaining dimensions must match
+                          the dimensions of the corresponding input tensor
+                          starting from the second dimension.
+
+                        - callable function, optionally takes `inputs` as an
+                          argument and either returns a single tensor
+                          or a tuple of those.
+
+                        It is recommended that the number of samples in the baselines'
+                        tensors is larger than one.
+            n_samples (int, optional):  The number of randomly generated examples
+                        per sample in the input batch. Random examples are
+                        generated by adding gaussian random noise to each sample.
+                        Default: `5` if `n_samples` is not provided.
+            stdevs    (float, or a tuple of floats optional): The standard deviation
+                        of gaussian noise with zero mean that is added to each
+                        input in the batch. If `stdevs` is a single float value
+                        then that same value is used for all inputs. If it is
+                        a tuple, then it must have the same length as the inputs
+                        tuple. In this case, each stdev value in the stdevs tuple
+                        corresponds to the input with the same index in the inputs
+                        tuple.
+                        Default: 0.0
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It can contain a tuple of ND tensors or
+                        any arbitrary python type of any shape.
+                        In case of the ND tensor the first dimension of the
+                        tensor must correspond to the batch size. It will be
+                        repeated for each `n_steps` for each randomly generated
+                        input sample.
+                        Note that the gradients are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attribution score computed based on GradientSHAP with respect
+                        to each input feature. Attributions will always be
+                        the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple is provided for inputs, a tuple of
+                        corresponding sized tensors is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        This is computed using the property that the total
+                        sum of forward_func(inputs) - forward_func(baselines)
+                        must be very close to the total sum of the attributions
+                        based on GradientSHAP.
+                        Delta is calculated for each example in the input after adding
+                        `n_samples` times gaussian noise to each of them. Therefore,
+                        the dimensionality of the deltas tensor is equal to the
+                        `number of examples in the input` * `n_samples`
+                        The deltas are ordered by each input example and `n_samples`
+                        noisy samples generated for it.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> gradient_shap = GradientShap(net)
+            >>> input = torch.randn(3, 3, 32, 32, requires_grad=True)
+            >>> # choosing baselines randomly
+            >>> baselines = torch.randn(20, 3, 32, 32)
+            >>> # Computes gradient shap for the input
+            >>> # Attribution size matches input size: 3x3x32x32
+            >>> attribution = gradient_shap.attribute(input, baselines,
+                                                                target=5)
+
+        """
+        # since `baselines` is a distribution, we can generate it using a function
+        # rather than passing it as an input argument
+        baselines = _format_callable_baseline(baselines, inputs)
+        assert isinstance(baselines[0], torch.Tensor), (
+            "Baselines distribution has to be provided in a form "
+            "of a torch.Tensor {}.".format(baselines[0])
+        )
+
+        input_min_baseline_x_grad = InputBaselineXGradient(
+            self.forward_func, self.multiplies_by_inputs
+        )
+        input_min_baseline_x_grad.gradient_func = self.gradient_func
+
+        nt = NoiseTunnel(input_min_baseline_x_grad)
+
+        # NOTE: using attribute.__wrapped__ to not log
+        attributions = nt.attribute.__wrapped__(
+            nt,  # self
+            inputs,
+            nt_type="smoothgrad",
+            nt_samples=n_samples,
+            stdevs=stdevs,
+            draw_baseline_from_distrib=True,
+            baselines=baselines,
+            target=target,
+            additional_forward_args=additional_forward_args,
+            return_convergence_delta=return_convergence_delta,
+        )
+
+        return attributions
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+
+class InputBaselineXGradient(GradientAttribution):
+    def __init__(self, forward_func: Callable, multiply_by_inputs=True) -> None:
+        r"""
+        Args:
+
+            forward_func (function): The forward function of the model or
+                        any modification of it
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in
+                        then this type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of gradient shap, if `multiply_by_inputs`
+                        is set to True, the sensitivity scores of scaled inputs
+                        are being multiplied by (inputs - baselines).
+
+        """
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        ...
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+
+        rand_coefficient = torch.tensor(
+            np.random.uniform(0.0, 1.0, inputs[0].shape[0]),
+            device=inputs[0].device,
+            dtype=inputs[0].dtype,
+        )
+
+        input_baseline_scaled = tuple(
+            _scale_input(input, baseline, rand_coefficient)
+            for input, baseline in zip(inputs, baselines)
+        )
+        grads = self.gradient_func(
+            self.forward_func, input_baseline_scaled, target, additional_forward_args
+        )
+
+        if self.multiplies_by_inputs:
+            input_baseline_diffs = tuple(
+                input - baseline for input, baseline in zip(inputs, baselines)
+            )
+            attributions = tuple(
+                input_baseline_diff * grad
+                for input_baseline_diff, grad in zip(input_baseline_diffs, grads)
+            )
+        else:
+            attributions = grads
+
+        return _compute_conv_delta_and_format_attrs(
+            self,
+            return_convergence_delta,
+            attributions,
+            baselines,
+            inputs,
+            additional_forward_args,
+            target,
+            is_inputs_tuple,
+        )
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+
+def _scale_input(
+    input: Tensor, baseline: Union[Tensor, int, float], rand_coefficient: Tensor
+) -> Tensor:
+    # batch size
+    bsz = input.shape[0]
+    inp_shape_wo_bsz = input.shape[1:]
+    inp_shape = (bsz,) + tuple([1] * len(inp_shape_wo_bsz))
+
+    # expand and reshape the indices
+    rand_coefficient = rand_coefficient.view(inp_shape)
+
+    input_baseline_scaled = (
+        rand_coefficient * input + (1.0 - rand_coefficient) * baseline
+    ).requires_grad_()
+    return input_baseline_scaled

captum/attr/_core/guided_backprop_deconvnet.py

@@ -0,0 +1,322 @@
+#!/usr/bin/env python3
+import warnings
+from typing import Any, List, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from captum._utils.common import (
+    _format_output,
+    _format_tensor_into_tuples,
+    _is_tuple,
+    _register_backward_hook,
+)
+from captum._utils.gradient import (
+    apply_gradient_requirements,
+    undo_gradient_requirements,
+)
+from captum._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._utils.attribution import GradientAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+from torch.utils.hooks import RemovableHandle
+
+
+class ModifiedReluGradientAttribution(GradientAttribution):
+    def __init__(self, model: Module, use_relu_grad_output: bool = False) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance.
+        """
+        GradientAttribution.__init__(self, model)
+        self.model = model
+        self.backward_hooks: List[RemovableHandle] = []
+        self.use_relu_grad_output = use_relu_grad_output
+        assert isinstance(self.model, torch.nn.Module), (
+            "Given model must be an instance of torch.nn.Module to properly hook"
+            " ReLU layers."
+        )
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Computes attribution by overriding relu gradients. Based on constructor
+        flag use_relu_grad_output, performs either GuidedBackpropagation if False
+        and Deconvolution if True. This class is the parent class of both these
+        methods, more information on usage can be found in the docstrings for each
+        implementing class.
+        """
+
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+
+        inputs = _format_tensor_into_tuples(inputs)
+        gradient_mask = apply_gradient_requirements(inputs)
+
+        # set hooks for overriding ReLU gradients
+        warnings.warn(
+            "Setting backward hooks on ReLU activations."
+            "The hooks will be removed after the attribution is finished"
+        )
+        try:
+            self.model.apply(self._register_hooks)
+
+            gradients = self.gradient_func(
+                self.forward_func, inputs, target, additional_forward_args
+            )
+        finally:
+            self._remove_hooks()
+
+        undo_gradient_requirements(inputs, gradient_mask)
+        return _format_output(is_inputs_tuple, gradients)
+
+    def _register_hooks(self, module: Module):
+        if isinstance(module, torch.nn.ReLU):
+            hook = _register_backward_hook(module, self._backward_hook, self)
+            self.backward_hooks.append(hook)
+
+    def _backward_hook(
+        self,
+        module: Module,
+        grad_input: Union[Tensor, Tuple[Tensor, ...]],
+        grad_output: Union[Tensor, Tuple[Tensor, ...]],
+    ):
+        to_override_grads = grad_output if self.use_relu_grad_output else grad_input
+        if isinstance(to_override_grads, tuple):
+            return tuple(
+                F.relu(to_override_grad) for to_override_grad in to_override_grads
+            )
+        else:
+            return F.relu(to_override_grads)
+
+    def _remove_hooks(self):
+        for hook in self.backward_hooks:
+            hook.remove()
+
+
+class GuidedBackprop(ModifiedReluGradientAttribution):
+    r"""
+    Computes attribution using guided backpropagation. Guided backpropagation
+    computes the gradient of the target output with respect to the input,
+    but gradients of ReLU functions are overridden so that only
+    non-negative gradients are backpropagated.
+
+    More details regarding the guided backpropagation algorithm can be found
+    in the original paper here:
+    https://arxiv.org/abs/1412.6806
+
+    Warning: Ensure that all ReLU operations in the forward function of the
+    given model are performed using a module (nn.module.ReLU).
+    If nn.functional.ReLU is used, gradients are not overridden appropriately.
+    """
+
+    def __init__(self, model: Module) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place ReLU submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API.
+        """
+        ModifiedReluGradientAttribution.__init__(
+            self, model, use_relu_grad_output=False
+        )
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        The guided backprop gradients with respect to each
+                        input feature. Attributions will always
+                        be the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple is provided for inputs, a tuple of
+                        corresponding sized tensors is returned.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> gbp = GuidedBackprop(net)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes Guided Backprop attribution scores for class 3.
+            >>> attribution = gbp.attribute(input, target=3)
+        """
+        return super().attribute.__wrapped__(
+            self, inputs, target, additional_forward_args
+        )
+
+
+class Deconvolution(ModifiedReluGradientAttribution):
+    r"""
+    Computes attribution using deconvolution. Deconvolution
+    computes the gradient of the target output with respect to the input,
+    but gradients of ReLU functions are overridden so that the gradient
+    of the ReLU input is simply computed taking ReLU of the output gradient,
+    essentially only propagating non-negative gradients (without
+    dependence on the sign of the ReLU input).
+
+    More details regarding the deconvolution algorithm can be found
+    in these papers:
+    https://arxiv.org/abs/1311.2901
+    https://link.springer.com/chapter/10.1007/978-3-319-46466-4_8
+
+    Warning: Ensure that all ReLU operations in the forward function of the
+    given model are performed using a module (nn.module.ReLU).
+    If nn.functional.ReLU is used, gradients are not overridden appropriately.
+    """
+
+    def __init__(self, model: Module) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place ReLU submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API.
+        """
+        ModifiedReluGradientAttribution.__init__(self, model, use_relu_grad_output=True)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        The deconvolution attributions with respect to each
+                        input feature. Attributions will always
+                        be the same size as the provided inputs, with each value
+                        providing the attribution of the corresponding input index.
+                        If a single tensor is provided as inputs, a single tensor is
+                        returned. If a tuple is provided for inputs, a tuple of
+                        corresponding sized tensors is returned.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> deconv = Deconvolution(net)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes Deconvolution attribution scores for class 3.
+            >>> attribution = deconv.attribute(input, target=3)
+        """
+        return super().attribute.__wrapped__(
+            self, inputs, target, additional_forward_args
+        )

captum/attr/_core/guided_grad_cam.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+import warnings
+from typing import Any, List, Union
+
+import torch
+from captum._utils.common import _format_output, _format_tensor_into_tuples, _is_tuple
+from captum._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._core.guided_backprop_deconvnet import GuidedBackprop
+from captum.attr._core.layer.grad_cam import LayerGradCam
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class GuidedGradCam(GradientAttribution):
+    r"""
+    Computes element-wise product of guided backpropagation attributions
+    with upsampled (non-negative) GradCAM attributions.
+    GradCAM attributions are computed with respect to the layer
+    provided in the constructor, and attributions
+    are upsampled to match the input size. GradCAM is designed for
+    convolutional neural networks, and is usually applied to the last
+    convolutional layer.
+
+    Note that if multiple input tensors are provided, attributions for
+    each input tensor are computed by upsampling the GradCAM
+    attributions to match that input's dimensions. If interpolation is
+    not possible for the input tensor dimensions and interpolation mode,
+    then an empty tensor is returned in the attributions for the
+    corresponding position of that input tensor. This can occur if the
+    input tensor does not have the same number of dimensions as the chosen
+    layer's output or is not either 3D, 4D or 5D.
+
+    Note that attributions are only meaningful for input tensors
+    which are spatially alligned with the chosen layer, e.g. an input
+    image tensor for a convolutional layer.
+
+    More details regarding GuidedGradCAM can be found in the original
+    GradCAM paper here:
+    https://arxiv.org/pdf/1610.02391.pdf
+
+    Warning: Ensure that all ReLU operations in the forward function of the
+    given model are performed using a module (nn.module.ReLU).
+    If nn.functional.ReLU is used, gradients are not overridden appropriately.
+    """
+
+    def __init__(
+        self, model: Module, layer: Module, device_ids: Union[None, List[int]] = None
+    ) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place ReLU submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API
+                        starting from PyTorch v1.9.
+            layer (torch.nn.Module): Layer for which GradCAM attributions are computed.
+                          Currently, only layers with a single tensor output are
+                          supported.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself,
+                          then it is not necessary to provide this argument.
+        """
+        GradientAttribution.__init__(self, model)
+        self.grad_cam = LayerGradCam(model, layer, device_ids)
+        self.guided_backprop = GuidedBackprop(model)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        interpolate_mode: str = "nearest",
+        attribute_to_layer_input: bool = False,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which attributions
+                        are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            interpolate_mode (str, optional): Method for interpolation, which
+                        must be a valid input interpolation mode for
+                        torch.nn.functional. These methods are
+                        "nearest", "area", "linear" (3D-only), "bilinear"
+                        (4D-only), "bicubic" (4D-only), "trilinear" (5D-only)
+                        based on the number of dimensions of the chosen layer
+                        output (which must also match the number of
+                        dimensions for the input tensor). Note that
+                        the original GradCAM paper uses "bilinear"
+                        interpolation, but we default to "nearest" for
+                        applicability to any of 3D, 4D or 5D tensors.
+                        Default: "nearest"
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output in `LayerGradCam`.
+                        If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer inputs, otherwise it will be computed with respect
+                        to layer outputs.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+
+        Returns:
+            *tensor* of **attributions**:
+            - **attributions** (*tensor*):
+                    Element-wise product of (upsampled) GradCAM
+                    and Guided Backprop attributions.
+                    If a single tensor is provided as inputs, a single tensor is
+                    returned. If a tuple is provided for inputs, a tuple of
+                    corresponding sized tensors is returned.
+                    Attributions will be the same size as the provided inputs,
+                    with each value providing the attribution of the
+                    corresponding input index.
+                    If the GradCAM attributions cannot be upsampled to the shape
+                    of a given input tensor, None is returned in the corresponding
+                    index position.
+
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains an attribute conv4, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx50x8x8.
+            >>> # It is the last convolution layer, which is the recommended
+            >>> # use case for GuidedGradCAM.
+            >>> net = ImageClassifier()
+            >>> guided_gc = GuidedGradCam(net, net.conv4)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes guided GradCAM attributions for class 3.
+            >>> # attribution size matches input size, Nx3x32x32
+            >>> attribution = guided_gc.attribute(input, 3)
+        """
+        is_inputs_tuple = _is_tuple(inputs)
+        inputs = _format_tensor_into_tuples(inputs)
+        grad_cam_attr = self.grad_cam.attribute.__wrapped__(
+            self.grad_cam,  # self
+            inputs=inputs,
+            target=target,
+            additional_forward_args=additional_forward_args,
+            attribute_to_layer_input=attribute_to_layer_input,
+            relu_attributions=True,
+        )
+        if isinstance(grad_cam_attr, tuple):
+            assert len(grad_cam_attr) == 1, (
+                "GuidedGradCAM attributions for layer with multiple inputs / "
+                "outputs is not supported."
+            )
+            grad_cam_attr = grad_cam_attr[0]
+
+        guided_backprop_attr = self.guided_backprop.attribute.__wrapped__(
+            self.guided_backprop,  # self
+            inputs=inputs,
+            target=target,
+            additional_forward_args=additional_forward_args,
+        )
+        output_attr: List[Tensor] = []
+        for i in range(len(inputs)):
+            try:
+                output_attr.append(
+                    guided_backprop_attr[i]
+                    * LayerAttribution.interpolate(
+                        grad_cam_attr,
+                        inputs[i].shape[2:],
+                        interpolate_mode=interpolate_mode,
+                    )
+                )
+            except Exception:
+                warnings.warn(
+                    "Couldn't appropriately interpolate GradCAM attributions for some "
+                    "input tensors, returning empty tensor for corresponding "
+                    "attributions."
+                )
+                output_attr.append(torch.empty(0))
+
+        return _format_output(is_inputs_tuple, tuple(output_attr))

captum/attr/_core/input_x_gradient.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+from typing import Any, Callable
+
+from captum._utils.common import _format_output, _format_tensor_into_tuples, _is_tuple
+from captum._utils.gradient import (
+    apply_gradient_requirements,
+    undo_gradient_requirements,
+)
+from captum._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._utils.attribution import GradientAttribution
+from captum.log import log_usage
+
+
+class InputXGradient(GradientAttribution):
+    r"""
+    A baseline approach for computing the attribution. It multiplies input with
+    the gradient with respect to input.
+    https://arxiv.org/abs/1605.01713
+    """
+
+    def __init__(self, forward_func: Callable) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+        """
+        GradientAttribution.__init__(self, forward_func)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+
+        Returns:
+                *tensor* or tuple of *tensors* of **attributions**:
+                - **attributions** (*tensor* or tuple of *tensors*):
+                            The input x gradient with
+                            respect to each input feature. Attributions will always be
+                            the same size as the provided inputs, with each value
+                            providing the attribution of the corresponding input index.
+                            If a single tensor is provided as inputs, a single tensor is
+                            returned. If a tuple is provided for inputs, a tuple of
+                            corresponding sized tensors is returned.
+
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> # Generating random input with size 2x3x3x32
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Defining InputXGradient interpreter
+            >>> input_x_gradient = InputXGradient(net)
+            >>> # Computes inputXgradient for class 4.
+            >>> attribution = input_x_gradient.attribute(input, target=4)
+        """
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+
+        inputs = _format_tensor_into_tuples(inputs)
+        gradient_mask = apply_gradient_requirements(inputs)
+
+        gradients = self.gradient_func(
+            self.forward_func, inputs, target, additional_forward_args
+        )
+
+        attributions = tuple(
+            input * gradient for input, gradient in zip(inputs, gradients)
+        )
+
+        undo_gradient_requirements(inputs, gradient_mask)
+        return _format_output(is_inputs_tuple, attributions)
+
+    @property
+    def multiplies_by_inputs(self):
+        return True

captum/attr/_core/integrated_gradients.py

@@ -0,0 +1,390 @@
+#!/usr/bin/env python3
+import typing
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _expand_additional_forward_args,
+    _expand_target,
+    _format_additional_forward_args,
+    _format_output,
+    _is_tuple,
+)
+from captum._utils.typing import (
+    BaselineType,
+    Literal,
+    TargetType,
+    TensorOrTupleOfTensorsGeneric,
+)
+from captum.attr._utils.approximation_methods import approximation_parameters
+from captum.attr._utils.attribution import GradientAttribution
+from captum.attr._utils.batching import _batch_attribution
+from captum.attr._utils.common import (
+    _format_input_baseline,
+    _reshape_and_sum,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+
+
+class IntegratedGradients(GradientAttribution):
+    r"""
+    Integrated Gradients is an axiomatic model interpretability algorithm that
+    assigns an importance score to each input feature by approximating the
+    integral of gradients of the model's output with respect to the inputs
+    along the path (straight line) from given baselines / references to inputs.
+
+    Baselines can be provided as input arguments to attribute method.
+    To approximate the integral we can choose to use either a variant of
+    Riemann sum or Gauss-Legendre quadrature rule.
+
+    More details regarding the integrated gradients method can be found in the
+    original paper:
+    https://arxiv.org/abs/1703.01365
+
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                    modification of it
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                    model inputs' multiplier in the final attribution scores.
+                    In the literature this is also known as local vs global
+                    attribution. If inputs' multiplier isn't factored in,
+                    then that type of attribution method is also called local
+                    attribution. If it is, then that type of attribution
+                    method is called global.
+                    More detailed can be found here:
+                    https://arxiv.org/abs/1711.06104
+
+                    In case of integrated gradients, if `multiply_by_inputs`
+                    is set to True, final sensitivity scores are being multiplied by
+                    (inputs - baselines).
+        """
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    # The following overloaded method signatures correspond to the case where
+    # return_convergence_delta is False, then only attributions are returned,
+    # and when return_convergence_delta is True, the return type is
+    # a tuple with both attributions and deltas.
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: Literal[False] = False,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        *,
+        return_convergence_delta: Literal[True],
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: bool = False,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
+        r"""
+        This method attributes the output of the model with given target index
+        (in case it is provided, otherwise it assumes that output is a
+        scalar) to the inputs of the model using the approach described above.
+
+        In addition to that it also returns, if `return_convergence_delta` is
+        set to True, integral approximation delta based on the completeness
+        property of integrated gradients.
+
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which integrated
+                        gradients are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define the starting point from which integral
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. It will be
+                        repeated for each of `n_steps` along the integrated
+                        path. For all other types, the given argument is used
+                        for all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            n_steps (int, optional): The number of steps used by the approximation
+                        method. Default: 50.
+            method (string, optional): Method for approximating the integral,
+                        one of `riemann_right`, `riemann_left`, `riemann_middle`,
+                        `riemann_trapezoid` or `gausslegendre`.
+                        Default: `gausslegendre` if no method is provided.
+            internal_batch_size (int, optional): Divides total #steps * #examples
+                        data points into chunks of size at most internal_batch_size,
+                        which are computed (forward / backward passes)
+                        sequentially. internal_batch_size must be at least equal to
+                        #examples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain internal_batch_size / num_devices examples.
+                        If internal_batch_size is None, then all evaluations are
+                        processed in one batch.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                    convergence delta or not. If `return_convergence_delta`
+                    is set to True convergence delta will be returned in
+                    a tuple following attributions.
+                    Default: False
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                    Integrated gradients with respect to each input feature.
+                    attributions will always be the same size as the provided
+                    inputs, with each value providing the attribution of the
+                    corresponding input index.
+                    If a single tensor is provided as inputs, a single tensor is
+                    returned. If a tuple is provided for inputs, a tuple of
+                    corresponding sized tensors is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                    The difference between the total approximated and true
+                    integrated gradients. This is computed using the property
+                    that the total sum of forward_func(inputs) -
+                    forward_func(baselines) must equal the total sum of the
+                    integrated gradient.
+                    Delta is calculated per example, meaning that the number of
+                    elements in returned delta tensor is equal to the number of
+                    of examples in inputs.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> ig = IntegratedGradients(net)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes integrated gradients for class 3.
+            >>> attribution = ig.attribute(input, target=3)
+        """
+        # Keeps track whether original input is a tuple or not before
+        # converting it into a tuple.
+        is_inputs_tuple = _is_tuple(inputs)
+
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+
+        _validate_input(inputs, baselines, n_steps, method)
+
+        if internal_batch_size is not None:
+            num_examples = inputs[0].shape[0]
+            attributions = _batch_attribution(
+                self,
+                num_examples,
+                internal_batch_size,
+                n_steps,
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                method=method,
+            )
+        else:
+            attributions = self._attribute(
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                n_steps=n_steps,
+                method=method,
+            )
+
+        if return_convergence_delta:
+            start_point, end_point = baselines, inputs
+            # computes approximation error based on the completeness axiom
+            delta = self.compute_convergence_delta(
+                attributions,
+                start_point,
+                end_point,
+                additional_forward_args=additional_forward_args,
+                target=target,
+            )
+            return _format_output(is_inputs_tuple, attributions), delta
+        return _format_output(is_inputs_tuple, attributions)
+
+    def _attribute(
+        self,
+        inputs: Tuple[Tensor, ...],
+        baselines: Tuple[Union[Tensor, int, float], ...],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+    ) -> Tuple[Tensor, ...]:
+        if step_sizes_and_alphas is None:
+            # retrieve step size and scaling factor for specified
+            # approximation method
+            step_sizes_func, alphas_func = approximation_parameters(method)
+            step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
+        else:
+            step_sizes, alphas = step_sizes_and_alphas
+
+        # scale features and compute gradients. (batch size is abbreviated as bsz)
+        # scaled_features' dim -> (bsz * #steps x inputs[0].shape[1:], ...)
+        scaled_features_tpl = tuple(
+            torch.cat(
+                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
+            ).requires_grad_()
+            for input, baseline in zip(inputs, baselines)
+        )
+
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        # apply number of steps to additional forward args
+        # currently, number of steps is applied only to additional forward arguments
+        # that are nd-tensors. It is assumed that the first dimension is
+        # the number of batches.
+        # dim -> (bsz * #steps x additional_forward_args[0].shape[1:], ...)
+        input_additional_args = (
+            _expand_additional_forward_args(additional_forward_args, n_steps)
+            if additional_forward_args is not None
+            else None
+        )
+        expanded_target = _expand_target(target, n_steps)
+
+        # grads: dim -> (bsz * #steps x inputs[0].shape[1:], ...)
+        grads = self.gradient_func(
+            forward_fn=self.forward_func,
+            inputs=scaled_features_tpl,
+            target_ind=expanded_target,
+            additional_forward_args=input_additional_args,
+        )
+
+        # flattening grads so that we can multilpy it with step-size
+        # calling contiguous to avoid `memory whole` problems
+        scaled_grads = [
+            grad.contiguous().view(n_steps, -1)
+            * torch.tensor(step_sizes).view(n_steps, 1).to(grad.device)
+            for grad in grads
+        ]
+
+        # aggregates across all steps for each tensor in the input tuple
+        # total_grads has the same dimensionality as inputs
+        total_grads = tuple(
+            _reshape_and_sum(
+                scaled_grad, n_steps, grad.shape[0] // n_steps, grad.shape[1:]
+            )
+            for (scaled_grad, grad) in zip(scaled_grads, grads)
+        )
+
+        # computes attribution for each tensor in input tuple
+        # attributions has the same dimensionality as inputs
+        if not self.multiplies_by_inputs:
+            attributions = total_grads
+        else:
+            attributions = tuple(
+                total_grad * (input - baseline)
+                for total_grad, input, baseline in zip(total_grads, inputs, baselines)
+            )
+        return attributions
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs

captum/attr/_core/kernel_shap.py

@@ -0,0 +1,348 @@
+#!/usr/bin/env python3
+
+from typing import Any, Callable, Generator, Tuple, Union
+
+import torch
+from captum._utils.models.linear_model import SkLearnLinearRegression
+from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._core.lime import construct_feature_mask, Lime
+from captum.attr._utils.common import _format_input_baseline
+from captum.log import log_usage
+from torch import Tensor
+from torch.distributions.categorical import Categorical
+
+
+class KernelShap(Lime):
+    r"""
+    Kernel SHAP is a method that uses the LIME framework to compute
+    Shapley Values. Setting the loss function, weighting kernel and
+    regularization terms appropriately in the LIME framework allows
+    theoretically obtaining Shapley Values more efficiently than
+    directly computing Shapley Values.
+
+    More information regarding this method and proof of equivalence
+    can be found in the original paper here:
+    https://arxiv.org/abs/1705.07874
+    """
+
+    def __init__(self, forward_func: Callable) -> None:
+        r"""
+        Args:
+
+            forward_func (callable): The forward function of the model or
+                        any modification of it
+        """
+        Lime.__init__(
+            self,
+            forward_func,
+            interpretable_model=SkLearnLinearRegression(),
+            similarity_func=self.kernel_shap_similarity_kernel,
+            perturb_func=self.kernel_shap_perturb_generator,
+        )
+        self.inf_weight = 1000000.0
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        feature_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
+        n_samples: int = 25,
+        perturbations_per_eval: int = 1,
+        return_input_shape: bool = True,
+        show_progress: bool = False,
+    ) -> TensorOrTupleOfTensorsGeneric:
+        r"""
+        This method attributes the output of the model with given target index
+        (in case it is provided, otherwise it assumes that output is a
+        scalar) to the inputs of the model using the approach described above,
+        training an interpretable model based on KernelSHAP and returning a
+        representation of the interpretable model.
+
+        It is recommended to only provide a single example as input (tensors
+        with first dimension or batch size = 1). This is because LIME / KernelShap
+        is generally used for sample-based interpretability, training a separate
+        interpretable model to explain a model's prediction on each individual example.
+
+        A batch of inputs can also be provided as inputs, similar to
+        other perturbation-based attribution methods. In this case, if forward_fn
+        returns a scalar per example, attributions will be computed for each
+        example independently, with a separate interpretable model trained for each
+        example. Note that provided similarity and perturbation functions will be
+        provided each example separately (first dimension = 1) in this case.
+        If forward_fn returns a scalar per batch (e.g. loss), attributions will
+        still be computed using a single interpretable model for the full batch.
+        In this case, similarity and perturbation functions will be provided the
+        same original input containing the full batch.
+
+        The number of interpretable features is determined from the provided
+        feature mask, or if none is provided, from the default feature mask,
+        which considers each scalar input as a separate feature. It is
+        generally recommended to provide a feature mask which groups features
+        into a small number of interpretable features / components (e.g.
+        superpixels in images).
+
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which KernelShap
+                        is computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define the reference value which replaces each
+                        feature when the corresponding interpretable feature
+                        is set to 0.
+                        Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which surrogate model is trained
+                        (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. It will be
+                        repeated for each of `n_steps` along the integrated
+                        path. For all other types, the given argument is used
+                        for all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            feature_mask (tensor or tuple of tensors, optional):
+                        feature_mask defines a mask for the input, grouping
+                        features which correspond to the same
+                        interpretable feature. feature_mask
+                        should contain the same number of tensors as inputs.
+                        Each tensor should
+                        be the same size as the corresponding input or
+                        broadcastable to match the input tensor. Values across
+                        all tensors should be integers in the range 0 to
+                        num_interp_features - 1, and indices corresponding to the
+                        same feature should have the same value.
+                        Note that features are grouped across tensors
+                        (unlike feature ablation and occlusion), so
+                        if the same index is used in different tensors, those
+                        features are still grouped and added simultaneously.
+                        If None, then a feature mask is constructed which assigns
+                        each scalar within a tensor as a separate feature.
+                        Default: None
+            n_samples (int, optional):  The number of samples of the original
+                        model used to train the surrogate interpretable model.
+                        Default: `50` if `n_samples` is not provided.
+            perturbations_per_eval (int, optional): Allows multiple samples
+                        to be processed simultaneously in one call to forward_fn.
+                        Each forward pass will contain a maximum of
+                        perturbations_per_eval * #examples samples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain at most
+                        (perturbations_per_eval * #examples) / num_devices
+                        samples.
+                        If the forward function returns a single scalar per batch,
+                        perturbations_per_eval must be set to 1.
+                        Default: 1
+            return_input_shape (bool, optional): Determines whether the returned
+                        tensor(s) only contain the coefficients for each interp-
+                        retable feature from the trained surrogate model, or
+                        whether the returned attributions match the input shape.
+                        When return_input_shape is True, the return type of attribute
+                        matches the input shape, with each element containing the
+                        coefficient of the corresponding interpretable feature.
+                        All elements with the same value in the feature mask
+                        will contain the same coefficient in the returned
+                        attributions. If return_input_shape is False, a 1D
+                        tensor is returned, containing only the coefficients
+                        of the trained interpretable model, with length
+                        num_interp_features.
+            show_progress (bool, optional): Displays the progress of computation.
+                        It will try to use tqdm if available for advanced features
+                        (e.g. time estimation). Otherwise, it will fallback to
+                        a simple output of progress.
+                        Default: False
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        The attributions with respect to each input feature.
+                        If return_input_shape = True, attributions will be
+                        the same size as the provided inputs, with each value
+                        providing the coefficient of the corresponding
+                        interpretale feature.
+                        If return_input_shape is False, a 1D
+                        tensor is returned, containing only the coefficients
+                        of the trained interpreatable models, with length
+                        num_interp_features.
+        Examples::
+            >>> # SimpleClassifier takes a single input tensor of size Nx4x4,
+            >>> # and returns an Nx3 tensor of class probabilities.
+            >>> net = SimpleClassifier()
+
+            >>> # Generating random input with size 1 x 4 x 4
+            >>> input = torch.randn(1, 4, 4)
+
+            >>> # Defining KernelShap interpreter
+            >>> ks = KernelShap(net)
+            >>> # Computes attribution, with each of the 4 x 4 = 16
+            >>> # features as a separate interpretable feature
+            >>> attr = ks.attribute(input, target=1, n_samples=200)
+
+            >>> # Alternatively, we can group each 2x2 square of the inputs
+            >>> # as one 'interpretable' feature and perturb them together.
+            >>> # This can be done by creating a feature mask as follows, which
+            >>> # defines the feature groups, e.g.:
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 0 | 0 | 1 | 1 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # | 2 | 2 | 3 | 3 |
+            >>> # +---+---+---+---+
+            >>> # With this mask, all inputs with the same value are set to their
+            >>> # baseline value, when the corresponding binary interpretable
+            >>> # feature is set to 0.
+            >>> # The attributions can be calculated as follows:
+            >>> # feature mask has dimensions 1 x 4 x 4
+            >>> feature_mask = torch.tensor([[[0,0,1,1],[0,0,1,1],
+            >>>                             [2,2,3,3],[2,2,3,3]]])
+
+            >>> # Computes KernelSHAP attributions with feature mask.
+            >>> attr = ks.attribute(input, target=1, feature_mask=feature_mask)
+        """
+        formatted_inputs, baselines = _format_input_baseline(inputs, baselines)
+        feature_mask, num_interp_features = construct_feature_mask(
+            feature_mask, formatted_inputs
+        )
+        num_features_list = torch.arange(num_interp_features, dtype=torch.float)
+        denom = num_features_list * (num_interp_features - num_features_list)
+        probs = (num_interp_features - 1) / denom
+        probs[0] = 0.0
+        return self._attribute_kwargs(
+            inputs=inputs,
+            baselines=baselines,
+            target=target,
+            additional_forward_args=additional_forward_args,
+            feature_mask=feature_mask,
+            n_samples=n_samples,
+            perturbations_per_eval=perturbations_per_eval,
+            return_input_shape=return_input_shape,
+            num_select_distribution=Categorical(probs),
+            show_progress=show_progress,
+        )
+
+    def kernel_shap_similarity_kernel(
+        self, _, __, interpretable_sample: Tensor, **kwargs
+    ) -> Tensor:
+        assert (
+            "num_interp_features" in kwargs
+        ), "Must provide num_interp_features to use default similarity kernel"
+        num_selected_features = int(interpretable_sample.sum(dim=1).item())
+        num_features = kwargs["num_interp_features"]
+        if num_selected_features == 0 or num_selected_features == num_features:
+            # weight should be theoretically infinite when
+            # num_selected_features = 0 or num_features
+            # enforcing that trained linear model must satisfy
+            # end-point criteria. In practice, it is sufficient to
+            # make this weight substantially larger so setting this
+            # weight to 1000000 (all other weights are 1).
+            similarities = self.inf_weight
+        else:
+            similarities = 1.0
+        return torch.tensor([similarities])
+
+    def kernel_shap_perturb_generator(
+        self, original_inp: Union[Tensor, Tuple[Tensor, ...]], **kwargs
+    ) -> Generator[Tensor, None, None]:
+        r"""
+        Perturbations are sampled by the following process:
+         - Choose k (number of selected features), based on the distribution
+                p(k) = (M - 1) / (k * (M - k))
+            where M is the total number of features in the interpretable space
+         - Randomly select a binary vector with k ones, each sample is equally
+            likely. This is done by generating a random vector of normal
+            values and thresholding based on the top k elements.
+
+         Since there are M choose k vectors with k ones, this weighted sampling
+         is equivalent to applying the Shapley kernel for the sample weight,
+         defined as:
+         k(M, k) = (M - 1) / (k * (M - k) * (M choose k))
+        """
+        assert (
+            "num_select_distribution" in kwargs and "num_interp_features" in kwargs
+        ), (
+            "num_select_distribution and num_interp_features are necessary"
+            " to use kernel_shap_perturb_func"
+        )
+        if isinstance(original_inp, Tensor):
+            device = original_inp.device
+        else:
+            device = original_inp[0].device
+        num_features = kwargs["num_interp_features"]
+        yield torch.ones(1, num_features, device=device, dtype=torch.long)
+        yield torch.zeros(1, num_features, device=device, dtype=torch.long)
+        while True:
+            num_selected_features = kwargs["num_select_distribution"].sample()
+            rand_vals = torch.randn(1, num_features)
+            threshold = torch.kthvalue(
+                rand_vals, num_features - num_selected_features
+            ).values.item()
+            yield (rand_vals > threshold).to(device=device).long()

captum/attr/_core/layer/grad_cam.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from captum._utils.common import (
+    _format_additional_forward_args,
+    _format_output,
+    _format_tensor_into_tuples,
+)
+from captum._utils.gradient import compute_layer_gradients_and_eval
+from captum._utils.typing import TargetType
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerGradCam(LayerAttribution, GradientAttribution):
+    r"""
+    Computes GradCAM attribution for chosen layer. GradCAM is designed for
+    convolutional neural networks, and is usually applied to the last
+    convolutional layer.
+
+    GradCAM computes the gradients of the target output with respect to
+    the given layer, averages for each output channel (dimension 2 of
+    output), and multiplies the average gradient for each channel by the
+    layer activations. The results are summed over all channels.
+
+    Note that in the original GradCAM algorithm described in the paper,
+    ReLU is applied to the output, returning only non-negative attributions.
+    For providing more flexibility to the user, we choose to not perform the
+    ReLU internally by default and return the sign information. To match the
+    original GradCAM algorithm, it is necessary to pass the parameter
+    relu_attributions=True to apply ReLU on the final
+    attributions or alternatively only visualize the positive attributions.
+
+    Note: this procedure sums over the second dimension (# of channels),
+    so the output of GradCAM attributions will have a second
+    dimension of 1, but all other dimensions will match that of the layer
+    output.
+
+    GradCAM attributions are generally upsampled and can be viewed as a
+    mask to the input, since a convolutional layer output generally
+    matches the input image spatially. This upsampling can be performed
+    using LayerAttribution.interpolate, as shown in the example below.
+
+    More details regarding the GradCAM method can be found in the
+    original paper here:
+    https://arxiv.org/pdf/1610.02391.pdf
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                          Output size of attribute matches this layer's output
+                          dimensions, except for dimension 2, which will be 1,
+                          since GradCAM sums over channels.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself,
+                          then it is not necessary to provide this argument.
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        attribute_to_layer_input: bool = False,
+        relu_attributions: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which attributions
+                        are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attributions with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to the
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Note that currently it is assumed that either the input
+                        or the outputs of internal layers, depending on whether we
+                        attribute to the input or output, are single tensors.
+                        Support for multiple tensors will be added later.
+                        Default: False
+            relu_attributions (bool, optional): Indicates whether to
+                        apply a ReLU operation on the final attribution,
+                        returning only non-negative attributions. Setting this
+                        flag to True matches the original GradCAM algorithm,
+                        otherwise, by default, both positive and negative
+                        attributions are returned.
+                        Default: False
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attributions based on GradCAM method.
+                        Attributions will be the same size as the
+                        output of the given layer, except for dimension 2,
+                        which will be 1 due to summing over channels.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains a layer conv4, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx50x8x8.
+            >>> # It is the last convolution layer, which is the recommended
+            >>> # use case for GradCAM.
+            >>> net = ImageClassifier()
+            >>> layer_gc = LayerGradCam(net, net.conv4)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes layer GradCAM for class 3.
+            >>> # attribution size matches layer output except for dimension
+            >>> # 1, so dimensions of attr would be Nx1x8x8.
+            >>> attr = layer_gc.attribute(input, 3)
+            >>> # GradCAM attributions are often upsampled and viewed as a
+            >>> # mask to the input, since the convolutional layer output
+            >>> # spatially matches the original input image.
+            >>> # This can be done with LayerAttribution's interpolate method.
+            >>> upsampled_attr = LayerAttribution.interpolate(attr, (32, 32))
+        """
+        inputs = _format_tensor_into_tuples(inputs)
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        # Returns gradient of output with respect to
+        # hidden layer and hidden layer evaluated at each input.
+        layer_gradients, layer_evals = compute_layer_gradients_and_eval(
+            self.forward_func,
+            self.layer,
+            inputs,
+            target,
+            additional_forward_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        summed_grads = tuple(
+            torch.mean(
+                layer_grad,
+                dim=tuple(x for x in range(2, len(layer_grad.shape))),
+                keepdim=True,
+            )
+            if len(layer_grad.shape) > 2
+            else layer_grad
+            for layer_grad in layer_gradients
+        )
+
+        scaled_acts = tuple(
+            torch.sum(summed_grad * layer_eval, dim=1, keepdim=True)
+            for summed_grad, layer_eval in zip(summed_grads, layer_evals)
+        )
+        if relu_attributions:
+            scaled_acts = tuple(F.relu(scaled_act) for scaled_act in scaled_acts)
+        return _format_output(len(scaled_acts) > 1, scaled_acts)

captum/attr/_core/layer/internal_influence.py

@@ -0,0 +1,309 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _expand_additional_forward_args,
+    _expand_target,
+    _format_additional_forward_args,
+    _format_output,
+)
+from captum._utils.gradient import compute_layer_gradients_and_eval
+from captum._utils.typing import BaselineType, TargetType
+from captum.attr._utils.approximation_methods import approximation_parameters
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.attr._utils.batching import _batch_attribution
+from captum.attr._utils.common import (
+    _format_input_baseline,
+    _reshape_and_sum,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class InternalInfluence(LayerAttribution, GradientAttribution):
+    r"""
+    Computes internal influence by approximating the integral of gradients
+    for a particular layer along the path from a baseline input to the
+    given input.
+    If no baseline is provided, the default baseline is the zero tensor.
+    More details on this approach can be found here:
+    https://arxiv.org/pdf/1802.03788.pdf
+
+    Note that this method is similar to applying integrated gradients and
+    taking the layer as input, integrating the gradient of the layer with
+    respect to the output.
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                          Output size of attribute matches this layer's input or
+                          output dimensions, depending on whether we attribute to
+                          the inputs or outputs of the layer, corresponding to
+                          attribution of each neuron in the input or output of
+                          this layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself,
+                          then it is not necessary to provide this argument.
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which internal
+                        influence is computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define a starting point from which integral
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. It will be
+                        repeated for each of `n_steps` along the integrated
+                        path. For all other types, the given argument is used
+                        for all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            n_steps (int, optional): The number of steps used by the approximation
+                        method. Default: 50.
+            method (string, optional): Method for approximating the integral,
+                        one of `riemann_right`, `riemann_left`, `riemann_middle`,
+                        `riemann_trapezoid` or `gausslegendre`.
+                        Default: `gausslegendre` if no method is provided.
+            internal_batch_size (int, optional): Divides total #steps * #examples
+                        data points into chunks of size at most internal_batch_size,
+                        which are computed (forward / backward passes)
+                        sequentially. internal_batch_size must be at least equal to
+                        #examples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain internal_batch_size / num_devices examples.
+                        If internal_batch_size is None, then all evaluations
+                        are processed in one batch.
+                        Default: None
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer inputs, otherwise it will be computed with respect
+                        to layer outputs.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Internal influence of each neuron in given
+                        layer output. Attributions will always be the same size
+                        as the output or input of the given layer depending on
+                        whether `attribute_to_layer_input` is set to `False` or
+                        `True`respectively.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx12x32x32.
+            >>> net = ImageClassifier()
+            >>> layer_int_inf = InternalInfluence(net, net.conv1)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes layer internal influence.
+            >>> # attribution size matches layer output, Nx12x32x32
+            >>> attribution = layer_int_inf.attribute(input)
+        """
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs, baselines, n_steps, method)
+        if internal_batch_size is not None:
+            num_examples = inputs[0].shape[0]
+            attrs = _batch_attribution(
+                self,
+                num_examples,
+                internal_batch_size,
+                n_steps,
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                method=method,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+        else:
+            attrs = self._attribute(
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                n_steps=n_steps,
+                method=method,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+
+        return attrs
+
+    def _attribute(
+        self,
+        inputs: Tuple[Tensor, ...],
+        baselines: Tuple[Union[Tensor, int, float], ...],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        attribute_to_layer_input: bool = False,
+        step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        if step_sizes_and_alphas is None:
+            # retrieve step size and scaling factor for specified approximation method
+            step_sizes_func, alphas_func = approximation_parameters(method)
+            step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
+        else:
+            step_sizes, alphas = step_sizes_and_alphas
+
+        # Compute scaled inputs from baseline to final input.
+        scaled_features_tpl = tuple(
+            torch.cat(
+                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
+            ).requires_grad_()
+            for input, baseline in zip(inputs, baselines)
+        )
+
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        # apply number of steps to additional forward args
+        # currently, number of steps is applied only to additional forward arguments
+        # that are nd-tensors. It is assumed that the first dimension is
+        # the number of batches.
+        # dim -> (bsz * #steps x additional_forward_args[0].shape[1:], ...)
+        input_additional_args = (
+            _expand_additional_forward_args(additional_forward_args, n_steps)
+            if additional_forward_args is not None
+            else None
+        )
+        expanded_target = _expand_target(target, n_steps)
+
+        # Returns gradient of output with respect to hidden layer.
+        layer_gradients, _ = compute_layer_gradients_and_eval(
+            forward_fn=self.forward_func,
+            layer=self.layer,
+            inputs=scaled_features_tpl,
+            target_ind=expanded_target,
+            additional_forward_args=input_additional_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+        # flattening grads so that we can multiply it with step-size
+        # calling contiguous to avoid `memory whole` problems
+        scaled_grads = tuple(
+            layer_grad.contiguous().view(n_steps, -1)
+            * torch.tensor(step_sizes).view(n_steps, 1).to(layer_grad.device)
+            for layer_grad in layer_gradients
+        )
+
+        # aggregates across all steps for each tensor in the input tuple
+        attrs = tuple(
+            _reshape_and_sum(
+                scaled_grad, n_steps, inputs[0].shape[0], layer_grad.shape[1:]
+            )
+            for scaled_grad, layer_grad in zip(scaled_grads, layer_gradients)
+        )
+        return _format_output(len(attrs) > 1, attrs)

captum/attr/_core/layer/layer_activation.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from captum._utils.common import _format_output
+from captum._utils.gradient import _forward_layer_eval
+from captum._utils.typing import ModuleOrModuleList
+from captum.attr._utils.attribution import LayerAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerActivation(LayerAttribution):
+    r"""
+    Computes activation of selected layer for given input.
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: ModuleOrModuleList,
+        device_ids: Union[None, List[int]] = None,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+            layer (torch.nn.Module or list(torch.nn.Module)): Layer or layers
+                          for which attributions are computed.
+                          Output size of attribute matches this layer's input or
+                          output dimensions, depending on whether we attribute to
+                          the inputs or outputs of the layer, corresponding to
+                          attribution of each neuron in the input or output of
+                          this layer. If multiple layers are provided, attributions
+                          are returned as a list, each element corresponding to the
+                          activations of the corresponding layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself,
+                          then it is not necessary to provide this argument.
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        additional_forward_args: Any = None,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer
+                        activation is computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+
+        Returns:
+            *tensor* or tuple of *tensors* or *list* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors* or *list*):
+                        Activation of each neuron in given layer output.
+                        Attributions will always be the same size as the
+                        output of the given layer.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+                        If multiple layers are provided, attributions
+                        are returned as a list, each element corresponding to the
+                        activations of the corresponding layer.
+
+
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx12x32x32.
+            >>> net = ImageClassifier()
+            >>> layer_act = LayerActivation(net, net.conv1)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes layer activation.
+            >>> # attribution is layer output, with size Nx12x32x32
+            >>> attribution = layer_cond.attribute(input)
+        """
+        with torch.no_grad():
+            layer_eval = _forward_layer_eval(
+                self.forward_func,
+                inputs,
+                self.layer,
+                additional_forward_args,
+                device_ids=self.device_ids,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+        if isinstance(self.layer, Module):
+            return _format_output(len(layer_eval) > 1, layer_eval)
+        else:
+            return [
+                _format_output(len(single_layer_eval) > 1, single_layer_eval)
+                for single_layer_eval in layer_eval
+            ]
+
+    @property
+    def multiplies_by_inputs(self):
+        return True

captum/attr/_core/layer/layer_conductance.py

@@ -0,0 +1,395 @@
+#!/usr/bin/env python3
+import typing
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _expand_additional_forward_args,
+    _expand_target,
+    _format_additional_forward_args,
+    _format_output,
+)
+from captum._utils.gradient import compute_layer_gradients_and_eval
+from captum._utils.typing import BaselineType, Literal, TargetType
+from captum.attr._utils.approximation_methods import approximation_parameters
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.attr._utils.batching import _batch_attribution
+from captum.attr._utils.common import (
+    _format_input_baseline,
+    _reshape_and_sum,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerConductance(LayerAttribution, GradientAttribution):
+    r"""
+    Computes conductance with respect to the given layer. The
+    returned output is in the shape of the layer's output, showing the total
+    conductance of each hidden layer neuron.
+
+    The details of the approach can be found here:
+    https://arxiv.org/abs/1805.12233
+    https://arxiv.org/pdf/1807.09946.pdf
+
+    Note that this provides the total conductance of each neuron in the
+    layer's output. To obtain the breakdown of a neuron's conductance by input
+    features, utilize NeuronConductance instead, and provide the target
+    neuron index.
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                          Output size of attribute matches this layer's input or
+                          output dimensions, depending on whether we attribute to
+                          the inputs or outputs of the layer, corresponding to
+                          attribution of each neuron in the input or output of
+                          this layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself,
+                          then it is not necessary to provide this argument.
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        *,
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool = False,
+    ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: Literal[False] = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[
+            None, int, float, Tensor, Tuple[Union[int, float, Tensor], ...]
+        ] = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[
+        Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer
+                        conductance is computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define the starting point from which integral
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. It will be repeated
+                        for each of `n_steps` along the integrated path.
+                        For all other types, the given argument is used for
+                        all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            n_steps (int, optional): The number of steps used by the approximation
+                        method. Default: 50.
+            method (string, optional): Method for approximating the integral,
+                        one of `riemann_right`, `riemann_left`, `riemann_middle`,
+                        `riemann_trapezoid` or `gausslegendre`.
+                        Default: `gausslegendre` if no method is provided.
+            internal_batch_size (int, optional): Divides total #steps * #examples
+                        data points into chunks of size at most internal_batch_size,
+                        which are computed (forward / backward passes)
+                        sequentially. internal_batch_size must be at least equal to
+                        2 * #examples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain internal_batch_size / num_devices examples.
+                        If internal_batch_size is None, then all evaluations are
+                        processed in one batch.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer inputs, otherwise it will be computed with respect
+                        to layer outputs.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Conductance of each neuron in given layer input or
+                        output. Attributions will always be the same size as
+                        the input or output of the given layer, depending on
+                        whether we attribute to the inputs or outputs
+                        of the layer which is decided by the input flag
+                        `attribute_to_layer_input`.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        The difference between the total
+                        approximated and true conductance.
+                        This is computed using the property that the total sum of
+                        forward_func(inputs) - forward_func(baselines) must equal
+                        the total sum of the attributions.
+                        Delta is calculated per example, meaning that the number of
+                        elements in returned delta tensor is equal to the number of
+                        of examples in inputs.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx12x32x32.
+            >>> net = ImageClassifier()
+            >>> layer_cond = LayerConductance(net, net.conv1)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes layer conductance for class 3.
+            >>> # attribution size matches layer output, Nx12x32x32
+            >>> attribution = layer_cond.attribute(input, target=3)
+        """
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs, baselines, n_steps, method)
+
+        num_examples = inputs[0].shape[0]
+        if internal_batch_size is not None:
+            num_examples = inputs[0].shape[0]
+            attrs = _batch_attribution(
+                self,
+                num_examples,
+                internal_batch_size,
+                n_steps + 1,
+                include_endpoint=True,
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                method=method,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+
+        else:
+            attrs = self._attribute(
+                inputs=inputs,
+                baselines=baselines,
+                target=target,
+                additional_forward_args=additional_forward_args,
+                n_steps=n_steps,
+                method=method,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+
+        is_layer_tuple = isinstance(attrs, tuple)
+        attributions = attrs if is_layer_tuple else (attrs,)
+
+        if return_convergence_delta:
+            start_point, end_point = baselines, inputs
+            delta = self.compute_convergence_delta(
+                attributions,
+                start_point,
+                end_point,
+                target=target,
+                additional_forward_args=additional_forward_args,
+            )
+            return _format_output(is_layer_tuple, attributions), delta
+        return _format_output(is_layer_tuple, attributions)
+
+    def _attribute(
+        self,
+        inputs: Tuple[Tensor, ...],
+        baselines: Tuple[Union[Tensor, int, float], ...],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        attribute_to_layer_input: bool = False,
+        step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        num_examples = inputs[0].shape[0]
+        if step_sizes_and_alphas is None:
+            # Retrieve scaling factors for specified approximation method
+            step_sizes_func, alphas_func = approximation_parameters(method)
+            alphas = alphas_func(n_steps + 1)
+        else:
+            _, alphas = step_sizes_and_alphas
+        # Compute scaled inputs from baseline to final input.
+        scaled_features_tpl = tuple(
+            torch.cat(
+                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
+            ).requires_grad_()
+            for input, baseline in zip(inputs, baselines)
+        )
+
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        # apply number of steps to additional forward args
+        # currently, number of steps is applied only to additional forward arguments
+        # that are nd-tensors. It is assumed that the first dimension is
+        # the number of batches.
+        # dim -> (#examples * #steps x additional_forward_args[0].shape[1:], ...)
+        input_additional_args = (
+            _expand_additional_forward_args(additional_forward_args, n_steps + 1)
+            if additional_forward_args is not None
+            else None
+        )
+        expanded_target = _expand_target(target, n_steps + 1)
+
+        # Conductance Gradients - Returns gradient of output with respect to
+        # hidden layer and hidden layer evaluated at each input.
+        (layer_gradients, layer_evals,) = compute_layer_gradients_and_eval(
+            forward_fn=self.forward_func,
+            layer=self.layer,
+            inputs=scaled_features_tpl,
+            additional_forward_args=input_additional_args,
+            target_ind=expanded_target,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        # Compute differences between consecutive evaluations of layer_eval.
+        # This approximates the total input gradient of each step multiplied
+        # by the step size.
+        grad_diffs = tuple(
+            layer_eval[num_examples:] - layer_eval[:-num_examples]
+            for layer_eval in layer_evals
+        )
+
+        # Element-wise multiply gradient of output with respect to hidden layer
+        # and summed gradients with respect to input (chain rule) and sum
+        # across stepped inputs.
+        attributions = tuple(
+            _reshape_and_sum(
+                grad_diff * layer_gradient[:-num_examples],
+                n_steps,
+                num_examples,
+                layer_eval.shape[1:],
+            )
+            for layer_gradient, layer_eval, grad_diff in zip(
+                layer_gradients, layer_evals, grad_diffs
+            )
+        )
+        return _format_output(len(attributions) > 1, attributions)
+
+    @property
+    def multiplies_by_inputs(self):
+        return True

captum/attr/_core/layer/layer_deep_lift.py

@@ -0,0 +1,682 @@
+#!/usr/bin/env python3
+import typing
+from typing import Any, Callable, cast, Sequence, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _expand_target,
+    _format_additional_forward_args,
+    _format_baseline,
+    _format_tensor_into_tuples,
+    ExpansionTypes,
+)
+from captum._utils.gradient import compute_layer_gradients_and_eval
+from captum._utils.typing import (
+    BaselineType,
+    Literal,
+    TargetType,
+    TensorOrTupleOfTensorsGeneric,
+)
+from captum.attr._core.deep_lift import DeepLift, DeepLiftShap
+from captum.attr._utils.attribution import LayerAttribution
+from captum.attr._utils.common import (
+    _call_custom_attribution_func,
+    _compute_conv_delta_and_format_attrs,
+    _format_callable_baseline,
+    _tensorize_baseline,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerDeepLift(LayerAttribution, DeepLift):
+    r"""
+    Implements DeepLIFT algorithm for the layer based on the following paper:
+    Learning Important Features Through Propagating Activation Differences,
+    Avanti Shrikumar, et. al.
+    https://arxiv.org/abs/1704.02685
+
+    and the gradient formulation proposed in:
+    Towards better understanding of gradient-based attribution methods for
+    deep neural networks,  Marco Ancona, et.al.
+    https://openreview.net/pdf?id=Sy21R9JAW
+
+    This implementation supports only Rescale rule. RevealCancel rule will
+    be supported in later releases.
+    Although DeepLIFT's(Rescale Rule) attribution quality is comparable with
+    Integrated Gradients, it runs significantly faster than Integrated
+    Gradients and is preferred for large datasets.
+
+    Currently we only support a limited number of non-linear activations
+    but the plan is to expand the list in the future.
+
+    Note: As we know, currently we cannot access the building blocks,
+    of PyTorch's built-in LSTM, RNNs and GRUs such as Tanh and Sigmoid.
+    Nonetheless, it is possible to build custom LSTMs, RNNS and GRUs
+    with performance similar to built-in ones using TorchScript.
+    More details on how to build custom RNNs can be found here:
+    https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/
+    """
+
+    def __init__(
+        self,
+        model: Module,
+        layer: Module,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place nonlinear submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API
+                        starting from PyTorch v1.9.
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                        The size and dimensionality of the attributions
+                        corresponds to the size and dimensionality of the layer's
+                        input or output depending on whether we attribute to the
+                        inputs or outputs of the layer.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in
+                        then that type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of Layer DeepLift, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores
+                        are being multiplied by
+                        layer activations for inputs - layer activations for baselines.
+                        This flag applies only if `custom_attribution_func` is
+                        set to None.
+        """
+        LayerAttribution.__init__(self, model, layer)
+        DeepLift.__init__(self, model)
+        self.model = model
+        self._multiply_by_inputs = multiply_by_inputs
+
+    # Ignoring mypy error for inconsistent signature with DeepLift
+    @typing.overload  # type: ignore
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer
+                        attributions are computed. If forward_func takes a
+                        single tensor as input, a single input tensor should be
+                        provided. If forward_func takes multiple tensors as input,
+                        a tuple of the input tensors should be provided. It is
+                        assumed that for all given input tensors, dimension 0
+                        corresponds to the number of examples (aka batch size),
+                        and if multiple input tensors are provided, the examples
+                        must be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define reference samples that are compared with
+                        the inputs. In order to assign attribution scores DeepLift
+                        computes the differences between the inputs/outputs and
+                        corresponding references.
+                        Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+
+                          - either a tensor with matching dimensions to
+                            corresponding tensor in the inputs' tuple
+                            or the first dimension is one and the remaining
+                            dimensions match with the corresponding
+                            input tensor.
+
+                          - or a scalar, corresponding to a tensor in the
+                            inputs' tuple. This scalar value is broadcasted
+                            for corresponding input tensor.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+            custom_attribution_func (callable, optional): A custom function for
+                        computing final attribution scores. This function can take
+                        at least one and at most three arguments with the
+                        following signature:
+
+                        - custom_attribution_func(multipliers)
+                        - custom_attribution_func(multipliers, inputs)
+                        - custom_attribution_func(multipliers, inputs, baselines)
+
+                        In case this function is not provided, we use the default
+                        logic defined as: multipliers * (inputs - baselines)
+                        It is assumed that all input arguments, `multipliers`,
+                        `inputs` and `baselines` are provided in tuples of same length.
+                        `custom_attribution_func` returns a tuple of attribution
+                        tensors that have the same length as the `inputs`.
+                        Default: None
+
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                Attribution score computed based on DeepLift's rescale rule with
+                respect to layer's inputs or outputs. Attributions will always be the
+                same size as the provided layer's inputs or outputs, depending on
+                whether we attribute to the inputs or outputs of the layer.
+                If the layer input / output is a single tensor, then
+                just a tensor is returned; if the layer input / output
+                has multiple tensors, then a corresponding tuple
+                of tensors is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                This is computed using the property that the total sum of
+                forward_func(inputs) - forward_func(baselines) must equal the
+                total sum of the attributions computed based on DeepLift's
+                rescale rule.
+                Delta is calculated per example, meaning that the number of
+                elements in returned delta tensor is equal to the number of
+                of examples in input.
+                Note that the logic described for deltas is guaranteed
+                when the default logic for attribution computations is used,
+                meaning that the `custom_attribution_func=None`, otherwise
+                it is not guaranteed and depends on the specifics of the
+                `custom_attribution_func`.
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> # creates an instance of LayerDeepLift to interpret target
+            >>> # class 1 with respect to conv4 layer.
+            >>> dl = LayerDeepLift(net, net.conv4)
+            >>> input = torch.randn(1, 3, 32, 32, requires_grad=True)
+            >>> # Computes deeplift attribution scores for conv4 layer and class 3.
+            >>> attribution = dl.attribute(input, target=1)
+        """
+        inputs = _format_tensor_into_tuples(inputs)
+        baselines = _format_baseline(baselines, inputs)
+        _validate_input(inputs, baselines)
+
+        baselines = _tensorize_baseline(inputs, baselines)
+
+        main_model_hooks = []
+        try:
+            main_model_hooks = self._hook_main_model()
+
+            self.model.apply(
+                lambda mod: self._register_hooks(
+                    mod, attribute_to_layer_input=attribute_to_layer_input
+                )
+            )
+
+            additional_forward_args = _format_additional_forward_args(
+                additional_forward_args
+            )
+            expanded_target = _expand_target(
+                target, 2, expansion_type=ExpansionTypes.repeat
+            )
+            wrapped_forward_func = self._construct_forward_func(
+                self.model,
+                (inputs, baselines),
+                expanded_target,
+                additional_forward_args,
+            )
+
+            def chunk_output_fn(out: TensorOrTupleOfTensorsGeneric) -> Sequence:
+                if isinstance(out, Tensor):
+                    return out.chunk(2)
+                return tuple(out_sub.chunk(2) for out_sub in out)
+
+            gradients, attrs = compute_layer_gradients_and_eval(
+                wrapped_forward_func,
+                self.layer,
+                inputs,
+                attribute_to_layer_input=attribute_to_layer_input,
+                output_fn=lambda out: chunk_output_fn(out),
+            )
+
+            attr_inputs = tuple(map(lambda attr: attr[0], attrs))
+            attr_baselines = tuple(map(lambda attr: attr[1], attrs))
+            gradients = tuple(map(lambda grad: grad[0], gradients))
+
+            if custom_attribution_func is None:
+                if self.multiplies_by_inputs:
+                    attributions = tuple(
+                        (input - baseline) * gradient
+                        for input, baseline, gradient in zip(
+                            attr_inputs, attr_baselines, gradients
+                        )
+                    )
+                else:
+                    attributions = gradients
+            else:
+                attributions = _call_custom_attribution_func(
+                    custom_attribution_func, gradients, attr_inputs, attr_baselines
+                )
+        finally:
+            # remove hooks from all activations
+            self._remove_hooks(main_model_hooks)
+
+        return _compute_conv_delta_and_format_attrs(
+            self,
+            return_convergence_delta,
+            attributions,
+            baselines,
+            inputs,
+            additional_forward_args,
+            target,
+            cast(Union[Literal[True], Literal[False]], len(attributions) > 1),
+        )
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+
+class LayerDeepLiftShap(LayerDeepLift, DeepLiftShap):
+    r"""
+    Extends LayerDeepLift and DeepLiftShap algorithms and approximates SHAP
+    values for given input `layer`.
+    For each input sample - baseline pair it computes DeepLift attributions
+    with respect to inputs or outputs of given `layer` averages
+    resulting attributions across baselines. Whether to compute the attributions
+    with respect to the inputs or outputs of the layer is defined by the
+    input flag `attribute_to_layer_input`.
+    More details about the algorithm can be found here:
+
+    http://papers.nips.cc/paper/7062-a-unified-approach-to-interpreting-model-predictions.pdf
+
+    Note that the explanation model:
+        1. Assumes that input features are independent of one another
+        2. Is linear, meaning that the explanations are modeled through
+            the additive composition of feature effects.
+    Although, it assumes a linear model for each explanation, the overall
+    model across multiple explanations can be complex and non-linear.
+    """
+
+    def __init__(
+        self,
+        model: Module,
+        layer: Module,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            model (nn.Module):  The reference to PyTorch model instance. Model cannot
+                        contain any in-place nonlinear submodules; these are not
+                        supported by the register_full_backward_hook PyTorch API
+                        starting from PyTorch v1.9.
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                        The size and dimensionality of the attributions
+                        corresponds to the size and dimensionality of the layer's
+                        input or output depending on whether we attribute to the
+                        inputs or outputs of the layer.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in
+                        then that type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of LayerDeepLiftShap, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores are being
+                        multiplied by
+                        layer activations for inputs - layer activations for baselines
+                        This flag applies only if `custom_attribution_func` is
+                        set to None.
+        """
+        LayerDeepLift.__init__(self, model, layer)
+        DeepLiftShap.__init__(self, model, multiply_by_inputs)
+
+    # Ignoring mypy error for inconsistent signature with DeepLiftShap
+    @typing.overload  # type: ignore
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[
+            Tensor, Tuple[Tensor, ...], Callable[..., Union[Tensor, Tuple[Tensor, ...]]]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[
+            Tensor, Tuple[Tensor, ...], Callable[..., Union[Tensor, Tuple[Tensor, ...]]]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[
+            Tensor, Tuple[Tensor, ...], Callable[..., Union[Tensor, Tuple[Tensor, ...]]]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples (aka batch size), and if
+                        multiple input tensors are provided, the examples must
+                        be aligned appropriately.
+            baselines (tensor, tuple of tensors, callable):
+                        Baselines define reference samples that are compared with
+                        the inputs. In order to assign attribution scores DeepLift
+                        computes the differences between the inputs/outputs and
+                        corresponding references. Baselines can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          the first dimension equal to the number of examples
+                          in the baselines' distribution. The remaining dimensions
+                          must match with input tensor's dimension starting from
+                          the second dimension.
+
+                        - a tuple of tensors, if inputs is a tuple of tensors,
+                          with the first dimension of any tensor inside the tuple
+                          equal to the number of examples in the baseline's
+                          distribution. The remaining dimensions must match
+                          the dimensions of the corresponding input tensor
+                          starting from the second dimension.
+
+                        - callable function, optionally takes `inputs` as an
+                          argument and either returns a single tensor
+                          or a tuple of those.
+
+                        It is recommended that the number of samples in the baselines'
+                        tensors is larger than one.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a tuple
+                        containing multiple additional arguments including tensors
+                        or any arbitrary python types. These arguments are provided to
+                        forward_func in order, following the arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attributions with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer inputs, otherwise it will be computed with respect
+                        to layer outputs.
+                        Note that currently it assumes that both the inputs and
+                        outputs of internal layers are single tensors.
+                        Support for multiple tensors will be added later.
+                        Default: False
+            custom_attribution_func (callable, optional): A custom function for
+                        computing final attribution scores. This function can take
+                        at least one and at most three arguments with the
+                        following signature:
+
+                        - custom_attribution_func(multipliers)
+                        - custom_attribution_func(multipliers, inputs)
+                        - custom_attribution_func(multipliers, inputs, baselines)
+
+                        In case this function is not provided, we use the default
+                        logic defined as: multipliers * (inputs - baselines)
+                        It is assumed that all input arguments, `multipliers`,
+                        `inputs` and `baselines` are provided in tuples of same
+                        length. `custom_attribution_func` returns a tuple of
+                        attribution tensors that have the same length as the
+                        `inputs`.
+                        Default: None
+
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attribution score computed based on DeepLift's rescale rule
+                        with respect to layer's inputs or outputs. Attributions
+                        will always be the same size as the provided layer's inputs
+                        or outputs, depending on whether we attribute to the inputs
+                        or outputs of the layer.
+                        Attributions are returned in a tuple based on whether
+                        the layer inputs / outputs are contained in a tuple
+                        from a forward hook. For standard modules, inputs of
+                        a single tensor are usually wrapped in a tuple, while
+                        outputs of a single tensor are not.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        This is computed using the property that the
+                        total sum of forward_func(inputs) - forward_func(baselines)
+                        must be very close to the total sum of attributions
+                        computed based on approximated SHAP values using
+                        DeepLift's rescale rule.
+                        Delta is calculated for each example input and baseline pair,
+                        meaning that the number of elements in returned delta tensor
+                        is equal to the
+                        `number of examples in input` * `number of examples
+                        in baseline`. The deltas are ordered in the first place by
+                        input example, followed by the baseline.
+                        Note that the logic described for deltas is guaranteed
+                        when the default logic for attribution computations is used,
+                        meaning that the `custom_attribution_func=None`, otherwise
+                        it is not guaranteed and depends on the specifics of the
+                        `custom_attribution_func`.
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> net = ImageClassifier()
+            >>> # creates an instance of LayerDeepLift to interpret target
+            >>> # class 1 with respect to conv4 layer.
+            >>> dl = LayerDeepLiftShap(net, net.conv4)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes shap values using deeplift for class 3.
+            >>> attribution = dl.attribute(input, target=3)
+        """
+        inputs = _format_tensor_into_tuples(inputs)
+        baselines = _format_callable_baseline(baselines, inputs)
+
+        assert isinstance(baselines[0], torch.Tensor) and baselines[0].shape[0] > 1, (
+            "Baselines distribution has to be provided in form of a torch.Tensor"
+            " with more than one example but found: {}."
+            " If baselines are provided in shape of scalars or with a single"
+            " baseline example, `LayerDeepLift`"
+            " approach can be used instead.".format(baselines[0])
+        )
+
+        # batch sizes
+        inp_bsz = inputs[0].shape[0]
+        base_bsz = baselines[0].shape[0]
+
+        (
+            exp_inp,
+            exp_base,
+            exp_target,
+            exp_addit_args,
+        ) = DeepLiftShap._expand_inputs_baselines_targets(
+            self, baselines, inputs, target, additional_forward_args
+        )
+        attributions = LayerDeepLift.attribute.__wrapped__(  # type: ignore
+            self,
+            exp_inp,
+            exp_base,
+            target=exp_target,
+            additional_forward_args=exp_addit_args,
+            return_convergence_delta=cast(
+                Literal[True, False], return_convergence_delta
+            ),
+            attribute_to_layer_input=attribute_to_layer_input,
+            custom_attribution_func=custom_attribution_func,
+        )
+        if return_convergence_delta:
+            attributions, delta = attributions
+        if isinstance(attributions, tuple):
+            attributions = tuple(
+                DeepLiftShap._compute_mean_across_baselines(
+                    self, inp_bsz, base_bsz, cast(Tensor, attrib)
+                )
+                for attrib in attributions
+            )
+        else:
+            attributions = DeepLiftShap._compute_mean_across_baselines(
+                self, inp_bsz, base_bsz, attributions
+            )
+        if return_convergence_delta:
+            return attributions, delta
+        else:
+            return attributions
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs

captum/attr/_core/layer/layer_feature_ablation.py

@@ -0,0 +1,302 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _extract_device,
+    _format_additional_forward_args,
+    _format_output,
+    _format_tensor_into_tuples,
+    _run_forward,
+)
+from captum._utils.gradient import _forward_layer_eval
+from captum._utils.typing import BaselineType, TargetType
+from captum.attr._core.feature_ablation import FeatureAblation
+from captum.attr._utils.attribution import LayerAttribution, PerturbationAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+from torch.nn.parallel.scatter_gather import scatter
+
+
+class LayerFeatureAblation(LayerAttribution, PerturbationAttribution):
+    r"""
+    A perturbation based approach to computing layer attribution, involving
+    replacing values in the input / output of a layer with a given baseline /
+    reference, and computing the difference in output. By default, each
+    neuron (scalar input / output value) within the layer is replaced
+    independently.
+    Passing a layer mask allows grouping neurons to be
+    ablated together.
+    Each neuron in the group will be given the same attribution value
+    equal to the change in target as a result of ablating the entire neuron
+    group.
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                          modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                          Output size of attribute matches this layer's input or
+                          output dimensions, depending on whether we attribute to
+                          the inputs or outputs of the layer, corresponding to
+                          attribution of each neuron in the input or output of
+                          this layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                          applies a DataParallel model. This allows reconstruction of
+                          intermediate outputs from batched results across devices.
+                          If forward_func is given as the DataParallel model itself
+                          (or otherwise has a device_ids attribute with the device
+                          ID list), then it is not necessary to provide this
+                          argument.
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        PerturbationAttribution.__init__(self, forward_func)
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        layer_baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        layer_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
+        attribute_to_layer_input: bool = False,
+        perturbations_per_eval: int = 1,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer
+                        attributions are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            layer_baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Layer baselines define reference values which replace each
+                        layer input / output value when ablated.
+                        Layer baselines should be a single tensor with dimensions
+                        matching the input / output of the target layer (or
+                        broadcastable to match it), based
+                        on whether we are attributing to the input or output
+                        of the target layer.
+                        In the cases when `baselines` is not provided, we internally
+                        use zero as the baseline for each neuron.
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            layer_mask (tensor or tuple of tensors, optional):
+                        layer_mask defines a mask for the layer, grouping
+                        elements of the layer input / output which should be
+                        ablated together.
+                        layer_mask should be a single tensor with dimensions
+                        matching the input / output of the target layer (or
+                        broadcastable to match it), based
+                        on whether we are attributing to the input or output
+                        of the target layer. layer_mask
+                        should contain integers in the range 0 to num_groups
+                        - 1, and all elements with the same value are
+                        considered to be in the same group.
+                        If None, then a layer mask is constructed which assigns
+                        each neuron within the layer as a separate group, which
+                        is ablated independently.
+                        Default: None
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attributions with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer's inputs, otherwise it will be computed with respect
+                        to layer's outputs.
+                        Note that currently it is assumed that either the input
+                        or the output of the layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+            perturbations_per_eval (int, optional): Allows ablation of multiple
+                        neuron (groups) to be processed simultaneously in one
+                        call to forward_fn.
+                        Each forward pass will contain a maximum of
+                        perturbations_per_eval * #examples samples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain at most
+                        (perturbations_per_eval * #examples) / num_devices
+                        samples.
+                        Default: 1
+
+        Returns:
+            *tensor* or tuple of *tensors* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attribution of each neuron in given layer input or
+                        output. Attributions will always be the same size as
+                        the input or output of the given layer, depending on
+                        whether we attribute to the inputs or outputs
+                        of the layer which is decided by the input flag
+                        `attribute_to_layer_input`
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+
+
+        Examples::
+
+        >>> # SimpleClassifier takes a single input tensor of size Nx4x4,
+        >>> # and returns an Nx3 tensor of class probabilities.
+        >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+        >>> # and the output of this layer has dimensions Nx12x3x3.
+        >>> net = SimpleClassifier()
+        >>> # Generating random input with size 2 x 4 x 4
+        >>> input = torch.randn(2, 4, 4)
+        >>> # Defining LayerFeatureAblation interpreter
+        >>> ablator = LayerFeatureAblation(net, net.conv1)
+        >>> # Computes ablation attribution, ablating each of the 108
+        >>> # neurons independently.
+        >>> attr = ablator.attribute(input, target=1)
+
+        >>> # Alternatively, we may want to ablate neurons in groups, e.g.
+        >>> # grouping all the layer outputs in the same row.
+        >>> # This can be done by creating a layer mask as follows, which
+        >>> # defines the groups of layer inputs / outouts, e.g.:
+        >>> # +---+---+---+
+        >>> # | 0 | 0 | 0 |
+        >>> # +---+---+---+
+        >>> # | 1 | 1 | 1 |
+        >>> # +---+---+---+
+        >>> # | 2 | 2 | 2 |
+        >>> # +---+---+---+
+        >>> # With this mask, all the 36 neurons in a row / channel are ablated
+        >>> # simultaneously, and the attribution for each neuron in the same
+        >>> # group (0 - 2) per example are the same.
+        >>> # The attributions can be calculated as follows:
+        >>> # layer mask has dimensions 1 x 3 x 3
+        >>> layer_mask = torch.tensor([[[0,0,0],[1,1,1],
+        >>>                             [2,2,2]]])
+        >>> attr = ablator.attribute(input, target=1,
+        >>>                          layer_mask=layer_mask)
+        """
+
+        def layer_forward_func(*args):
+            layer_length = args[-1]
+            layer_input = args[:layer_length]
+            original_inputs = args[layer_length:-1]
+
+            device_ids = self.device_ids
+            if device_ids is None:
+                device_ids = getattr(self.forward_func, "device_ids", None)
+
+            all_layer_inputs = {}
+            if device_ids is not None:
+                scattered_layer_input = scatter(layer_input, target_gpus=device_ids)
+                for device_tensors in scattered_layer_input:
+                    all_layer_inputs[device_tensors[0].device] = device_tensors
+            else:
+                all_layer_inputs[layer_input[0].device] = layer_input
+
+            def forward_hook(module, inp, out=None):
+                device = _extract_device(module, inp, out)
+                is_layer_tuple = (
+                    isinstance(out, tuple)
+                    if out is not None
+                    else isinstance(inp, tuple)
+                )
+                if device not in all_layer_inputs:
+                    raise AssertionError(
+                        "Layer input not placed on appropriate "
+                        "device. If using a DataParallel model, either provide the "
+                        "DataParallel model as forward_func or provide device ids"
+                        " to the constructor."
+                    )
+                if not is_layer_tuple:
+                    return all_layer_inputs[device][0]
+                return all_layer_inputs[device]
+
+            hook = None
+            try:
+                if attribute_to_layer_input:
+                    hook = self.layer.register_forward_pre_hook(forward_hook)
+                else:
+                    hook = self.layer.register_forward_hook(forward_hook)
+                eval = _run_forward(self.forward_func, original_inputs, target=target)
+            finally:
+                if hook is not None:
+                    hook.remove()
+            return eval
+
+        with torch.no_grad():
+            inputs = _format_tensor_into_tuples(inputs)
+            additional_forward_args = _format_additional_forward_args(
+                additional_forward_args
+            )
+            layer_eval = _forward_layer_eval(
+                self.forward_func,
+                inputs,
+                self.layer,
+                additional_forward_args,
+                device_ids=self.device_ids,
+                attribute_to_layer_input=attribute_to_layer_input,
+            )
+            layer_eval_len = (len(layer_eval),)
+            all_inputs = (
+                (inputs + additional_forward_args + layer_eval_len)
+                if additional_forward_args is not None
+                else inputs + layer_eval_len
+            )
+
+            ablator = FeatureAblation(layer_forward_func)
+
+            layer_attribs = ablator.attribute.__wrapped__(
+                ablator,  # self
+                layer_eval,
+                baselines=layer_baselines,
+                additional_forward_args=all_inputs,
+                feature_mask=layer_mask,
+                perturbations_per_eval=perturbations_per_eval,
+            )
+            _attr = _format_output(len(layer_attribs) > 1, layer_attribs)
+        return _attr

captum/attr/_core/layer/layer_gradient_shap.py

@@ -0,0 +1,474 @@
+#!/usr/bin/env python3
+
+import typing
+from typing import Any, Callable, cast, List, Tuple, Union
+
+import numpy as np
+import torch
+from captum._utils.gradient import _forward_layer_eval, compute_layer_gradients_and_eval
+from captum._utils.typing import Literal, TargetType, TensorOrTupleOfTensorsGeneric
+from captum.attr._core.gradient_shap import _scale_input
+from captum.attr._core.noise_tunnel import NoiseTunnel
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.attr._utils.common import (
+    _compute_conv_delta_and_format_attrs,
+    _format_callable_baseline,
+    _format_input_baseline,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerGradientShap(LayerAttribution, GradientAttribution):
+    r"""
+    Implements gradient SHAP for layer based on the implementation from SHAP's
+    primary author. For reference, please, view:
+
+    https://github.com/slundberg/shap\
+    #deep-learning-example-with-gradientexplainer-tensorflowkeraspytorch-models
+
+    A Unified Approach to Interpreting Model Predictions
+    http://papers.nips.cc/paper\
+    7062-a-unified-approach-to-interpreting-model-predictions
+
+    GradientShap approximates SHAP values by computing the expectations of
+    gradients by randomly sampling from the distribution of baselines/references.
+    It adds white noise to each input sample `n_samples` times, selects a
+    random baseline from baselines' distribution and a random point along the
+    path between the baseline and the input, and computes the gradient of
+    outputs with respect to selected random points in chosen `layer`.
+    The final SHAP values represent the expected values of
+    `gradients * (layer_attr_inputs - layer_attr_baselines)`.
+
+    GradientShap makes an assumption that the input features are independent
+    and that the explanation model is linear, meaning that the explanations
+    are modeled through the additive composition of feature effects.
+    Under those assumptions, SHAP value can be approximated as the expectation
+    of gradients that are computed for randomly generated `n_samples` input
+    samples after adding gaussian noise `n_samples` times to each input for
+    different baselines/references.
+
+    In some sense it can be viewed as an approximation of integrated gradients
+    by computing the expectations of gradients for different baselines.
+
+    Current implementation uses Smoothgrad from `NoiseTunnel` in order to
+    randomly draw samples from the distribution of baselines, add noise to input
+    samples and compute the expectation (smoothgrad).
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                        modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                        Output size of attribute matches this layer's input or
+                        output dimensions, depending on whether we attribute to
+                        the inputs or outputs of the layer, corresponding to
+                        attribution of each neuron in the input or output of
+                        this layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                        applies a DataParallel model. This allows reconstruction of
+                        intermediate outputs from batched results across devices.
+                        If forward_func is given as the DataParallel model itself,
+                        then it is not necessary to provide this argument.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in,
+                        then this type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of layer gradient shap, if `multiply_by_inputs`
+                        is set to True, the sensitivity scores for scaled inputs
+                        are being multiplied by
+                        layer activations for inputs - layer activations for baselines.
+
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[TensorOrTupleOfTensorsGeneric, Callable],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool = False,
+    ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[TensorOrTupleOfTensorsGeneric, Callable],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[TensorOrTupleOfTensorsGeneric, Callable],
+        n_samples: int = 5,
+        stdevs: Union[float, Tuple[float, ...]] = 0.0,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[
+        Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
+    ]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input which are used to compute
+                        SHAP attribution values for a given `layer`. If `forward_func`
+                        takes a single tensor as input, a single input tensor should
+                        be provided.
+                        If `forward_func` takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (tensor, tuple of tensors, callable):
+                        Baselines define the starting point from which expectation
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          the first dimension equal to the number of examples
+                          in the baselines' distribution. The remaining dimensions
+                          must match with input tensor's dimension starting from
+                          the second dimension.
+
+                        - a tuple of tensors, if inputs is a tuple of tensors,
+                          with the first dimension of any tensor inside the tuple
+                          equal to the number of examples in the baseline's
+                          distribution. The remaining dimensions must match
+                          the dimensions of the corresponding input tensor
+                          starting from the second dimension.
+
+                        - callable function, optionally takes `inputs` as an
+                          argument and either returns a single tensor
+                          or a tuple of those.
+
+                        It is recommended that the number of samples in the baselines'
+                        tensors is larger than one.
+            n_samples (int, optional):  The number of randomly generated examples
+                        per sample in the input batch. Random examples are
+                        generated by adding gaussian random noise to each sample.
+                        Default: `5` if `n_samples` is not provided.
+            stdevs    (float, or a tuple of floats optional): The standard deviation
+                        of gaussian noise with zero mean that is added to each
+                        input in the batch. If `stdevs` is a single float value
+                        then that same value is used for all inputs. If it is
+                        a tuple, then it must have the same length as the inputs
+                        tuple. In this case, each stdev value in the stdevs tuple
+                        corresponds to the input with the same index in the inputs
+                        tuple.
+                        Default: 0.0
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It can contain a tuple of ND tensors or
+                        any arbitrary python type of any shape.
+                        In case of the ND tensor the first dimension of the
+                        tensor must correspond to the batch size. It will be
+                        repeated for each `n_steps` for each randomly generated
+                        input sample.
+                        Note that the attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+        Returns:
+            **attributions** or 2-element tuple of **attributions**, **delta**:
+            - **attributions** (*tensor* or tuple of *tensors*):
+                        Attribution score computed based on GradientSHAP with
+                        respect to layer's input or output. Attributions will always
+                        be the same size as the provided layer's inputs or outputs,
+                        depending on whether we attribute to the inputs or outputs
+                        of the layer.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+            - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        This is computed using the property that the total
+                        sum of forward_func(inputs) - forward_func(baselines)
+                        must be very close to the total sum of the attributions
+                        based on layer gradient SHAP.
+                        Delta is calculated for each example in the input after adding
+                        `n_samples` times gaussian noise to each of them. Therefore,
+                        the dimensionality of the deltas tensor is equal to the
+                        `number of examples in the input` * `n_samples`
+                        The deltas are ordered by each input example and `n_samples`
+                        noisy samples generated for it.
+
+            Examples::
+
+                >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+                >>> # and returns an Nx10 tensor of class probabilities.
+                >>> net = ImageClassifier()
+                >>> layer_grad_shap = LayerGradientShap(net, net.linear1)
+                >>> input = torch.randn(3, 3, 32, 32, requires_grad=True)
+                >>> # choosing baselines randomly
+                >>> baselines = torch.randn(20, 3, 32, 32)
+                >>> # Computes gradient SHAP of output layer when target is equal
+                >>> # to 0 with respect to the layer linear1.
+                >>> # Attribution size matches to the size of the linear1 layer
+                >>> attribution = layer_grad_shap.attribute(input, baselines,
+                                                            target=5)
+
+        """
+        # since `baselines` is a distribution, we can generate it using a function
+        # rather than passing it as an input argument
+        baselines = _format_callable_baseline(baselines, inputs)
+        assert isinstance(baselines[0], torch.Tensor), (
+            "Baselines distribution has to be provided in a form "
+            "of a torch.Tensor {}.".format(baselines[0])
+        )
+
+        input_min_baseline_x_grad = LayerInputBaselineXGradient(
+            self.forward_func,
+            self.layer,
+            device_ids=self.device_ids,
+            multiply_by_inputs=self.multiplies_by_inputs,
+        )
+
+        nt = NoiseTunnel(input_min_baseline_x_grad)
+
+        attributions = nt.attribute.__wrapped__(
+            nt,  # self
+            inputs,
+            nt_type="smoothgrad",
+            nt_samples=n_samples,
+            stdevs=stdevs,
+            draw_baseline_from_distrib=True,
+            baselines=baselines,
+            target=target,
+            additional_forward_args=additional_forward_args,
+            return_convergence_delta=return_convergence_delta,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        return attributions
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+
+class LayerInputBaselineXGradient(LayerAttribution, GradientAttribution):
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: Module,
+        device_ids: Union[None, List[int]] = None,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                        modification of it
+            layer (torch.nn.Module): Layer for which attributions are computed.
+                        Output size of attribute matches this layer's input or
+                        output dimensions, depending on whether we attribute to
+                        the inputs or outputs of the layer, corresponding to
+                        attribution of each neuron in the input or output of
+                        this layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                        applies a DataParallel model. This allows reconstruction of
+                        intermediate outputs from batched results across devices.
+                        If forward_func is given as the DataParallel model itself,
+                        then it is not necessary to provide this argument.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in,
+                        then this type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of layer input minus baseline x gradient,
+                        if `multiply_by_inputs` is set to True, the sensitivity scores
+                        for scaled inputs are being multiplied by
+                        layer activations for inputs - layer activations for baselines.
+
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[Tensor, Tuple[Tensor, ...]],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: Literal[False] = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
+        ...
+
+    @typing.overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[Tensor, Tuple[Tensor, ...]],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        *,
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool = False,
+    ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]:
+        ...
+
+    @log_usage()
+    def attribute(  # type: ignore
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: Union[Tensor, Tuple[Tensor, ...]],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[
+        Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
+    ]:
+        inputs, baselines = _format_input_baseline(inputs, baselines)
+        rand_coefficient = torch.tensor(
+            np.random.uniform(0.0, 1.0, inputs[0].shape[0]),
+            device=inputs[0].device,
+            dtype=inputs[0].dtype,
+        )
+
+        input_baseline_scaled = tuple(
+            _scale_input(input, baseline, rand_coefficient)
+            for input, baseline in zip(inputs, baselines)
+        )
+        grads, _ = compute_layer_gradients_and_eval(
+            self.forward_func,
+            self.layer,
+            input_baseline_scaled,
+            target,
+            additional_forward_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        attr_baselines = _forward_layer_eval(
+            self.forward_func,
+            baselines,
+            self.layer,
+            additional_forward_args=additional_forward_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        attr_inputs = _forward_layer_eval(
+            self.forward_func,
+            inputs,
+            self.layer,
+            additional_forward_args=additional_forward_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        if self.multiplies_by_inputs:
+            input_baseline_diffs = tuple(
+                input - baseline for input, baseline in zip(attr_inputs, attr_baselines)
+            )
+            attributions = tuple(
+                input_baseline_diff * grad
+                for input_baseline_diff, grad in zip(input_baseline_diffs, grads)
+            )
+        else:
+            attributions = grads
+
+        return _compute_conv_delta_and_format_attrs(
+            self,
+            return_convergence_delta,
+            attributions,
+            baselines,
+            inputs,
+            additional_forward_args,
+            target,
+            cast(Union[Literal[True], Literal[False]], len(attributions) > 1),
+        )
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs

captum/attr/_core/layer/layer_gradient_x_activation.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+from typing import Any, Callable, List, Tuple, Union
+
+from captum._utils.common import (
+    _format_additional_forward_args,
+    _format_output,
+    _format_tensor_into_tuples,
+)
+from captum._utils.gradient import compute_layer_gradients_and_eval
+from captum._utils.typing import ModuleOrModuleList, TargetType
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn import Module
+
+
+class LayerGradientXActivation(LayerAttribution, GradientAttribution):
+    r"""
+    Computes element-wise product of gradient and activation for selected
+    layer on given inputs.
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: ModuleOrModuleList,
+        device_ids: Union[None, List[int]] = None,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+
+            forward_func (callable):  The forward function of the model or any
+                        modification of it
+            layer (torch.nn.Module or list(torch.nn.Module)): Layer or layers
+                          for which attributions are computed.
+                          Output size of attribute matches this layer's input or
+                          output dimensions, depending on whether we attribute to
+                          the inputs or outputs of the layer, corresponding to
+                          attribution of each neuron in the input or output of
+                          this layer. If multiple layers are provided, attributions
+                          are returned as a list, each element corresponding to the
+                          attributions of the corresponding layer.
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                        applies a DataParallel model. This allows reconstruction of
+                        intermediate outputs from batched results across devices.
+                        If forward_func is given as the DataParallel model itself,
+                        then it is not necessary to provide this argument.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in,
+                        then this type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of layer gradient x activation, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores are being multiplied by
+                        layer activations for inputs.
+
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids)
+        GradientAttribution.__init__(self, forward_func)
+        self._multiply_by_inputs = multiply_by_inputs
+
+    @property
+    def multiplies_by_inputs(self):
+        return self._multiply_by_inputs
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
+        r"""
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which attributions
+                        are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Default: False
+
+        Returns:
+            *tensor* or tuple of *tensors* or *list* of **attributions**:
+            - **attributions** (*tensor* or tuple of *tensors* or *list*):
+                        Product of gradient and activation for each
+                        neuron in given layer output.
+                        Attributions will always be the same size as the
+                        output of the given layer.
+                        Attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+                        If multiple layers are provided, attributions
+                        are returned as a list, each element corresponding to the
+                        activations of the corresponding layer.
+
+
+        Examples::
+
+            >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+            >>> # and returns an Nx10 tensor of class probabilities.
+            >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+            >>> # and the output of this layer has dimensions Nx12x32x32.
+            >>> net = ImageClassifier()
+            >>> layer_ga = LayerGradientXActivation(net, net.conv1)
+            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+            >>> # Computes layer activation x gradient for class 3.
+            >>> # attribution size matches layer output, Nx12x32x32
+            >>> attribution = layer_ga.attribute(input, 3)
+        """
+        inputs = _format_tensor_into_tuples(inputs)
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+        # Returns gradient of output with respect to
+        # hidden layer and hidden layer evaluated at each input.
+        layer_gradients, layer_evals = compute_layer_gradients_and_eval(
+            self.forward_func,
+            self.layer,
+            inputs,
+            target,
+            additional_forward_args,
+            device_ids=self.device_ids,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+        if isinstance(self.layer, Module):
+            return _format_output(
+                len(layer_evals) > 1,
+                self.multiply_gradient_acts(layer_gradients, layer_evals),
+            )
+        else:
+            return [
+                _format_output(
+                    len(layer_evals[i]) > 1,
+                    self.multiply_gradient_acts(layer_gradients[i], layer_evals[i]),
+                )
+                for i in range(len(self.layer))
+            ]
+
+    def multiply_gradient_acts(
+        self, gradients: Tuple[Tensor, ...], evals: Tuple[Tensor, ...]
+    ) -> Tuple[Tensor, ...]:
+        return tuple(
+            single_gradient * single_eval
+            if self.multiplies_by_inputs
+            else single_gradient
+            for single_gradient, single_eval in zip(gradients, evals)
+        )

captum/attr/_core/layer/layer_integrated_gradients.py

@@ -0,0 +1,528 @@
+#!/usr/bin/env python3
+import functools
+import warnings
+from typing import Any, Callable, List, overload, Tuple, Union
+
+import torch
+from captum._utils.common import (
+    _extract_device,
+    _format_additional_forward_args,
+    _format_outputs,
+)
+from captum._utils.gradient import _forward_layer_eval, _run_forward
+from captum._utils.typing import BaselineType, Literal, ModuleOrModuleList, TargetType
+from captum.attr._core.integrated_gradients import IntegratedGradients
+from captum.attr._utils.attribution import GradientAttribution, LayerAttribution
+from captum.attr._utils.common import (
+    _format_input_baseline,
+    _tensorize_baseline,
+    _validate_input,
+)
+from captum.log import log_usage
+from torch import Tensor
+from torch.nn.parallel.scatter_gather import scatter
+
+
+class LayerIntegratedGradients(LayerAttribution, GradientAttribution):
+    r"""
+    Layer Integrated Gradients is a variant of Integrated Gradients that assigns
+    an importance score to layer inputs or outputs, depending on whether we
+    attribute to the former or to the latter one.
+
+    Integrated Gradients is an axiomatic model interpretability algorithm that
+    attributes / assigns an importance score to each input feature by approximating
+    the integral of gradients of the model's output with respect to the inputs
+    along the path (straight line) from given baselines / references to inputs.
+
+    Baselines can be provided as input arguments to attribute method.
+    To approximate the integral we can choose to use either a variant of
+    Riemann sum or Gauss-Legendre quadrature rule.
+
+    More details regarding the integrated gradients method can be found in the
+    original paper:
+    https://arxiv.org/abs/1703.01365
+
+    """
+
+    def __init__(
+        self,
+        forward_func: Callable,
+        layer: ModuleOrModuleList,
+        device_ids: Union[None, List[int]] = None,
+        multiply_by_inputs: bool = True,
+    ) -> None:
+        r"""
+        Args:
+            forward_func (callable):  The forward function of the model or any
+                        modification of it
+            layer (ModuleOrModuleList):
+                        Layer or list of layers for which attributions are computed.
+                        For each layer the output size of the attribute matches
+                        this layer's input or output dimensions, depending on
+                        whether we attribute to the inputs or outputs of the
+                        layer, corresponding to the attribution of each neuron
+                        in the input or output of this layer.
+
+                        Please note that layers to attribute on cannot be
+                        dependent on each other. That is, a subset of layers in
+                        `layer` cannot produce the inputs for another layer.
+
+                        For example, if your model is of a simple linked-list
+                        based graph structure (think nn.Sequence), e.g. x -> l1
+                        -> l2 -> l3 -> output. If you pass in any one of those
+                        layers, you cannot pass in another due to the
+                        dependence, e.g.  if you pass in l2 you cannot pass in
+                        l1 or l3.
+
+            device_ids (list(int)): Device ID list, necessary only if forward_func
+                        applies a DataParallel model. This allows reconstruction of
+                        intermediate outputs from batched results across devices.
+                        If forward_func is given as the DataParallel model itself,
+                        then it is not necessary to provide this argument.
+            multiply_by_inputs (bool, optional): Indicates whether to factor
+                        model inputs' multiplier in the final attribution scores.
+                        In the literature this is also known as local vs global
+                        attribution. If inputs' multiplier isn't factored in,
+                        then this type of attribution method is also called local
+                        attribution. If it is, then that type of attribution
+                        method is called global.
+                        More detailed can be found here:
+                        https://arxiv.org/abs/1711.06104
+
+                        In case of layer integrated gradients, if `multiply_by_inputs`
+                        is set to True, final sensitivity scores are being multiplied by
+                        layer activations for inputs - layer activations for baselines.
+
+        """
+        LayerAttribution.__init__(self, forward_func, layer, device_ids=device_ids)
+        GradientAttribution.__init__(self, forward_func)
+        self.ig = IntegratedGradients(forward_func, multiply_by_inputs)
+
+        if isinstance(layer, list) and len(layer) > 1:
+            warnings.warn(
+                "Multiple layers provided. Please ensure that each layer is"
+                "**not** solely solely dependent on the outputs of"
+                "another layer. Please refer to the documentation for more"
+                "detail."
+            )
+
+    @overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType,
+        target: TargetType,
+        additional_forward_args: Any,
+        n_steps: int,
+        method: str,
+        internal_batch_size: Union[None, int],
+        return_convergence_delta: Literal[False],
+        attribute_to_layer_input: bool,
+    ) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
+        ...
+
+    @overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType,
+        target: TargetType,
+        additional_forward_args: Any,
+        n_steps: int,
+        method: str,
+        internal_batch_size: Union[None, int],
+        return_convergence_delta: Literal[True],
+        attribute_to_layer_input: bool,
+    ) -> Tuple[
+        Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]],
+        Tensor,
+    ]:
+        ...
+
+    @overload
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[
+        Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]],
+        Tuple[
+            Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]],
+            Tensor,
+        ],
+    ]:
+        ...
+
+    @log_usage()
+    def attribute(
+        self,
+        inputs: Union[Tensor, Tuple[Tensor, ...]],
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        n_steps: int = 50,
+        method: str = "gausslegendre",
+        internal_batch_size: Union[None, int] = None,
+        return_convergence_delta: bool = False,
+        attribute_to_layer_input: bool = False,
+    ) -> Union[
+        Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]],
+        Tuple[
+            Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]],
+            Tensor,
+        ],
+    ]:
+        r"""
+        This method attributes the output of the model with given target index
+        (in case it is provided, otherwise it assumes that output is a
+        scalar) to layer inputs or outputs of the model, depending on whether
+        `attribute_to_layer_input` is set to True or False, using the approach
+        described above.
+
+        In addition to that it also returns, if `return_convergence_delta` is
+        set to True, integral approximation delta based on the completeness
+        property of integrated gradients.
+
+        Args:
+
+            inputs (tensor or tuple of tensors):  Input for which layer integrated
+                        gradients are computed. If forward_func takes a single
+                        tensor as input, a single input tensor should be provided.
+                        If forward_func takes multiple tensors as input, a tuple
+                        of the input tensors should be provided. It is assumed
+                        that for all given input tensors, dimension 0 corresponds
+                        to the number of examples, and if multiple input tensors
+                        are provided, the examples must be aligned appropriately.
+            baselines (scalar, tensor, tuple of scalars or tensors, optional):
+                        Baselines define the starting point from which integral
+                        is computed and can be provided as:
+
+                        - a single tensor, if inputs is a single tensor, with
+                          exactly the same dimensions as inputs or the first
+                          dimension is one and the remaining dimensions match
+                          with inputs.
+
+                        - a single scalar, if inputs is a single tensor, which will
+                          be broadcasted for each input value in input tensor.
+
+                        - a tuple of tensors or scalars, the baseline corresponding
+                          to each tensor in the inputs' tuple can be:
+                            - either a tensor with matching dimensions to
+                              corresponding tensor in the inputs' tuple
+                              or the first dimension is one and the remaining
+                              dimensions match with the corresponding
+                              input tensor.
+                            - or a scalar, corresponding to a tensor in the
+                              inputs' tuple. This scalar value is broadcasted
+                              for corresponding input tensor.
+
+                        In the cases when `baselines` is not provided, we internally
+                        use zero scalar corresponding to each input tensor.
+
+                        Default: None
+            target (int, tuple, tensor or list, optional):  Output indices for
+                        which gradients are computed (for classification cases,
+                        this is usually the target class).
+                        If the network returns a scalar value per example,
+                        no target index is necessary.
+                        For general 2D outputs, targets can be either:
+
+                        - a single integer or a tensor containing a single
+                          integer, which is applied to all input examples
+
+                        - a list of integers or a 1D tensor, with length matching
+                          the number of examples in inputs (dim 0). Each integer
+                          is applied as the target for the corresponding example.
+
+                        For outputs with > 2 dimensions, targets can be either:
+
+                        - A single tuple, which contains #output_dims - 1
+                          elements. This target index is applied to all examples.
+
+                        - A list of tuples with length equal to the number of
+                          examples in inputs (dim 0), and each tuple containing
+                          #output_dims - 1 elements. Each tuple is applied as the
+                          target for the corresponding example.
+
+                        Default: None
+            additional_forward_args (any, optional): If the forward function
+                        requires additional arguments other than the inputs for
+                        which attributions should not be computed, this argument
+                        can be provided. It must be either a single additional
+                        argument of a Tensor or arbitrary (non-tuple) type or a
+                        tuple containing multiple additional arguments including
+                        tensors or any arbitrary python types. These arguments
+                        are provided to forward_func in order following the
+                        arguments in inputs.
+                        For a tensor, the first dimension of the tensor must
+                        correspond to the number of examples. It will be
+                        repeated for each of `n_steps` along the integrated
+                        path. For all other types, the given argument is used
+                        for all forward evaluations.
+                        Note that attributions are not computed with respect
+                        to these arguments.
+                        Default: None
+            n_steps (int, optional): The number of steps used by the approximation
+                        method. Default: 50.
+            method (string, optional): Method for approximating the integral,
+                        one of `riemann_right`, `riemann_left`, `riemann_middle`,
+                        `riemann_trapezoid` or `gausslegendre`.
+                        Default: `gausslegendre` if no method is provided.
+            internal_batch_size (int, optional): Divides total #steps * #examples
+                        data points into chunks of size at most internal_batch_size,
+                        which are computed (forward / backward passes)
+                        sequentially. internal_batch_size must be at least equal to
+                        #examples.
+                        For DataParallel models, each batch is split among the
+                        available devices, so evaluations on each available
+                        device contain internal_batch_size / num_devices examples.
+                        If internal_batch_size is None, then all evaluations are
+                        processed in one batch.
+                        Default: None
+            return_convergence_delta (bool, optional): Indicates whether to return
+                        convergence delta or not. If `return_convergence_delta`
+                        is set to True convergence delta will be returned in
+                        a tuple following attributions.
+                        Default: False
+            attribute_to_layer_input (bool, optional): Indicates whether to
+                        compute the attribution with respect to the layer input
+                        or output. If `attribute_to_layer_input` is set to True
+                        then the attributions will be computed with respect to
+                        layer input, otherwise it will be computed with respect
+                        to layer output.
+                        Note that currently it is assumed that either the input
+                        or the output of internal layer, depending on whether we
+                        attribute to the input or output, is a single tensor.
+                        Support for multiple tensors will be added later.
+                        Default: False
+            Returns:
+                **attributions** or 2-element tuple of **attributions**, **delta**:
+                - **attributions** (*tensor*, tuple of *tensors* or tuple of *tensors*):
+                        Integrated gradients with respect to `layer`'s inputs or
+                        outputs. Attributions will always be the same size and
+                        dimensionality as the input or output of the given layer,
+                        depending on whether we attribute to the inputs or outputs
+                        of the layer which is decided by the input flag
+                        `attribute_to_layer_input`.
+
+                        For a single layer, attributions are returned in a tuple if
+                        the layer inputs / outputs contain multiple tensors,
+                        otherwise a single tensor is returned.
+
+                        For multiple layers, attributions will always be
+                        returned as a list. Each element in this list will be
+                        equivalent to that of a single layer output, i.e. in the
+                        case that one layer, in the given layers, inputs / outputs
+                        multiple tensors: the corresponding output element will be
+                        a tuple of tensors. The ordering of the outputs will be
+                        the same order as the layers given in the constructor.
+                - **delta** (*tensor*, returned if return_convergence_delta=True):
+                        The difference between the total approximated and true
+                        integrated gradients. This is computed using the property
+                        that the total sum of forward_func(inputs) -
+                        forward_func(baselines) must equal the total sum of the
+                        integrated gradient.
+                        Delta is calculated per example, meaning that the number of
+                        elements in returned delta tensor is equal to the number of
+                        of examples in inputs.
+
+            Examples::
+
+                >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
+                >>> # and returns an Nx10 tensor of class probabilities.
+                >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
+                >>> # and the output of this layer has dimensions Nx12x32x32.
+                >>> net = ImageClassifier()
+                >>> lig = LayerIntegratedGradients(net, net.conv1)
+                >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
+                >>> # Computes layer integrated gradients for class 3.
+                >>> # attribution size matches layer output, Nx12x32x32
+                >>> attribution = lig.attribute(input, target=3)
+        """
+        inps, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inps, baselines, n_steps, method)
+
+        baselines = _tensorize_baseline(inps, baselines)
+        additional_forward_args = _format_additional_forward_args(
+            additional_forward_args
+        )
+
+        def flatten_tuple(tup):
+            return tuple(
+                sum((list(x) if isinstance(x, (tuple, list)) else [x] for x in tup), [])
+            )
+
+        if self.device_ids is None:
+            self.device_ids = getattr(self.forward_func, "device_ids", None)
+
+        inputs_layer = _forward_layer_eval(
+            self.forward_func,
+            inps,
+            self.layer,
+            device_ids=self.device_ids,
+            additional_forward_args=additional_forward_args,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+
+        # if we have one output
+        if not isinstance(self.layer, list):
+            inputs_layer = (inputs_layer,)
+
+        num_outputs = [1 if isinstance(x, Tensor) else len(x) for x in inputs_layer]
+        num_outputs_cumsum = torch.cumsum(
+            torch.IntTensor([0] + num_outputs), dim=0  # type: ignore
+        )
+        inputs_layer = flatten_tuple(inputs_layer)
+
+        baselines_layer = _forward_layer_eval(
+            self.forward_func,
+            baselines,
+            self.layer,
+            device_ids=self.device_ids,
+            additional_forward_args=additional_forward_args,
+            attribute_to_layer_input=attribute_to_layer_input,
+        )
+        baselines_layer = flatten_tuple(baselines_layer)
+
+        # inputs -> these inputs are scaled
+        def gradient_func(
+            forward_fn: Callable,
+            inputs: Union[Tensor, Tuple[Tensor, ...]],
+            target_ind: TargetType = None,
+            additional_forward_args: Any = None,
+        ) -> Tuple[Tensor, ...]:
+            if self.device_ids is None or len(self.device_ids) == 0:
+                scattered_inputs = (inputs,)
+            else:
+                # scatter method does not have a precise enough return type in its
+                # stub, so suppress the type warning.
+                scattered_inputs = scatter(  # type:ignore
+                    inputs, target_gpus=self.device_ids
+                )
+
+            scattered_inputs_dict = {
+                scattered_input[0].device: scattered_input
+                for scattered_input in scattered_inputs
+            }
+
+            with torch.autograd.set_grad_enabled(True):
+
+                def layer_forward_hook(
+                    module, hook_inputs, hook_outputs=None, layer_idx=0
+                ):
+                    device = _extract_device(module, hook_inputs, hook_outputs)
+                    is_layer_tuple = (
+                        isinstance(hook_outputs, tuple)
+                        # hook_outputs is None if attribute_to_layer_input == True
+                        if hook_outputs is not None
+                        else isinstance(hook_inputs, tuple)
+                    )
+
+                    if is_layer_tuple:
+                        return scattered_inputs_dict[device][
+                            num_outputs_cumsum[layer_idx] : num_outputs_cumsum[
+                                layer_idx + 1
+                            ]
+                        ]
+
+                    return scattered_inputs_dict[device][num_outputs_cumsum[layer_idx]]
+
+                hooks = []
+                try:
+
+                    layers = self.layer
+                    if not isinstance(layers, list):
+                        layers = [self.layer]
+
+                    for layer_idx, layer in enumerate(layers):
+                        hook = None
+                        # TODO:
+                        # Allow multiple attribute_to_layer_input flags for
+                        # each layer, i.e. attribute_to_layer_input[layer_idx]
+                        if attribute_to_layer_input:
+                            hook = layer.register_forward_pre_hook(
+                                functools.partial(
+                                    layer_forward_hook, layer_idx=layer_idx
+                                )
+                            )
+                        else:
+                            hook = layer.register_forward_hook(
+                                functools.partial(
+                                    layer_forward_hook, layer_idx=layer_idx
+                                )
+                            )
+
+                        hooks.append(hook)
+
+                    output = _run_forward(
+                        self.forward_func, tuple(), target_ind, additional_forward_args
+                    )
+                finally:
+                    for hook in hooks:
+                        if hook is not None:
+                            hook.remove()
+
+                assert output[0].numel() == 1, (
+                    "Target not provided when necessary, cannot"
+                    " take gradient with respect to multiple outputs."
+                )
+                # torch.unbind(forward_out) is a list of scalar tensor tuples and
+                # contains batch_size * #steps elements
+                grads = torch.autograd.grad(torch.unbind(output), inputs)
+            return grads
+
+        self.ig.gradient_func = gradient_func
+        all_inputs = (
+            (inps + additional_forward_args)
+            if additional_forward_args is not None
+            else inps
+        )
+
+        attributions = self.ig.attribute.__wrapped__(  # type: ignore
+            self.ig,  # self
+            inputs_layer,
+            baselines=baselines_layer,
+            target=target,
+            additional_forward_args=all_inputs,
+            n_steps=n_steps,
+            method=method,
+            internal_batch_size=internal_batch_size,
+            return_convergence_delta=False,
+        )
+
+        # handle multiple outputs
+        output: List[Tuple[Tensor, ...]] = [
+            tuple(
+                attributions[
+                    int(num_outputs_cumsum[i]) : int(num_outputs_cumsum[i + 1])
+                ]
+            )
+            for i in range(len(num_outputs))
+        ]
+
+        if return_convergence_delta:
+            start_point, end_point = baselines, inps
+            # computes approximation error based on the completeness axiom
+            delta = self.compute_convergence_delta(
+                attributions,
+                start_point,
+                end_point,
+                additional_forward_args=additional_forward_args,
+                target=target,
+            )
+            return _format_outputs(isinstance(self.layer, list), output), delta
+        return _format_outputs(isinstance(self.layer, list), output)
+
+    def has_convergence_delta(self) -> bool:
+        return True
+
+    @property
+    def multiplies_by_inputs(self):
+        return self.ig.multiplies_by_inputs