v1

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+*.params filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+*/__pycache__/*
+*/__pycache__/
+.ipynb_checkpoints
+__pycache__
+.venv/

AdaptiveWingLoss/.gitignore

@@ -0,0 +1,8 @@
+# Python generated files
+*.pyc
+
+# Project related files
+ckpt/*.pth
+dataset/*
+!dataset/!.py
+experiments/*

AdaptiveWingLoss/core/coord_conv.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn as nn
+
+
+class AddCoordsTh(nn.Module):
+    def __init__(self, x_dim=64, y_dim=64, with_r=False, with_boundary=False):
+        super(AddCoordsTh, self).__init__()
+        self.x_dim = x_dim
+        self.y_dim = y_dim
+        self.with_r = with_r
+        self.with_boundary = with_boundary
+
+    def forward(self, input_tensor, heatmap=None):
+        """
+        input_tensor: (batch, c, x_dim, y_dim)
+        """
+        batch_size_tensor = input_tensor.shape[0]
+
+        xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32).cuda()
+        xx_ones = xx_ones.unsqueeze(-1)
+
+        xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0).cuda()
+        xx_range = xx_range.unsqueeze(1)
+
+        xx_channel = torch.matmul(xx_ones.float(), xx_range.float())
+        xx_channel = xx_channel.unsqueeze(-1)
+
+
+        yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32).cuda()
+        yy_ones = yy_ones.unsqueeze(1)
+
+        yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0).cuda()
+        yy_range = yy_range.unsqueeze(-1)
+
+        yy_channel = torch.matmul(yy_range.float(), yy_ones.float())
+        yy_channel = yy_channel.unsqueeze(-1)
+
+        xx_channel = xx_channel.permute(0, 3, 2, 1)
+        yy_channel = yy_channel.permute(0, 3, 2, 1)
+
+        xx_channel = xx_channel / (self.x_dim - 1)
+        yy_channel = yy_channel / (self.y_dim - 1)
+
+        xx_channel = xx_channel * 2 - 1
+        yy_channel = yy_channel * 2 - 1
+
+        xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1)
+        yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1)
+
+        if self.with_boundary and type(heatmap) != type(None):
+            boundary_channel = torch.clamp(heatmap[:, -1:, :, :],
+                                        0.0, 1.0)
+
+            zero_tensor = torch.zeros_like(xx_channel)
+            xx_boundary_channel = torch.where(boundary_channel>0.05,
+                                              xx_channel, zero_tensor)
+            yy_boundary_channel = torch.where(boundary_channel>0.05,
+                                              yy_channel, zero_tensor)
+        if self.with_boundary and type(heatmap) != type(None):
+            xx_boundary_channel = xx_boundary_channel.cuda()
+            yy_boundary_channel = yy_boundary_channel.cuda()
+        ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1)
+
+
+        if self.with_r:
+            rr = torch.sqrt(torch.pow(xx_channel, 2) + torch.pow(yy_channel, 2))
+            rr = rr / torch.max(rr)
+            ret = torch.cat([ret, rr], dim=1)
+
+        if self.with_boundary and type(heatmap) != type(None):
+            ret = torch.cat([ret, xx_boundary_channel,
+                             yy_boundary_channel], dim=1)
+        return ret
+
+
+class CoordConvTh(nn.Module):
+    """CoordConv layer as in the paper."""
+    def __init__(self, x_dim, y_dim, with_r, with_boundary,
+                 in_channels, first_one=False, *args, **kwargs):
+        super(CoordConvTh, self).__init__()
+        self.addcoords = AddCoordsTh(x_dim=x_dim, y_dim=y_dim, with_r=with_r,
+                                    with_boundary=with_boundary)
+        in_channels += 2
+        if with_r:
+            in_channels += 1
+        if with_boundary and not first_one:
+            in_channels += 2
+        self.conv = nn.Conv2d(in_channels=in_channels, *args, **kwargs)
+
+    def forward(self, input_tensor, heatmap=None):
+        ret = self.addcoords(input_tensor, heatmap)
+        last_channel = ret[:, -2:, :, :]
+        ret = self.conv(ret)
+        return ret, last_channel
+
+
+'''
+An alternative implementation for PyTorch with auto-infering the x-y dimensions.
+'''
+class AddCoords(nn.Module):
+
+    def __init__(self, with_r=False):
+        super().__init__()
+        self.with_r = with_r
+
+    def forward(self, input_tensor):
+        """
+        Args:
+            input_tensor: shape(batch, channel, x_dim, y_dim)
+        """
+        batch_size, _, x_dim, y_dim = input_tensor.size()
+
+        xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1)
+        yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2)
+
+        xx_channel = xx_channel / (x_dim - 1)
+        yy_channel = yy_channel / (y_dim - 1)
+
+        xx_channel = xx_channel * 2 - 1
+        yy_channel = yy_channel * 2 - 1
+
+        xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+        yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+
+        if input_tensor.is_cuda:
+            xx_channel = xx_channel.cuda()
+            yy_channel = yy_channel.cuda()
+
+        ret = torch.cat([
+            input_tensor,
+            xx_channel.type_as(input_tensor),
+            yy_channel.type_as(input_tensor)], dim=1)
+
+        if self.with_r:
+            rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
+            if input_tensor.is_cuda:
+                rr = rr.cuda()
+            ret = torch.cat([ret, rr], dim=1)
+
+        return ret
+
+
+class CoordConv(nn.Module):
+
+    def __init__(self, in_channels, out_channels, with_r=False, **kwargs):
+        super().__init__()
+        self.addcoords = AddCoords(with_r=with_r)
+        self.conv = nn.Conv2d(in_channels + 2, out_channels, **kwargs)
+
+    def forward(self, x):
+        ret = self.addcoords(x)
+        ret = self.conv(ret)
+        return ret

AdaptiveWingLoss/core/dataloader.py

@@ -0,0 +1,368 @@
+import sys
+import os
+import random
+import glob
+import torch
+from skimage import io
+from skimage import transform as ski_transform
+from skimage.color import rgb2gray
+import scipy.io as sio
+from scipy import interpolate
+import numpy as np
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms.functional import adjust_brightness, adjust_contrast, adjust_saturation, adjust_hue
+from utils.utils import cv_crop, cv_rotate, draw_gaussian, transform, power_transform, shuffle_lr, fig2data, generate_weight_map
+from PIL import Image
+import cv2
+import copy
+import math
+from imgaug import augmenters as iaa
+
+
+class AddBoundary(object):
+    def __init__(self, num_landmarks=68):
+        self.num_landmarks = num_landmarks
+
+    def __call__(self, sample):
+        landmarks_64 = np.floor(sample['landmarks'] / 4.0)
+        if self.num_landmarks == 68:
+            boundaries = {}
+            boundaries['cheek'] = landmarks_64[0:17]
+            boundaries['left_eyebrow'] = landmarks_64[17:22]
+            boundaries['right_eyebrow'] = landmarks_64[22:27]
+            boundaries['uper_left_eyelid'] = landmarks_64[36:40]
+            boundaries['lower_left_eyelid'] = np.array([landmarks_64[i] for i in [36, 41, 40, 39]])
+            boundaries['upper_right_eyelid'] = landmarks_64[42:46]
+            boundaries['lower_right_eyelid'] = np.array([landmarks_64[i] for i in [42, 47, 46, 45]])
+            boundaries['noise'] = landmarks_64[27:31]
+            boundaries['noise_bot'] = landmarks_64[31:36]
+            boundaries['upper_outer_lip'] = landmarks_64[48:55]
+            boundaries['upper_inner_lip'] = np.array([landmarks_64[i] for i in [60, 61, 62, 63, 64]])
+            boundaries['lower_outer_lip'] = np.array([landmarks_64[i] for i in [48, 59, 58, 57, 56, 55, 54]])
+            boundaries['lower_inner_lip'] = np.array([landmarks_64[i] for i in [60, 67, 66, 65, 64]])
+        elif self.num_landmarks == 98:
+            boundaries = {}
+            boundaries['cheek'] = landmarks_64[0:33]
+            boundaries['left_eyebrow'] = landmarks_64[33:38]
+            boundaries['right_eyebrow'] = landmarks_64[42:47]
+            boundaries['uper_left_eyelid'] = landmarks_64[60:65]
+            boundaries['lower_left_eyelid'] = np.array([landmarks_64[i] for i in [60, 67, 66, 65, 64]])
+            boundaries['upper_right_eyelid'] = landmarks_64[68:73]
+            boundaries['lower_right_eyelid'] = np.array([landmarks_64[i] for i in [68, 75, 74, 73, 72]])
+            boundaries['noise'] = landmarks_64[51:55]
+            boundaries['noise_bot'] = landmarks_64[55:60]
+            boundaries['upper_outer_lip'] = landmarks_64[76:83]
+            boundaries['upper_inner_lip'] = np.array([landmarks_64[i] for i in [88, 89, 90, 91, 92]])
+            boundaries['lower_outer_lip'] = np.array([landmarks_64[i] for i in [76, 87, 86, 85, 84, 83, 82]])
+            boundaries['lower_inner_lip'] = np.array([landmarks_64[i] for i in [88, 95, 94, 93, 92]])
+        elif self.num_landmarks == 19:
+            boundaries = {}
+            boundaries['left_eyebrow'] = landmarks_64[0:3]
+            boundaries['right_eyebrow'] = landmarks_64[3:5]
+            boundaries['left_eye'] = landmarks_64[6:9]
+            boundaries['right_eye'] = landmarks_64[9:12]
+            boundaries['noise'] = landmarks_64[12:15]
+
+        elif self.num_landmarks == 29:
+            boundaries = {}
+            boundaries['upper_left_eyebrow'] = np.stack([
+                landmarks_64[0],
+                landmarks_64[4],
+                landmarks_64[2]
+            ], axis=0)
+            boundaries['lower_left_eyebrow'] = np.stack([
+                landmarks_64[0],
+                landmarks_64[5],
+                landmarks_64[2]
+            ], axis=0)
+            boundaries['upper_right_eyebrow'] = np.stack([
+                landmarks_64[1],
+                landmarks_64[6],
+                landmarks_64[3]
+            ], axis=0)
+            boundaries['lower_right_eyebrow'] = np.stack([
+                landmarks_64[1],
+                landmarks_64[7],
+                landmarks_64[3]
+            ], axis=0)
+            boundaries['upper_left_eye'] = np.stack([
+                landmarks_64[8],
+                landmarks_64[12],
+                landmarks_64[10]
+            ], axis=0)
+            boundaries['lower_left_eye'] = np.stack([
+                landmarks_64[8],
+                landmarks_64[13],
+                landmarks_64[10]
+            ], axis=0)
+            boundaries['upper_right_eye'] = np.stack([
+                landmarks_64[9],
+                landmarks_64[14],
+                landmarks_64[11]
+            ], axis=0)
+            boundaries['lower_right_eye'] = np.stack([
+                landmarks_64[9],
+                landmarks_64[15],
+                landmarks_64[11]
+            ], axis=0)
+            boundaries['noise'] = np.stack([
+                landmarks_64[18],
+                landmarks_64[21],
+                landmarks_64[19]
+            ], axis=0)
+            boundaries['outer_upper_lip'] = np.stack([
+                landmarks_64[22],
+                landmarks_64[24],
+                landmarks_64[23]
+            ], axis=0)
+            boundaries['inner_upper_lip'] = np.stack([
+                landmarks_64[22],
+                landmarks_64[25],
+                landmarks_64[23]
+            ], axis=0)
+            boundaries['outer_lower_lip'] = np.stack([
+                landmarks_64[22],
+                landmarks_64[26],
+                landmarks_64[23]
+            ], axis=0)
+            boundaries['inner_lower_lip'] = np.stack([
+                landmarks_64[22],
+                landmarks_64[27],
+                landmarks_64[23]
+            ], axis=0)
+        functions = {}
+
+        for key, points in boundaries.items():
+            temp = points[0]
+            new_points = points[0:1, :]
+            for point in points[1:]:
+                if point[0] == temp[0] and point[1] == temp[1]:
+                    continue
+                else:
+                    new_points = np.concatenate((new_points, np.expand_dims(point, 0)), axis=0)
+                    temp = point
+            points = new_points
+            if points.shape[0] == 1:
+                points = np.concatenate((points, points+0.001), axis=0)
+            k = min(4, points.shape[0])
+            functions[key] = interpolate.splprep([points[:, 0], points[:, 1]], k=k-1,s=0)
+
+        boundary_map = np.zeros((64, 64))
+
+        fig = plt.figure(figsize=[64/96.0, 64/96.0], dpi=96)
+
+        ax = fig.add_axes([0, 0, 1, 1])
+
+        ax.axis('off')
+
+        ax.imshow(boundary_map, interpolation='nearest', cmap='gray')
+        #ax.scatter(landmarks[:, 0], landmarks[:, 1], s=1, marker=',', c='w')
+
+        for key in functions.keys():
+            xnew = np.arange(0, 1, 0.01)
+            out = interpolate.splev(xnew, functions[key][0], der=0)
+            plt.plot(out[0], out[1], ',', linewidth=1, color='w')
+
+        img = fig2data(fig)
+
+        plt.close()
+
+        sigma = 1
+        temp = 255-img[:,:,1]
+        temp = cv2.distanceTransform(temp, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
+        temp = temp.astype(np.float32)
+        temp = np.where(temp < 3*sigma, np.exp(-(temp*temp)/(2*sigma*sigma)), 0 )
+
+        fig = plt.figure(figsize=[64/96.0, 64/96.0], dpi=96)
+
+        ax = fig.add_axes([0, 0, 1, 1])
+
+        ax.axis('off')
+        ax.imshow(temp, cmap='gray')
+        plt.close()
+
+        boundary_map = fig2data(fig)
+
+        sample['boundary'] = boundary_map[:, :, 0]
+
+        return sample
+
+class AddWeightMap(object):
+    def __call__(self, sample):
+        heatmap= sample['heatmap']
+        boundary = sample['boundary']
+        heatmap = np.concatenate((heatmap, np.expand_dims(boundary, axis=0)), 0)
+        weight_map = np.zeros_like(heatmap)
+        for i in range(heatmap.shape[0]):
+            weight_map[i] = generate_weight_map(weight_map[i],
+                                                heatmap[i])
+        sample['weight_map'] = weight_map
+        return sample
+
+class ToTensor(object):
+    """Convert ndarrays in sample to Tensors."""
+
+    def __call__(self, sample):
+        image, heatmap, landmarks, boundary, weight_map= sample['image'], sample['heatmap'], sample['landmarks'], sample['boundary'], sample['weight_map']
+
+        # swap color axis because
+        # numpy image: H x W x C
+        # torch image: C X H X W
+        if len(image.shape) == 2:
+            image = np.expand_dims(image, axis=2)
+            image_small = np.expand_dims(image_small, axis=2)
+        image = image.transpose((2, 0, 1))
+        boundary = np.expand_dims(boundary, axis=2)
+        boundary = boundary.transpose((2, 0, 1))
+        return {'image': torch.from_numpy(image).float().div(255.0),
+                'heatmap': torch.from_numpy(heatmap).float(),
+                'landmarks': torch.from_numpy(landmarks).float(),
+                'boundary': torch.from_numpy(boundary).float().div(255.0),
+                'weight_map': torch.from_numpy(weight_map).float()}
+
+class FaceLandmarksDataset(Dataset):
+    """Face Landmarks dataset."""
+
+    def __init__(self, img_dir, landmarks_dir, num_landmarks=68, gray_scale=False,
+                 detect_face=False, enhance=False, center_shift=0,
+                 transform=None,):
+        """
+        Args:
+            landmark_dir (string): Path to the mat file with landmarks saved.
+            img_dir (string): Directory with all the images.
+            transform (callable, optional): Optional transform to be applied
+                on a sample.
+        """
+        self.img_dir = img_dir
+        self.landmarks_dir = landmarks_dir
+        self.num_lanmdkars = num_landmarks
+        self.transform = transform
+        self.img_names = glob.glob(self.img_dir+'*.jpg') + \
+                         glob.glob(self.img_dir+'*.png')
+        self.gray_scale = gray_scale
+        self.detect_face = detect_face
+        self.enhance = enhance
+        self.center_shift = center_shift
+        if self.detect_face:
+            self.face_detector = MTCNN(thresh=[0.5, 0.6, 0.7])
+    def __len__(self):
+        return len(self.img_names)
+
+    def __getitem__(self, idx):
+        img_name = self.img_names[idx]
+        pil_image = Image.open(img_name)
+        if pil_image.mode != "RGB":
+            # if input is grayscale image, convert it to 3 channel image
+            if self.enhance:
+                pil_image = power_transform(pil_image, 0.5)
+            temp_image = Image.new('RGB', pil_image.size)
+            temp_image.paste(pil_image)
+            pil_image = temp_image
+        image = np.array(pil_image)
+        if self.gray_scale:
+            image = rgb2gray(image)
+            image = np.expand_dims(image, axis=2)
+            image = np.concatenate((image, image, image), axis=2)
+            image = image * 255.0
+            image = image.astype(np.uint8)
+        if not self.detect_face:
+            center = [450//2, 450//2+0]
+            if self.center_shift != 0:
+                center[0] += int(np.random.uniform(-self.center_shift,
+                                               self.center_shift))
+                center[1] += int(np.random.uniform(-self.center_shift,
+                                               self.center_shift))
+            scale = 1.8
+        else:
+            detected_faces = self.face_detector.detect_image(image)
+            if len(detected_faces) > 0:
+                box = detected_faces[0]
+                left, top, right, bottom, _ = box
+                center = [right - (right - left) / 2.0,
+                        bottom - (bottom - top) / 2.0]
+                center[1] = center[1] - (bottom - top) * 0.12
+                scale = (right - left + bottom - top) / 195.0
+            else:
+                center = [450//2, 450//2+0]
+                scale = 1.8
+            if self.center_shift != 0:
+                shift = self.center * self.center_shift / 450
+                center[0] += int(np.random.uniform(-shift, shift))
+                center[1] += int(np.random.uniform(-shift, shift))
+        base_name = os.path.basename(img_name)
+        landmarks_base_name = base_name[:-4] + '_pts.mat'
+        landmarks_name = os.path.join(self.landmarks_dir, landmarks_base_name)
+        if os.path.isfile(landmarks_name):
+            mat_data = sio.loadmat(landmarks_name)
+            landmarks = mat_data['pts_2d']
+        elif os.path.isfile(landmarks_name[:-8] + '.pts.npy'):
+            landmarks = np.load(landmarks_name[:-8] + '.pts.npy')
+        else:
+            landmarks = []
+            heatmap = []
+
+        if landmarks != []:
+            new_image, new_landmarks = cv_crop(image, landmarks, center,
+                                               scale, 256, self.center_shift)
+            tries = 0
+            while self.center_shift != 0 and tries < 5 and (np.max(new_landmarks) > 240 or np.min(new_landmarks) < 15):
+                center = [450//2, 450//2+0]
+                scale += 0.05
+                center[0] += int(np.random.uniform(-self.center_shift,
+                                            self.center_shift))
+                center[1] += int(np.random.uniform(-self.center_shift,
+                                            self.center_shift))
+
+                new_image, new_landmarks = cv_crop(image, landmarks,
+                                                    center, scale, 256,
+                                                    self.center_shift)
+                tries += 1
+            if np.max(new_landmarks) > 250 or np.min(new_landmarks) < 5:
+                center = [450//2, 450//2+0]
+                scale = 2.25
+                new_image, new_landmarks = cv_crop(image, landmarks,
+                                                    center, scale, 256,
+                                                    100)
+            assert (np.min(new_landmarks) > 0 and np.max(new_landmarks) < 256), \
+                "Landmarks out of boundary!"
+            image = new_image
+            landmarks = new_landmarks
+            heatmap = np.zeros((self.num_lanmdkars, 64, 64))
+            for i in range(self.num_lanmdkars):
+                if landmarks[i][0] > 0:
+                    heatmap[i] = draw_gaussian(heatmap[i], landmarks[i]/4.0+1, 1)
+        sample = {'image': image, 'heatmap': heatmap, 'landmarks': landmarks}
+        if self.transform:
+            sample = self.transform(sample)
+
+        return sample
+
+def get_dataset(val_img_dir, val_landmarks_dir, batch_size,
+                num_landmarks=68, rotation=0, scale=0,
+                center_shift=0, random_flip=False,
+                brightness=0, contrast=0, saturation=0,
+                blur=False, noise=False, jpeg_effect=False,
+                random_occlusion=False, gray_scale=False,
+                detect_face=False, enhance=False):
+    val_transforms = transforms.Compose([AddBoundary(num_landmarks),
+                                         AddWeightMap(),
+                                         ToTensor()])
+
+    val_dataset = FaceLandmarksDataset(val_img_dir, val_landmarks_dir,
+                                       num_landmarks=num_landmarks,
+                                       gray_scale=gray_scale,
+                                       detect_face=detect_face,
+                                       enhance=enhance,
+                                       transform=val_transforms)
+
+    val_dataloader = torch.utils.data.DataLoader(val_dataset,
+                                                   batch_size=batch_size,
+                                                   shuffle=False,
+                                                   num_workers=6)
+    data_loaders = {'val': val_dataloader}
+    dataset_sizes = {}
+    dataset_sizes['val'] = len(val_dataset)
+    return data_loaders, dataset_sizes

AdaptiveWingLoss/core/evaler.py

@@ -0,0 +1,116 @@
+import matplotlib
+matplotlib.use('Agg')
+import math
+import torch
+import copy
+import time
+from torch.autograd import Variable
+import shutil
+from skimage import io
+import numpy as np
+from utils.utils import fan_NME, show_landmarks, get_preds_fromhm
+from PIL import Image, ImageDraw
+import os
+import sys
+import cv2
+import matplotlib.pyplot as plt
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+def eval_model(model, dataloaders, dataset_sizes,
+               writer, use_gpu=True, epoches=5, dataset='val',
+               save_path='./', num_landmarks=68):
+    global_nme = 0
+    model.eval()
+    for epoch in range(epoches):
+        running_loss = 0
+        step = 0
+        total_nme = 0
+        total_count = 0
+        fail_count = 0
+        nmes = []
+        # running_corrects = 0
+
+        # Iterate over data.
+        with torch.no_grad():
+            for data in dataloaders[dataset]:
+                total_runtime = 0
+                run_count = 0
+                step_start = time.time()
+                step += 1
+                # get the inputs
+                inputs = data['image'].type(torch.FloatTensor)
+                labels_heatmap = data['heatmap'].type(torch.FloatTensor)
+                labels_boundary = data['boundary'].type(torch.FloatTensor)
+                landmarks = data['landmarks'].type(torch.FloatTensor)
+                loss_weight_map = data['weight_map'].type(torch.FloatTensor)
+                # wrap them in Variable
+                if use_gpu:
+                    inputs = inputs.to(device)
+                    labels_heatmap = labels_heatmap.to(device)
+                    labels_boundary = labels_boundary.to(device)
+                    loss_weight_map = loss_weight_map.to(device)
+                else:
+                    inputs, labels_heatmap = Variable(inputs), Variable(labels_heatmap)
+                    labels_boundary = Variable(labels_boundary)
+                labels = torch.cat((labels_heatmap, labels_boundary), 1)
+                single_start = time.time()
+                outputs, boundary_channels = model(inputs)
+                single_end = time.time()
+                total_runtime += time.time() - single_start
+                run_count += 1
+                step_end = time.time()
+                for i in range(inputs.shape[0]):
+                    img = inputs[i]
+                    img = img.cpu().numpy()
+                    img = img.transpose((1, 2, 0))*255.0
+                    img = img.astype(np.uint8)
+                    img = Image.fromarray(img)
+                    # pred_heatmap = outputs[-1][i].detach().cpu()[:-1, :, :]
+                    pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
+                    pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0))
+                    pred_landmarks = pred_landmarks.squeeze().numpy()
+
+                    gt_landmarks = data['landmarks'][i].numpy()
+                    if num_landmarks == 68:
+                        left_eye = np.average(gt_landmarks[36:42], axis=0)
+                        right_eye = np.average(gt_landmarks[42:48], axis=0)
+                        norm_factor = np.linalg.norm(left_eye - right_eye)
+                        # norm_factor = np.linalg.norm(gt_landmarks[36]- gt_landmarks[45])
+
+                    elif num_landmarks == 98:
+                        norm_factor = np.linalg.norm(gt_landmarks[60]- gt_landmarks[72])
+                    elif num_landmarks == 19:
+                        left, top = gt_landmarks[-2, :]
+                        right, bottom = gt_landmarks[-1, :]
+                        norm_factor = math.sqrt(abs(right - left)*abs(top-bottom))
+                        gt_landmarks = gt_landmarks[:-2, :]
+                    elif num_landmarks == 29:
+                        # norm_factor = np.linalg.norm(gt_landmarks[8]- gt_landmarks[9])
+                        norm_factor = np.linalg.norm(gt_landmarks[16]- gt_landmarks[17])
+                    single_nme = (np.sum(np.linalg.norm(pred_landmarks*4 - gt_landmarks, axis=1)) / pred_landmarks.shape[0]) / norm_factor
+
+                    nmes.append(single_nme)
+                    total_count += 1
+                    if single_nme > 0.1:
+                        fail_count += 1
+                if step % 10 == 0:
+                    print('Step {} Time: {:.6f} Input Mean: {:.6f} Output Mean: {:.6f}'.format(
+                        step, step_end - step_start,
+                        torch.mean(labels),
+                        torch.mean(outputs[0])))
+                # gt_landmarks = landmarks.numpy()
+                # pred_heatmap = outputs[-1].to('cpu').numpy()
+                gt_landmarks = landmarks
+                batch_nme = fan_NME(outputs[-1][:, :-1, :, :].detach().cpu(), gt_landmarks, num_landmarks)
+                # batch_nme = 0
+                total_nme += batch_nme
+        epoch_nme = total_nme / dataset_sizes['val']
+        global_nme += epoch_nme
+        nme_save_path = os.path.join(save_path, 'nme_log.npy')
+        np.save(nme_save_path, np.array(nmes))
+        print('NME: {:.6f} Failure Rate: {:.6f} Total Count: {:.6f} Fail Count: {:.6f}'.format(epoch_nme, fail_count/total_count, total_count, fail_count))
+    print('Evaluation done! Average NME: {:.6f}'.format(global_nme/epoches))
+    print('Everage runtime for a single batch: {:.6f}'.format(total_runtime/run_count))
+    return model

AdaptiveWingLoss/core/models.py

@@ -0,0 +1,228 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from AdaptiveWingLoss.core.coord_conv import CoordConvTh
+
+
+def conv3x3(in_planes, out_planes, strd=1, padding=1,
+            bias=False,dilation=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3,
+                     stride=strd, padding=padding, bias=bias,
+                     dilation=dilation)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        # self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        # self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        # out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        # out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_planes, out_planes):
+        super(ConvBlock, self).__init__()
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.conv1 = conv3x3(in_planes, int(out_planes / 2))
+        self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
+        self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4),
+                             padding=1, dilation=1)
+        self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
+        self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4),
+                             padding=1, dilation=1)
+
+        if in_planes != out_planes:
+            self.downsample = nn.Sequential(
+                nn.BatchNorm2d(in_planes),
+                nn.ReLU(True),
+                nn.Conv2d(in_planes, out_planes,
+                          kernel_size=1, stride=1, bias=False),
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        residual = x
+
+        out1 = self.bn1(x)
+        out1 = F.relu(out1, True)
+        out1 = self.conv1(out1)
+
+        out2 = self.bn2(out1)
+        out2 = F.relu(out2, True)
+        out2 = self.conv2(out2)
+
+        out3 = self.bn3(out2)
+        out3 = F.relu(out3, True)
+        out3 = self.conv3(out3)
+
+        out3 = torch.cat((out1, out2, out3), 1)
+
+        if self.downsample is not None:
+            residual = self.downsample(residual)
+
+        out3 += residual
+
+        return out3
+
+class HourGlass(nn.Module):
+    def __init__(self, num_modules, depth, num_features, first_one=False):
+        super(HourGlass, self).__init__()
+        self.num_modules = num_modules
+        self.depth = depth
+        self.features = num_features
+        self.coordconv = CoordConvTh(x_dim=64, y_dim=64,
+                                     with_r=True, with_boundary=True,
+                                     in_channels=256, first_one=first_one,
+                                     out_channels=256,
+                                     kernel_size=1,
+                                     stride=1, padding=0)
+        self._generate_network(self.depth)
+
+    def _generate_network(self, level):
+        self.add_module('b1_' + str(level), ConvBlock(256, 256))
+
+        self.add_module('b2_' + str(level), ConvBlock(256, 256))
+
+        if level > 1:
+            self._generate_network(level - 1)
+        else:
+            self.add_module('b2_plus_' + str(level), ConvBlock(256, 256))
+
+        self.add_module('b3_' + str(level), ConvBlock(256, 256))
+
+    def _forward(self, level, inp):
+        # Upper branch
+        up1 = inp
+        up1 = self._modules['b1_' + str(level)](up1)
+
+        # Lower branch
+        low1 = F.avg_pool2d(inp, 2, stride=2)
+        low1 = self._modules['b2_' + str(level)](low1)
+
+        if level > 1:
+            low2 = self._forward(level - 1, low1)
+        else:
+            low2 = low1
+            low2 = self._modules['b2_plus_' + str(level)](low2)
+
+        low3 = low2
+        low3 = self._modules['b3_' + str(level)](low3)
+
+        up2 = F.upsample(low3, scale_factor=2, mode='nearest')
+
+        return up1 + up2
+
+    def forward(self, x, heatmap):
+        x, last_channel = self.coordconv(x, heatmap)
+        return self._forward(self.depth, x), last_channel
+
+class FAN(nn.Module):
+
+    def __init__(self, num_modules=1, end_relu=False, gray_scale=False,
+                 num_landmarks=68):
+        super(FAN, self).__init__()
+        self.num_modules = num_modules
+        self.gray_scale = gray_scale
+        self.end_relu = end_relu
+        self.num_landmarks = num_landmarks
+
+        # Base part
+        if self.gray_scale:
+            self.conv1 = CoordConvTh(x_dim=256, y_dim=256,
+                                     with_r=True, with_boundary=False,
+                                     in_channels=3, out_channels=64,
+                                     kernel_size=7,
+                                     stride=2, padding=3)
+        else:
+            self.conv1 = CoordConvTh(x_dim=256, y_dim=256,
+                                     with_r=True, with_boundary=False,
+                                     in_channels=3, out_channels=64,
+                                     kernel_size=7,
+                                     stride=2, padding=3)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.conv2 = ConvBlock(64, 128)
+        self.conv3 = ConvBlock(128, 128)
+        self.conv4 = ConvBlock(128, 256)
+
+        # Stacking part
+        for hg_module in range(self.num_modules):
+            if hg_module == 0:
+                first_one = True
+            else:
+                first_one = False
+            self.add_module('m' + str(hg_module), HourGlass(1, 4, 256,
+                                                            first_one))
+            self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256))
+            self.add_module('conv_last' + str(hg_module),
+                            nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+            self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
+            self.add_module('l' + str(hg_module), nn.Conv2d(256,
+                                                            num_landmarks+1, kernel_size=1, stride=1, padding=0))
+
+            if hg_module < self.num_modules - 1:
+                self.add_module(
+                    'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+                self.add_module('al' + str(hg_module), nn.Conv2d(num_landmarks+1,
+                                                                 256, kernel_size=1, stride=1, padding=0))
+
+    def forward(self, x):
+        x, _ = self.conv1(x)
+        x = F.relu(self.bn1(x), True)
+        # x = F.relu(self.bn1(self.conv1(x)), True)
+        x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        previous = x
+
+        outputs = []
+        boundary_channels = []
+        tmp_out = None
+        for i in range(self.num_modules):
+            hg, boundary_channel = self._modules['m' + str(i)](previous,
+                                                               tmp_out)
+
+            ll = hg
+            ll = self._modules['top_m_' + str(i)](ll)
+
+            ll = F.relu(self._modules['bn_end' + str(i)]
+                        (self._modules['conv_last' + str(i)](ll)), True)
+
+            # Predict heatmaps
+            tmp_out = self._modules['l' + str(i)](ll)
+            if self.end_relu:
+                tmp_out = F.relu(tmp_out) # HACK: Added relu
+            outputs.append(tmp_out)
+            boundary_channels.append(boundary_channel)
+
+            if i < self.num_modules - 1:
+                ll = self._modules['bl' + str(i)](ll)
+                tmp_out_ = self._modules['al' + str(i)](tmp_out)
+                previous = previous + ll + tmp_out_
+
+        return outputs, boundary_channels

AdaptiveWingLoss/utils/utils.py

@@ -0,0 +1,355 @@
+from __future__ import print_function, division
+import os
+import sys
+import math
+import torch
+import cv2
+from PIL import Image
+from skimage import io
+from skimage import transform as ski_transform
+from scipy import ndimage
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+
+def _gaussian(
+        size=3, sigma=0.25, amplitude=1, normalize=False, width=None,
+        height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,
+        mean_vert=0.5):
+    # handle some defaults
+    if width is None:
+        width = size
+    if height is None:
+        height = size
+    if sigma_horz is None:
+        sigma_horz = sigma
+    if sigma_vert is None:
+        sigma_vert = sigma
+    center_x = mean_horz * width + 0.5
+    center_y = mean_vert * height + 0.5
+    gauss = np.empty((height, width), dtype=np.float32)
+    # generate kernel
+    for i in range(height):
+        for j in range(width):
+            gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (
+                sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))
+    if normalize:
+        gauss = gauss / np.sum(gauss)
+    return gauss
+
+def draw_gaussian(image, point, sigma):
+    # Check if the gaussian is inside
+    ul = [np.floor(np.floor(point[0]) - 3 * sigma),
+          np.floor(np.floor(point[1]) - 3 * sigma)]
+    br = [np.floor(np.floor(point[0]) + 3 * sigma),
+          np.floor(np.floor(point[1]) + 3 * sigma)]
+    if (ul[0] > image.shape[1] or ul[1] >
+            image.shape[0] or br[0] < 1 or br[1] < 1):
+        return image
+    size = 6 * sigma + 1
+    g = _gaussian(size)
+    g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) -
+           int(max(1, ul[0])) + int(max(1, -ul[0]))]
+    g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) -
+           int(max(1, ul[1])) + int(max(1, -ul[1]))]
+    img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+    img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+    assert (g_x[0] > 0 and g_y[1] > 0)
+    correct = False
+    while not correct:
+        try:
+            image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]
+            ] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]
+            correct = True
+        except:
+            print('img_x: {}, img_y: {}, g_x:{}, g_y:{}, point:{}, g_shape:{}, ul:{}, br:{}'.format(img_x, img_y, g_x, g_y, point, g.shape, ul, br))
+            ul = [np.floor(np.floor(point[0]) - 3 * sigma),
+                np.floor(np.floor(point[1]) - 3 * sigma)]
+            br = [np.floor(np.floor(point[0]) + 3 * sigma),
+                np.floor(np.floor(point[1]) + 3 * sigma)]
+            g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) -
+                int(max(1, ul[0])) + int(max(1, -ul[0]))]
+            g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) -
+                int(max(1, ul[1])) + int(max(1, -ul[1]))]
+            img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+            img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+            pass
+    image[image > 1] = 1
+    return image
+
+def transform(point, center, scale, resolution, rotation=0, invert=False):
+    _pt = np.ones(3)
+    _pt[0] = point[0]
+    _pt[1] = point[1]
+
+    h = 200.0 * scale
+    t = np.eye(3)
+    t[0, 0] = resolution / h
+    t[1, 1] = resolution / h
+    t[0, 2] = resolution * (-center[0] / h + 0.5)
+    t[1, 2] = resolution * (-center[1] / h + 0.5)
+
+    if rotation != 0:
+        rotation = -rotation
+        r = np.eye(3)
+        ang = rotation * math.pi / 180.0
+        s = math.sin(ang)
+        c = math.cos(ang)
+        r[0][0] = c
+        r[0][1] = -s
+        r[1][0] = s
+        r[1][1] = c
+
+        t_ = np.eye(3)
+        t_[0][2] = -resolution / 2.0
+        t_[1][2] = -resolution / 2.0
+        t_inv = torch.eye(3)
+        t_inv[0][2] = resolution / 2.0
+        t_inv[1][2] = resolution / 2.0
+        t = reduce(np.matmul, [t_inv, r, t_, t])
+
+    if invert:
+        t = np.linalg.inv(t)
+    new_point = (np.matmul(t, _pt))[0:2]
+
+    return new_point.astype(int)
+
+def cv_crop(image, landmarks, center, scale, resolution=256, center_shift=0):
+    new_image = cv2.copyMakeBorder(image, center_shift,
+                                   center_shift,
+                                   center_shift,
+                                   center_shift,
+                                   cv2.BORDER_CONSTANT, value=[0,0,0])
+    new_landmarks = landmarks.copy()
+    if center_shift != 0:
+        center[0] += center_shift
+        center[1] += center_shift
+        new_landmarks = new_landmarks + center_shift
+    length = 200 * scale
+    top = int(center[1] - length // 2)
+    bottom = int(center[1] + length // 2)
+    left = int(center[0] - length // 2)
+    right = int(center[0] + length // 2)
+    y_pad = abs(min(top, new_image.shape[0] - bottom, 0))
+    x_pad = abs(min(left, new_image.shape[1] - right, 0))
+    top, bottom, left, right = top + y_pad, bottom + y_pad, left + x_pad, right + x_pad
+    new_image = cv2.copyMakeBorder(new_image, y_pad,
+                                   y_pad,
+                                   x_pad,
+                                   x_pad,
+                                   cv2.BORDER_CONSTANT, value=[0,0,0])
+    new_image = new_image[top:bottom, left:right]
+    new_image = cv2.resize(new_image, dsize=(int(resolution), int(resolution)),
+                           interpolation=cv2.INTER_LINEAR)
+    new_landmarks[:, 0] = (new_landmarks[:, 0] + x_pad - left) * resolution / length
+    new_landmarks[:, 1] = (new_landmarks[:, 1] + y_pad - top) * resolution / length
+    return new_image, new_landmarks
+
+def cv_rotate(image, landmarks, heatmap, rot, scale, resolution=256):
+    img_mat = cv2.getRotationMatrix2D((resolution//2, resolution//2), rot, scale)
+    ones = np.ones(shape=(landmarks.shape[0], 1))
+    stacked_landmarks = np.hstack([landmarks, ones])
+    new_landmarks = img_mat.dot(stacked_landmarks.T).T
+    if np.max(new_landmarks) > 255 or np.min(new_landmarks) < 0:
+        return image, landmarks, heatmap
+    else:
+        new_image = cv2.warpAffine(image, img_mat, (resolution, resolution))
+        if heatmap is not None:
+            new_heatmap = np.zeros((heatmap.shape[0], 64, 64))
+            for i in range(heatmap.shape[0]):
+                if new_landmarks[i][0] > 0:
+                    new_heatmap[i] = draw_gaussian(new_heatmap[i],
+                                                   new_landmarks[i]/4.0+1, 1)
+        return new_image, new_landmarks, new_heatmap
+
+def show_landmarks(image, heatmap, gt_landmarks, gt_heatmap):
+    """Show image with pred_landmarks"""
+    pred_landmarks = []
+    pred_landmarks, _ = get_preds_fromhm(torch.from_numpy(heatmap).unsqueeze(0))
+    pred_landmarks = pred_landmarks.squeeze()*4
+
+    # pred_landmarks2 = get_preds_fromhm2(heatmap)
+    heatmap = np.max(gt_heatmap, axis=0)
+    heatmap = heatmap / np.max(heatmap)
+    # image = ski_transform.resize(image, (64, 64))*255
+    image = image.astype(np.uint8)
+    heatmap = np.max(gt_heatmap, axis=0)
+    heatmap = ski_transform.resize(heatmap, (image.shape[0], image.shape[1]))
+    heatmap *= 255
+    heatmap = heatmap.astype(np.uint8)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+    plt.imshow(image)
+    plt.scatter(gt_landmarks[:, 0], gt_landmarks[:, 1], s=0.5, marker='.', c='g')
+    plt.scatter(pred_landmarks[:, 0], pred_landmarks[:, 1], s=0.5, marker='.', c='r')
+    plt.pause(0.001)  # pause a bit so that plots are updated
+
+def fan_NME(pred_heatmaps, gt_landmarks, num_landmarks=68):
+    '''
+       Calculate total NME for a batch of data
+
+       Args:
+           pred_heatmaps: torch tensor of size [batch, points, height, width]
+           gt_landmarks: torch tesnsor of size [batch, points, x, y]
+
+       Returns:
+           nme: sum of nme for this batch
+    '''
+    nme = 0
+    pred_landmarks, _ = get_preds_fromhm(pred_heatmaps)
+    pred_landmarks = pred_landmarks.numpy()
+    gt_landmarks = gt_landmarks.numpy()
+    for i in range(pred_landmarks.shape[0]):
+        pred_landmark = pred_landmarks[i] * 4.0
+        gt_landmark = gt_landmarks[i]
+
+        if num_landmarks == 68:
+            left_eye = np.average(gt_landmark[36:42], axis=0)
+            right_eye = np.average(gt_landmark[42:48], axis=0)
+            norm_factor = np.linalg.norm(left_eye - right_eye)
+            # norm_factor = np.linalg.norm(gt_landmark[36]- gt_landmark[45])
+        elif num_landmarks == 98:
+            norm_factor = np.linalg.norm(gt_landmark[60]- gt_landmark[72])
+        elif num_landmarks == 19:
+            left, top = gt_landmark[-2, :]
+            right, bottom = gt_landmark[-1, :]
+            norm_factor = math.sqrt(abs(right - left)*abs(top-bottom))
+            gt_landmark = gt_landmark[:-2, :]
+        elif num_landmarks == 29:
+            # norm_factor = np.linalg.norm(gt_landmark[8]- gt_landmark[9])
+            norm_factor = np.linalg.norm(gt_landmark[16]- gt_landmark[17])
+        nme += (np.sum(np.linalg.norm(pred_landmark - gt_landmark, axis=1)) / pred_landmark.shape[0]) / norm_factor
+    return nme
+
+def fan_NME_hm(pred_heatmaps, gt_heatmaps, num_landmarks=68):
+    '''
+       Calculate total NME for a batch of data
+
+       Args:
+           pred_heatmaps: torch tensor of size [batch, points, height, width]
+           gt_landmarks: torch tesnsor of size [batch, points, x, y]
+
+       Returns:
+           nme: sum of nme for this batch
+    '''
+    nme = 0
+    pred_landmarks, _ = get_index_fromhm(pred_heatmaps)
+    pred_landmarks = pred_landmarks.numpy()
+    gt_landmarks = gt_landmarks.numpy()
+    for i in range(pred_landmarks.shape[0]):
+        pred_landmark = pred_landmarks[i] * 4.0
+        gt_landmark = gt_landmarks[i]
+        if num_landmarks == 68:
+            left_eye = np.average(gt_landmark[36:42], axis=0)
+            right_eye = np.average(gt_landmark[42:48], axis=0)
+            norm_factor = np.linalg.norm(left_eye - right_eye)
+        else:
+            norm_factor = np.linalg.norm(gt_landmark[60]- gt_landmark[72])
+        nme += (np.sum(np.linalg.norm(pred_landmark - gt_landmark, axis=1)) / pred_landmark.shape[0]) / norm_factor
+    return nme
+
+def power_transform(img, power):
+    img = np.array(img)
+    img_new = np.power((img/255.0), power) * 255.0
+    img_new = img_new.astype(np.uint8)
+    img_new = Image.fromarray(img_new)
+    return img_new
+
+def get_preds_fromhm(hm, center=None, scale=None, rot=None):
+    max, idx = torch.max(
+        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+    idx += 1
+    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+    preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
+    preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
+    
+    
+    for i in range(preds.size(0)):
+        for j in range(preds.size(1)):
+            hm_ = hm[i, j, :]
+            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                diff = torch.FloatTensor(
+                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],
+                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+                preds[i, j].add_(diff.sign_().mul_(.25))
+
+    preds.add_(-0.5)
+
+    preds_orig = torch.zeros(preds.size())
+    if center is not None and scale is not None:
+        for i in range(hm.size(0)):
+            for j in range(hm.size(1)):
+                preds_orig[i, j] = transform(
+                    preds[i, j], center, scale, hm.size(2), rot, True)
+
+    return preds, preds_orig
+
+def get_index_fromhm(hm):
+    max, idx = torch.max(
+        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+    preds[..., 0].remainder_(hm.size(3))
+    preds[..., 1].div_(hm.size(2)).floor_()
+
+    for i in range(preds.size(0)):
+        for j in range(preds.size(1)):
+            hm_ = hm[i, j, :]
+            pX, pY = int(preds[i, j, 0]), int(preds[i, j, 1])
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                diff = torch.FloatTensor(
+                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],
+                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+                preds[i, j].add_(diff.sign_().mul_(.25))
+
+    return preds
+
+def shuffle_lr(parts, num_landmarks=68, pairs=None):
+    if num_landmarks == 68:
+        if pairs is None:
+            pairs = [[0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11], [6, 10],
+                    [7, 9], [17, 26], [18, 25], [19, 24], [20, 23], [21, 22], [36, 45],
+                    [37, 44], [38, 43], [39, 42], [41, 46], [40, 47], [31, 35], [32, 34],
+                    [50, 52], [49, 53], [48, 54], [61, 63], [60, 64], [67, 65], [59, 55], [58, 56]]
+    elif num_landmarks == 98:
+        if pairs is None:
+            pairs = [[0, 32], [1,31], [2, 30], [3, 29], [4, 28], [5, 27], [6, 26], [7, 25], [8, 24], [9, 23], [10, 22], [11, 21], [12, 20], [13, 19], [14, 18], [15, 17], [33, 46], [34, 45], [35, 44], [36, 43], [37, 42], [38, 50], [39, 49], [40, 48], [41, 47], [60, 72], [61, 71], [62, 70], [63, 69], [64, 68], [65, 75], [66, 74], [67, 73], [96, 97], [55, 59], [56, 58], [76, 82], [77, 81], [78, 80], [88, 92], [89, 91], [95, 93], [87, 83], [86, 84]]
+    elif num_landmarks == 19:
+        if pairs is None:
+            pairs = [[0, 5], [1, 4], [2, 3], [6, 11], [7, 10], [8, 9], [12, 14], [15, 17]]
+    elif num_landmarks == 29:
+        if pairs is None:
+            pairs = [[0, 1], [4, 6], [5, 7], [2, 3], [8, 9], [12, 14], [16, 17], [13, 15], [10, 11], [18, 19], [22, 23]]
+    for matched_p in pairs:
+        idx1, idx2 = matched_p[0], matched_p[1]
+        tmp = np.copy(parts[idx1])
+        np.copyto(parts[idx1], parts[idx2])
+        np.copyto(parts[idx2], tmp)
+    return parts
+
+
+def generate_weight_map(weight_map,heatmap):
+
+    k_size = 3
+    dilate = ndimage.grey_dilation(heatmap ,size=(k_size,k_size))
+    weight_map[np.where(dilate>0.2)] = 1
+    return weight_map
+
+def fig2data(fig):
+    """
+    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
+    @param fig a matplotlib figure
+    @return a numpy 3D array of RGBA values
+    """
+    # draw the renderer
+    fig.canvas.draw ( )
+
+    # Get the RGB buffer from the figure
+    w,h = fig.canvas.get_width_height()
+    buf = np.fromstring (fig.canvas.tostring_rgb(), dtype=np.uint8)
+    buf.shape = (w, h, 3)
+
+    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
+    buf = np.roll (buf, 3, axis=2)
+    return buf

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

apex/.gitignore

@@ -0,0 +1,6 @@
+apex.egg-info
+dist
+build
+docs/build
+*~
+__pycache__

apex/.gitmodules

@@ -0,0 +1,7 @@
+[submodule "apex/contrib/csrc/multihead_attn/cutlass"]
+	path = apex/contrib/csrc/multihead_attn/cutlass
+	url = https://github.com/NVIDIA/cutlass.git
+	branch = v1.2.0
+[submodule "apex/contrib/csrc/cudnn-frontend"]
+	path = apex/contrib/csrc/cudnn-frontend
+	url = https://github.com/NVIDIA/cudnn-frontend.git

apex/LICENSE

@@ -0,0 +1,11 @@
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

apex/README.md

@@ -0,0 +1,146 @@
+# Introduction
+
+This repository holds NVIDIA-maintained utilities to streamline
+mixed precision and distributed training in Pytorch.
+Some of the code here will be included in upstream Pytorch eventually.
+The intention of Apex is to make up-to-date utilities available to
+users as quickly as possible.
+
+## Full API Documentation: [https://nvidia.github.io/apex](https://nvidia.github.io/apex)
+
+## [GTC 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/GTC_2019) and [Pytorch DevCon 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/Pytorch_Devcon_2019) Slides
+
+# Contents
+
+## 1. Amp:  Automatic Mixed Precision
+
+`apex.amp` is a tool to enable mixed precision training by changing only 3 lines of your script.
+Users can easily experiment with different pure and mixed precision training modes by supplying
+different flags to `amp.initialize`.
+
+[Webinar introducing Amp](https://info.nvidia.com/webinar-mixed-precision-with-pytorch-reg-page.html)
+(The flag `cast_batchnorm` has been renamed to `keep_batchnorm_fp32`).
+
+[API Documentation](https://nvidia.github.io/apex/amp.html)
+
+[Comprehensive Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+
+[DCGAN example coming soon...](https://github.com/NVIDIA/apex/tree/master/examples/dcgan)
+
+[Moving to the new Amp API](https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users) (for users of the deprecated "Amp" and "FP16_Optimizer" APIs)
+
+## 2. Distributed Training
+
+`apex.parallel.DistributedDataParallel` is a module wrapper, similar to
+`torch.nn.parallel.DistributedDataParallel`.  It enables convenient multiprocess distributed training,
+optimized for NVIDIA's NCCL communication library.
+
+[API Documentation](https://nvidia.github.io/apex/parallel.html)
+
+[Python Source](https://github.com/NVIDIA/apex/tree/master/apex/parallel)
+
+[Example/Walkthrough](https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed)
+
+The [Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+shows use of `apex.parallel.DistributedDataParallel` along with `apex.amp`.
+
+### Synchronized Batch Normalization
+
+`apex.parallel.SyncBatchNorm` extends `torch.nn.modules.batchnorm._BatchNorm` to
+support synchronized BN.
+It allreduces stats across processes during multiprocess (DistributedDataParallel) training.
+Synchronous BN has been used in cases where only a small
+local minibatch can fit on each GPU.
+Allreduced stats increase the effective batch size for the BN layer to the
+global batch size across all processes (which, technically, is the correct
+formulation).
+Synchronous BN has been observed to improve converged accuracy in some of our research models.
+
+### Checkpointing
+
+To properly save and load your `amp` training, we introduce the `amp.state_dict()`, which contains all `loss_scalers` and their corresponding unskipped steps,
+as well as `amp.load_state_dict()` to restore these attributes.
+
+In order to get bitwise accuracy, we recommend the following workflow:
+```python
+# Initialization
+opt_level = 'O1'
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+
+# Train your model
+...
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+    scaled_loss.backward()
+...
+
+# Save checkpoint
+checkpoint = {
+    'model': model.state_dict(),
+    'optimizer': optimizer.state_dict(),
+    'amp': amp.state_dict()
+}
+torch.save(checkpoint, 'amp_checkpoint.pt')
+...
+
+# Restore
+model = ...
+optimizer = ...
+checkpoint = torch.load('amp_checkpoint.pt')
+
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+model.load_state_dict(checkpoint['model'])
+optimizer.load_state_dict(checkpoint['optimizer'])
+amp.load_state_dict(checkpoint['amp'])
+
+# Continue training
+...
+```
+
+Note that we recommend restoring the model using the same `opt_level`. Also note that we recommend calling the `load_state_dict` methods after `amp.initialize`.
+
+# Requirements
+
+Python 3
+
+CUDA 9 or newer
+
+PyTorch 0.4 or newer.  The CUDA and C++ extensions require pytorch 1.0 or newer.
+
+We recommend the latest stable release, obtainable from
+[https://pytorch.org/](https://pytorch.org/).  We also test against the latest master branch, obtainable from [https://github.com/pytorch/pytorch](https://github.com/pytorch/pytorch).
+
+It's often convenient to use Apex in Docker containers.  Compatible options include:
+* [NVIDIA Pytorch containers from NGC](https://ngc.nvidia.com/catalog/containers/nvidia%2Fpytorch), which come with Apex preinstalled.  To use the latest Amp API, you may need to `pip uninstall apex` then reinstall Apex using the **Quick Start** commands below.
+* [official Pytorch -devel Dockerfiles](https://hub.docker.com/r/pytorch/pytorch/tags), e.g. `docker pull pytorch/pytorch:nightly-devel-cuda10.0-cudnn7`, in which you can install Apex using the **Quick Start** commands.
+
+See the [Docker example folder](https://github.com/NVIDIA/apex/tree/master/examples/docker) for details.
+
+# Quick Start
+
+### Linux
+
+For performance and full functionality, we recommend installing Apex with
+CUDA and C++ extensions via
+```
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+```
+
+Apex also supports a Python-only build (required with Pytorch 0.4) via
+```
+pip install -v --disable-pip-version-check --no-cache-dir ./
+```
+A Python-only build omits:
+- Fused kernels required to use `apex.optimizers.FusedAdam`.
+- Fused kernels required to use `apex.normalization.FusedLayerNorm`.
+- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.
+- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.
+`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.
+
+Pyprof support has been moved to its own [dedicated repository](https://github.com/NVIDIA/PyProf).
+The codebase is deprecated in Apex and will be removed soon.
+
+### Windows support
+Windows support is experimental, and Linux is recommended.  `pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .` may work if you were able to build Pytorch from source
+on your system.  `pip install -v --no-cache-dir .` (without CUDA/C++ extensions) is more likely to work.  If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.

apex/apex/RNN/README.md

@@ -0,0 +1 @@
+Under construction...

apex/apex/RNN/RNNBackend.py

@@ -0,0 +1,365 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+import torch.nn.functional as F
+
+import math
+
+
+def is_iterable(maybe_iterable):
+    return isinstance(maybe_iterable, list) or isinstance(maybe_iterable, tuple)
+
+
+def flatten_list(tens_list):
+    """
+    flatten_list
+    """
+    if not is_iterable(tens_list):
+        return tens_list
+    
+    return torch.cat(tens_list, dim=0).view(len(tens_list), *tens_list[0].size() )
+
+    
+#These modules always assumes batch_first
+class bidirectionalRNN(nn.Module):
+    """
+    bidirectionalRNN
+    """
+    def __init__(self, inputRNN, num_layers=1, dropout = 0):
+        super(bidirectionalRNN, self).__init__()
+        self.dropout = dropout
+        self.fwd = stackedRNN(inputRNN, num_layers=num_layers, dropout = dropout)
+        self.bckwrd = stackedRNN(inputRNN.new_like(), num_layers=num_layers, dropout = dropout)
+        self.rnns = nn.ModuleList([self.fwd, self.bckwrd])
+        
+    #collect hidden option will return all hidden/cell states from entire RNN
+    def forward(self, input, collect_hidden=False):
+        """
+        forward()
+        """
+        seq_len = input.size(0)
+        bsz = input.size(1)
+
+        fwd_out, fwd_hiddens = list(self.fwd(input, collect_hidden = collect_hidden))
+        bckwrd_out, bckwrd_hiddens = list(self.bckwrd(input, reverse=True, collect_hidden = collect_hidden))
+        
+        output = torch.cat( [fwd_out, bckwrd_out], -1 )
+        hiddens = tuple( torch.cat(hidden, -1) for hidden in zip( fwd_hiddens, bckwrd_hiddens) )
+
+        return output, hiddens
+
+    def reset_parameters(self):
+        """
+        reset_parameters()
+        """
+        for rnn in self.rnns:
+            rnn.reset_parameters()
+        
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.detachHidden()
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.reset_hidden(bsz)
+
+    def init_inference(self, bsz):    
+        """
+        init_inference()
+        """
+        for rnn in self.rnns:
+            rnn.init_inference(bsz)
+
+   
+#assumes hidden_state[0] of inputRNN is output hidden state
+#constructor either takes an RNNCell or list of RNN layers
+class stackedRNN(nn.Module):        
+    """
+    stackedRNN
+    """
+    def __init__(self, inputRNN, num_layers=1, dropout=0):
+        super(stackedRNN, self).__init__()
+        
+        self.dropout = dropout
+        
+        if isinstance(inputRNN, RNNCell):
+            self.rnns = [inputRNN]
+            for i in range(num_layers-1):
+                self.rnns.append(inputRNN.new_like(inputRNN.output_size))
+        elif isinstance(inputRNN, list):
+            assert len(inputRNN) == num_layers, "RNN list length must be equal to num_layers"
+            self.rnns=inputRNN
+        else:
+            raise RuntimeError()
+        
+        self.nLayers = len(self.rnns)
+        
+        self.rnns = nn.ModuleList(self.rnns)
+
+
+    '''
+    Returns output as hidden_state[0] Tensor([sequence steps][batch size][features])
+    If collect hidden will also return Tuple(
+        [n_hidden_states][sequence steps] Tensor([layer][batch size][features])
+    )
+    If not collect hidden will also return Tuple(
+        [n_hidden_states] Tensor([layer][batch size][features])
+    '''
+    def forward(self, input, collect_hidden=False, reverse=False):
+        """
+        forward()
+        """
+        seq_len = input.size(0)
+        bsz = input.size(1)
+        inp_iter = reversed(range(seq_len)) if reverse else range(seq_len)
+
+        hidden_states = [[] for i in range(self.nLayers)]
+        outputs = []
+
+        for seq in inp_iter:
+            for layer in range(self.nLayers):
+
+                if layer == 0:
+                    prev_out = input[seq]
+                    
+                outs = self.rnns[layer](prev_out)
+
+                if collect_hidden:
+                    hidden_states[layer].append(outs)
+                elif seq == seq_len-1:
+                    hidden_states[layer].append(outs)
+                    
+                prev_out = outs[0]
+
+            outputs.append(prev_out)
+
+        if reverse:
+            outputs = list(reversed(outputs))
+        '''
+        At this point outputs is in format:
+        list( [seq_length] x Tensor([bsz][features]) )
+        need to convert it to:
+        list( Tensor([seq_length][bsz][features]) )
+        '''
+        output = flatten_list(outputs)
+
+        '''
+        hidden_states at this point is in format:
+        list( [layer][seq_length][hidden_states] x Tensor([bsz][features]) )
+        need to convert it to:
+          For not collect hidden:
+            list( [hidden_states] x Tensor([layer][bsz][features]) )
+          For collect hidden:
+            list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+        '''
+        if not collect_hidden:
+            seq_len = 1
+        n_hid = self.rnns[0].n_hidden_states
+        new_hidden = [ [ [ None for k in range(self.nLayers)] for j in range(seq_len) ] for i in range(n_hid) ]
+
+
+        for i in range(n_hid):
+            for j in range(seq_len):
+                for k in range(self.nLayers):
+                    new_hidden[i][j][k] = hidden_states[k][j][i]
+
+        hidden_states = new_hidden
+        #Now in format list( [hidden_states][seq_length][layer] x Tensor([bsz][features]) )
+        #Reverse seq_length if reverse
+        if reverse:
+            hidden_states = list( list(reversed(list(entry))) for entry in hidden_states)
+
+        #flatten layer dimension into tensor
+        hiddens = list( list(
+            flatten_list(seq) for seq in hidden )
+                        for hidden in hidden_states )
+        
+        #Now in format list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+        #Remove seq_length dimension if not collect_hidden
+        if not collect_hidden:
+            hidden_states = list( entry[0] for entry in hidden_states)
+        return output, hidden_states
+    
+    def reset_parameters(self):
+        """
+        reset_parameters()
+        """
+        for rnn in self.rnns:
+            rnn.reset_parameters()
+        
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.detach_hidden()
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.reset_hidden(bsz)
+
+    def init_inference(self, bsz):    
+        """ 
+        init_inference()
+        """
+        for rnn in self.rnns:
+            rnn.init_inference(bsz)
+
+class RNNCell(nn.Module):
+    """ 
+    RNNCell 
+    gate_multiplier is related to the architecture you're working with
+    For LSTM-like it will be 4 and GRU-like will be 3.
+    Always assumes input is NOT batch_first.
+    Output size that's not hidden size will use output projection
+    Hidden_states is number of hidden states that are needed for cell
+    if one will go directly to cell as tensor, if more will go as list
+    """
+    def __init__(self, gate_multiplier, input_size, hidden_size, cell, n_hidden_states = 2, bias = False, output_size = None):
+        super(RNNCell, self).__init__()
+
+        self.gate_multiplier = gate_multiplier
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.cell = cell
+        self.bias = bias
+        self.output_size = output_size
+        if output_size is None:
+            self.output_size = hidden_size
+
+        self.gate_size = gate_multiplier * self.hidden_size
+        self.n_hidden_states = n_hidden_states
+
+        self.w_ih = nn.Parameter(torch.Tensor(self.gate_size, self.input_size))
+        self.w_hh = nn.Parameter(torch.Tensor(self.gate_size, self.output_size))
+
+        #Check if there's recurrent projection
+        if(self.output_size != self.hidden_size):
+            self.w_ho = nn.Parameter(torch.Tensor(self.output_size, self.hidden_size))
+
+        self.b_ih = self.b_hh = None
+        if self.bias:
+            self.b_ih = nn.Parameter(torch.Tensor(self.gate_size))
+            self.b_hh = nn.Parameter(torch.Tensor(self.gate_size))
+            
+        #hidden states for forward
+        self.hidden = [ None for states in range(self.n_hidden_states)]
+
+        self.reset_parameters()
+
+    def new_like(self, new_input_size=None):
+        """
+        new_like()
+        """
+        if new_input_size is None:
+            new_input_size = self.input_size
+            
+        return type(self)(self.gate_multiplier,
+                       new_input_size,
+                       self.hidden_size,
+                       self.cell,
+                       self.n_hidden_states,
+                       self.bias,
+                       self.output_size)
+
+    
+    #Use xavier where we can (weights), otherwise use uniform (bias)
+    def reset_parameters(self, gain=1):
+        """
+        reset_parameters()
+        """
+        stdev = 1.0 / math.sqrt(self.hidden_size)
+        for param in self.parameters():
+            param.data.uniform_(-stdev, stdev)
+    '''
+    Xavier reset:
+    def reset_parameters(self, gain=1):
+        stdv = 1.0 / math.sqrt(self.gate_size)
+
+        for param in self.parameters():
+            if (param.dim() > 1):
+                torch.nn.init.xavier_normal(param, gain)
+            else:
+                param.data.uniform_(-stdv, stdv)
+    '''
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for param in self.parameters():
+            if param is not None:
+                a_param = param
+                break
+
+        for i, _ in enumerate(self.hidden):
+            if(self.hidden[i] is None or self.hidden[i].data.size()[0] != bsz):
+
+                if i==0:
+                    hidden_size = self.output_size
+                else:
+                    hidden_size = self.hidden_size
+
+                tens = a_param.data.new(bsz, hidden_size).zero_()
+                self.hidden[i] = Variable(tens, requires_grad=False)
+            
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for i, _ in enumerate(self.hidden):
+            self.hidden[i] = None
+        self.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for i, _ in enumerate(self.hidden):
+            if self.hidden[i] is None:
+                raise RuntimeError("Must initialize hidden state before you can detach it")
+        for i, _ in enumerate(self.hidden):
+            self.hidden[i] = self.hidden[i].detach()
+        
+    def forward(self, input):
+        """
+        forward()
+        if not inited or bsz has changed this will create hidden states
+        """
+        self.init_hidden(input.size()[0])
+
+        hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+        self.hidden = self.cell(input, hidden_state, self.w_ih, self.w_hh, b_ih=self.b_ih, b_hh=self.b_hh)
+        if(self.n_hidden_states > 1):
+            self.hidden = list(self.hidden)
+        else:
+            self.hidden=[self.hidden]
+
+        if self.output_size != self.hidden_size:
+            self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+
+        return tuple(self.hidden)

apex/apex/RNN/__init__.py

@@ -0,0 +1,3 @@
+from .models import LSTM, GRU, ReLU, Tanh, mLSTM
+
+__all__ = ['models']

apex/apex/RNN/cells.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .RNNBackend import RNNCell
+
+from torch.nn._functions.thnn import rnnFusedPointwise as fusedBackend
+
+import math 
+
+
+class mLSTMRNNCell(RNNCell):
+    """
+    mLSTMRNNCell
+    """
+
+    def __init__(self, input_size, hidden_size, bias = False, output_size = None):
+        gate_multiplier = 4
+        super(mLSTMRNNCell, self).__init__(gate_multiplier, input_size, hidden_size, mLSTMCell, n_hidden_states = 2, bias = bias, output_size = output_size)
+
+        self.w_mih = nn.Parameter(torch.Tensor(self.output_size, self.input_size))
+        self.w_mhh = nn.Parameter(torch.Tensor(self.output_size, self.output_size))
+
+        self.reset_parameters()
+
+    def forward(self, input):
+        """
+        mLSTMRNNCell.forward()
+        """
+        #if not inited or bsz has changed this will create hidden states
+        self.init_hidden(input.size()[0])
+
+        hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+
+        self.hidden = list(
+                           self.cell(input, hidden_state, self.w_ih, self.w_hh, self.w_mih, self.w_mhh,
+                           b_ih=self.b_ih, b_hh=self.b_hh)
+        )
+        
+        if self.output_size != self.hidden_size:
+            self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+        return tuple(self.hidden)
+
+
+    def new_like(self, new_input_size=None):
+        if new_input_size is None:
+            new_input_size = self.input_size
+        
+        return type(self)(
+            new_input_size,
+            self.hidden_size,
+            self.bias,
+            self.output_size)
+
+def mLSTMCell(input, hidden, w_ih, w_hh, w_mih, w_mhh, b_ih=None, b_hh=None):
+    """
+    mLSTMCell
+    """
+
+    if input.is_cuda:
+        igates = F.linear(input, w_ih)
+        m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+        hgates = F.linear(m, w_hh)
+
+        state = fusedBackend.LSTMFused.apply
+        return state(igates, hgates, hidden[1], b_ih, b_hh)
+
+    hx, cx = hidden
+    
+    m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+    gates = F.linear(input, w_ih, b_ih) + F.linear(m, w_hh, b_hh)
+
+    ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+    ingate = F.sigmoid(ingate)
+    forgetgate = F.sigmoid(forgetgate)
+    cellgate = F.tanh(cellgate)
+    outgate = F.sigmoid(outgate)
+    
+    cy = (forgetgate * cx) + (ingate * cellgate)
+    hy = outgate * F.tanh(cy)
+    
+    return hy, cy
+                                                                            

apex/apex/RNN/models.py

@@ -0,0 +1,54 @@
+import torch
+
+from torch.nn._functions.rnn import LSTMCell, RNNReLUCell, RNNTanhCell, GRUCell
+
+from .RNNBackend import bidirectionalRNN, stackedRNN, RNNCell
+from .cells import mLSTMRNNCell, mLSTMCell
+
+def toRNNBackend(inputRNN, num_layers, bidirectional=False, dropout = 0):
+    """
+    :class:`toRNNBackend`
+    """
+
+    if bidirectional:
+        return bidirectionalRNN(inputRNN, num_layers, dropout = dropout)
+    else:
+        return stackedRNN(inputRNN, num_layers, dropout = dropout)
+
+
+def LSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`LSTM`
+    """
+    inputRNN = RNNCell(4, input_size, hidden_size, LSTMCell, 2, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def GRU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`GRU`
+    """
+    inputRNN = RNNCell(3, input_size, hidden_size, GRUCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def ReLU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`ReLU`
+    """
+    inputRNN = RNNCell(1, input_size, hidden_size, RNNReLUCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def Tanh(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`Tanh`
+    """
+    inputRNN = RNNCell(1, input_size, hidden_size, RNNTanhCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+        
+def mLSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`mLSTM`
+    """
+    inputRNN = mLSTMRNNCell(input_size, hidden_size, bias=bias, output_size=output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+

apex/apex/__init__.py

@@ -0,0 +1,20 @@
+# May help avoid undefined symbol errors https://pytorch.org/cppdocs/notes/faq.html#undefined-symbol-errors-from-pytorch-aten
+import torch
+import warnings
+
+if torch.distributed.is_available():
+    from . import parallel
+
+from . import amp
+from . import fp16_utils
+
+# For optimizers and normalization there is no Python fallback.
+# Absence of cuda backend is a hard error.
+# I would like the errors from importing fused_adam_cuda or fused_layer_norm_cuda
+# to be triggered lazily, because if someone has installed with --cpp_ext and --cuda_ext
+# so they expect those backends to be available, but for some reason they actually aren't
+# available (for example because they built improperly in a way that isn't revealed until
+# load time) the error message is timely and visible.
+from . import optimizers
+from . import normalization
+from . import pyprof

apex/apex/amp/README.md

@@ -0,0 +1,72 @@
+# amp: Automatic Mixed Precision
+
+## Annotating User Functions
+
+Nearly all PyTorch user code needs nothing more than the two steps
+above to use amp. After all, custom layers are built out of simpler
+PyTorch components, and amp already can see those.
+
+However, any custom C++ or CUDA code is outside of amp's (default)
+view of things. For example, suppose I implemented a new recurrent
+cell called a "forgetful recurrent unit" that calls directly into a
+CUDA backend:
+
+```python
+from backend import FRUBackend
+
+def fru(input, hidden, weight, bias):
+    # call to CUDA code
+    FRUBackend(input, hidden, weight, bias)
+```
+
+In this case, it is possible to get a runtime type mismatch. For
+example, you might have `input` in fp16, and `weight` in fp32, and amp
+doesn't have the visibility to insert an appropriate cast.
+
+amp exposes two ways to handle "invisible" backend code: function
+annotations and explicit registration.
+
+#### Function annotation
+
+The first way to handle backend code is a set of function annotations:
+
+- `@amp.half_function`
+- `@amp.float_function`
+- `@amp.promote_function`
+
+These correspond to:
+
+- Cast all arguments to fp16
+- Cast all argumnets fo fp32
+- If there are any type mismatches, cast everything to the widest type
+
+In our example, we believe that the FRU unit is fp16-safe and will get
+performance gains from casting its arguments to fp16, so we write:
+
+```python
+@amp.half_function
+def fru(input, hidden, weight, bias):
+    #...
+```
+
+#### Explicit registration
+
+The other way to handle backend code is with explicit function
+registration:
+
+- `amp.register_half_function(module, function_name)`
+- `amp.register_float_function(module, function_name)`
+- `amp.register_promote_function(module, function_name)`
+
+When using this API, `module` is the containing class or module for
+the function, and `function_name` is the _string_ name of the
+function. Note that the function must be registered before the call to
+`amp.initalize()`.
+
+For our FRU unit, we can register the backend function directly:
+
+```python
+import backend
+
+amp.register_half_function(backend, 'FRUBackend')
+```

apex/apex/amp/__init__.py

@@ -0,0 +1,5 @@
+from .amp import init, half_function, float_function, promote_function,\
+    register_half_function, register_float_function, register_promote_function
+from .handle import scale_loss, disable_casts
+from .frontend import initialize, state_dict, load_state_dict
+from ._amp_state import master_params, _amp_state

apex/apex/amp/__version__.py

@@ -0,0 +1,2 @@
+VERSION = (0, 1, 0)
+__version__ = '.'.join(map(str, VERSION))

apex/apex/amp/_amp_state.py

@@ -0,0 +1,69 @@
+# This is a "header object" that allows different amp modules to communicate.
+# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.
+# But apparently it's ok:
+# http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm
+import os
+import torch
+
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+
+if TORCH_MAJOR == 1 and TORCH_MINOR < 8:
+    from torch._six import container_abcs
+else:
+    import collections.abc as container_abcs
+
+
+class AmpState(object):
+    def __init__(self):
+        self.hard_override=False
+        self.allow_incoming_model_not_fp32 = False
+        self.verbosity=1
+
+
+# Attribute stash.  Could also just stash things as global module attributes.
+_amp_state = AmpState()
+
+
+def warn_or_err(msg):
+    if _amp_state.hard_override:
+        print("Warning:  " + msg)
+    else:
+        raise RuntimeError(msg)
+        # I'm not sure if allowing hard_override is a good idea.
+        # + "  If you're sure you know what you're doing, supply " +
+        #                    "hard_override=True to amp.initialize.")
+
+
+def maybe_print(msg, rank0=False):
+    distributed = torch.distributed.is_available() and \
+        torch.distributed.is_initialized() and \
+        torch.distributed.get_world_size() > 1
+    if _amp_state.verbosity > 0:
+        if rank0:
+            if distributed:
+                if torch.distributed.get_rank() == 0:
+                    print(msg)
+            else:
+                print(msg)
+        else:
+            print(msg)
+
+
+# def iter_params(param_groups):
+#     for group in param_groups:
+#         for p in group['params']:
+#             yield p
+
+
+def master_params(optimizer):
+    """
+    Generator expression that iterates over the params owned by ``optimizer``.
+
+    Args:
+        optimizer: An optimizer previously returned from ``amp.initialize``.
+    """
+    for group in optimizer.param_groups:
+        for p in group['params']:
+            yield p

apex/apex/amp/_initialize.py

@@ -0,0 +1,263 @@
+import torch
+from torch._six import string_classes
+import functools
+import numpy as np
+import sys
+from types import MethodType
+import warnings
+from ._amp_state import _amp_state, warn_or_err, container_abcs
+from .handle import disable_casts
+from .scaler import LossScaler
+from ._process_optimizer import _process_optimizer
+from apex.fp16_utils import convert_network
+from ..fp16_utils import FP16_Optimizer as FP16_Optimizer_general
+from ..contrib.optimizers import FP16_Optimizer as FP16_Optimizer_for_fused
+
+if torch.distributed.is_available():
+    from ..parallel import DistributedDataParallel as apex_DDP
+    from ..parallel.LARC import LARC
+
+
+def to_type(dtype, t):
+    if isinstance(t, torch.Tensor):
+        if not t.is_cuda:
+            # This should not be a hard error, since it may be legitimate.
+            warnings.warn("An input tensor was not cuda.")
+        # GANs require this.
+        # if t.requires_grad:
+        #     warn_or_err("input data requires grad.  Since input data is not a model parameter,\n"
+        #         "its gradients will not be properly allreduced by DDP.")
+        if t.is_floating_point():
+            return t.to(dtype)
+        return t
+    else:
+        # Trust the user's custom batch type, that's all I can do here.
+        return t.to(dtype)
+
+
+# Modified from torch.optim.optimizer.py.  This is a bit more general than casted_args in utils.py.
+def applier(value, fn):
+    if isinstance(value, torch.Tensor):
+        return fn(value)
+    elif isinstance(value, string_classes):
+        return value
+    elif isinstance(value, np.ndarray):
+        return value
+    elif hasattr(value, "to"): # Allow handling of custom batch classes
+        return fn(value)
+    elif isinstance(value, container_abcs.Mapping):
+        return {applier(k, fn) : applier(v, fn) for k, v in value.items()}
+    elif isinstance(value, container_abcs.Iterable):
+        return type(value)(applier(v, fn) for v in value)
+    else:
+        # Do I want this to fire off even if someone chooses to pass something ordinary like
+        # an int or float?  May be more annoying than it's worth.
+        # print("Warning:  unrecognized type in applier.  If your input data is a custom class, "
+        #     "provide it with a .to(dtype) method which converts its floating-point Tensors to dtype. "
+        #     "Amp will check for your custom to() and invoke it to cast the batch's "
+        #     "floating-point Tensors to the appropriate type. "
+        #     "Also, if your data is a custom class, it is your responsibility to ensure that "
+        #     "any Tensors you want to be cuda are already cuda."
+        return value
+
+
+def check_models(models):
+    for model in models:
+        parallel_type = None
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            parallel_type = "torch.nn.parallel.DistributedDataParallel"
+        if ('apex_DDP' in sys.modules) and isinstance(model, apex_DDP):
+            parallel_type = "apex.parallel.DistributedDataParallel"
+        if isinstance(model, torch.nn.parallel.DataParallel):
+            parallel_type = "torch.nn.parallel.DataParallel"
+        if parallel_type is not None:
+            raise RuntimeError("Incoming model is an instance of {}. ".format(parallel_type) +
+                "Parallel wrappers should only be applied to the model(s) AFTER \n"
+                "the model(s) have been returned from amp.initialize.")
+
+
+def check_params_fp32(models):
+    for model in models:
+        for name, param in model.named_parameters():
+            if param.is_floating_point():
+                if 'Half' in param.type():
+                    warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you do not need to call .half() on your model\n"
+                        "before passing it, no matter what optimization level you choose.".format(
+                        name, param.type()))
+                elif not param.is_cuda:
+                    warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you need to provide a model with parameters\n"
+                        "located on a CUDA device before passing it no matter what optimization level\n"
+                        "you chose. Use model.to('cuda') to use the default device.".format(
+                        name, param.type()))
+
+        # Backward compatibility for PyTorch 0.4
+        if hasattr(model, 'named_buffers'):
+            buf_iter = model.named_buffers()
+        else:
+            buf_iter = model._buffers
+        for obj in buf_iter:
+            if type(obj)==tuple:
+                name, buf = obj
+            else:
+                name, buf = obj, buf_iter[obj]
+            if buf.is_floating_point():
+                if 'Half' in buf.type():
+                    warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you do not need to call .half() on your model\n"
+                        "before passing it, no matter what optimization level you choose.".format(
+                        name, buf.type()))
+                elif not buf.is_cuda:
+                    warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you need to provide a model with buffers\n"
+                        "located on a CUDA device before passing it no matter what optimization level\n"
+                        "you chose. Use model.to('cuda') to use the default device.".format(
+                        name, buf.type()))
+
+
+def check_optimizers(optimizers):
+    for optim in optimizers:
+        bad_optim_type = None
+        if isinstance(optim, FP16_Optimizer_general):
+            bad_optim_type = "apex.fp16_utils.FP16_Optimizer"
+        if isinstance(optim, FP16_Optimizer_for_fused):
+            bad_optim_type = "apex.optimizers.FP16_Optimizer"
+        if bad_optim_type is not None:
+            raise RuntimeError("An incoming optimizer is an instance of {}. ".format(bad_optim_type) +
+                               "The optimizer(s) passed to amp.initialize() must be bare \n"
+                               "instances of either ordinary Pytorch optimizers, or Apex fused \n"
+                               "optimizers.\n")
+
+
+class O2StateDictHook(object):
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(self, module, state_dict, prefix, local_metadata):
+        for key in state_dict:
+            param = state_dict[key]
+            if 'Half' in param.type():
+                param = param.to(torch.float32)
+                state_dict[key] = param
+
+
+def _initialize(models, optimizers, properties, num_losses=1, cast_model_outputs=None):
+    from .amp import init as amp_init
+
+    optimizers_was_list = False
+    if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+        optimizers = [optimizers]
+    elif optimizers is None:
+        optimizers = []
+    elif isinstance(optimizers, list):
+        optimizers_was_list = True
+        check_optimizers(optimizers)
+    else:
+        check_optimizers([optimizers])
+        raise TypeError("optimizers must be either a single optimizer or a list of optimizers.")
+
+    if isinstance(models, torch.nn.Module):
+        models_was_list = False
+        models = [models]
+    elif isinstance(models, list):
+        models_was_list = True
+    else:
+        raise TypeError("models must be either a single model or a list of models.")
+
+    check_models(models)
+
+    if not _amp_state.allow_incoming_model_not_fp32:
+        check_params_fp32(models)
+
+    # In the future, when FP16_Optimizer can be deprecated and master weights can
+    # become an attribute, remember to stash master weights before casting the model.
+
+    if properties.cast_model_type:
+        if properties.keep_batchnorm_fp32:
+            for model in models:
+                convert_network(model, properties.cast_model_type)
+        else:
+            for model in models:
+                model.to(properties.cast_model_type)
+
+        input_caster = functools.partial(to_type, properties.cast_model_type)
+        if cast_model_outputs is not None:
+            output_caster = functools.partial(to_type, cast_model_outputs)
+        else:
+            output_caster = functools.partial(to_type, torch.float32)
+
+        for model in models:
+            # Patch the forward method to cast incoming data to the correct type, and
+            # outgoing data to float32, so "the user never needs to call .half()."
+            # I like writing things explicitly more than decorators.
+            def patch_forward(old_fwd):
+                def new_fwd(*args, **kwargs):
+                    output = old_fwd(*applier(args, input_caster),
+                                     **applier(kwargs, input_caster))
+                    return applier(output, output_caster)
+                return new_fwd
+
+            model.forward = patch_forward(model.forward)
+
+        # State dict trick to recast any preexisting per-param state tensors
+        for optimizer in optimizers:
+            optimizer.load_state_dict(optimizer.state_dict())
+
+        # patch model.state_dict() to return float32 params
+        for model in models:
+            for module in model.modules():
+                module._register_state_dict_hook(O2StateDictHook(functools.partial(to_type, torch.float32)))
+
+    elif cast_model_outputs is not None:
+        output_caster = functools.partial(to_type, cast_model_outputs)
+
+        for model in models:
+            def patch_forward(old_fwd):
+                def new_fwd(*args, **kwargs):
+                    output = old_fwd(*args, **kwargs)
+                    return applier(output, output_caster)
+                return new_fwd
+
+            model.forward = patch_forward(model.forward)
+
+    for i, optimizer in enumerate(optimizers):
+        optimizers[i] = _process_optimizer(optimizer, properties)
+
+    _amp_state.loss_scalers = []
+    for _ in range(num_losses):
+        _amp_state.loss_scalers.append(LossScaler(properties.loss_scale,
+                                                  min_loss_scale=_amp_state.min_loss_scale,
+                                                  max_loss_scale=_amp_state.max_loss_scale))
+
+    if properties.patch_torch_functions:
+        # handle is unused here. It's accessible later through a global value anyway.
+        handle = amp_init(loss_scale=properties.loss_scale, verbose=(_amp_state.verbosity == 2))
+        for optimizer in optimizers:
+            # Disable Amp casting for the optimizer step, because it should only be
+            # applied to FP32 master params anyway.
+            def patch_step(old_step):
+                def new_step(self, *args, **kwargs):
+                    with disable_casts():
+                        output = old_step(*args, **kwargs)
+                    return output
+                return new_step
+
+            optimizer.step = MethodType(patch_step(optimizer.step), optimizer)
+
+    if optimizers_was_list:
+        if models_was_list:
+            return models, optimizers
+        else:
+            return models[0], optimizers
+    else:
+        if models_was_list:
+            if len(optimizers) == 0:
+                return models
+            else:
+                return models, optimizers[0]
+        else:
+            if len(optimizers) == 0:
+                return models[0]
+            else:
+                return models[0], optimizers[0]

apex/apex/amp/_process_optimizer.py

@@ -0,0 +1,489 @@
+import types
+from ..fp16_utils import master_params_to_model_params
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import maybe_print
+import torch
+from ..optimizers import FusedSGD
+
+
+class AmpOptimizerState(object):
+    def __init__(self):
+        pass
+
+
+def _master_params_to_model_params(self):
+    stash = self._amp_stash
+    if multi_tensor_applier.available:
+        if len(stash.all_fp16_params) > 0:
+            multi_tensor_applier(
+                stash.multi_tensor_scale,
+                stash.dummy_overflow_buf,
+                [stash.all_fp32_from_fp16_params, stash.all_fp16_params],
+                1.0)
+    else:
+        for fp16_group, fp32_from_fp16_group in zip(stash.fp16_groups, stash.fp32_from_fp16_groups):
+            master_params_to_model_params(fp16_group, fp32_from_fp16_group)
+
+
+def lazy_init_with_master_weights(self):
+        stash = self._amp_stash
+        stash.fp16_groups = []
+        stash.fp32_from_fp16_groups = []
+        stash.fp32_from_fp32_groups = []
+        for i, param_group in enumerate(self.param_groups):
+            # maybe_print("FP16_Optimizer processing param group {}:".format(i))
+            fp16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_fp16_params_this_group = []
+            for i, param in enumerate(param_group['params']):
+                if param.requires_grad:
+                    if param.type() == 'torch.cuda.HalfTensor':
+                        # maybe_print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
+                        #             .format(param.size()))
+                        fp16_params_this_group.append(param)
+                        master_param = param.detach().clone().float()
+                        master_param.requires_grad = True
+                        param_group['params'][i] = master_param
+                        fp32_from_fp16_params_this_group.append(master_param)
+                        # Reset existing state dict key to the new master param.
+                        # We still need to recast per-param state tensors, if any, to FP32.
+                        if param in self.state:
+                           self.state[master_param] = self.state.pop(param)
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        # maybe_print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
+                        #             .format(param.size()))
+                        fp32_params_this_group.append(param)
+                        param_group['params'][i] = param
+                    else:
+                        raise TypeError("Optimizer's parameters must be either "
+                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                        "Received {}".format(param.type()))
+
+            stash.fp16_groups.append(fp16_params_this_group)
+            stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+            stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+        stash.all_fp16_params = []
+        for group in stash.fp16_groups:
+            stash.all_fp16_params += group
+
+        stash.all_fp32_from_fp16_params = []
+        for group in stash.fp32_from_fp16_groups:
+            stash.all_fp32_from_fp16_params += group
+
+        stash.all_fp32_from_fp32_params = []
+        for group in stash.fp32_from_fp32_groups:
+            stash.all_fp32_from_fp32_params += group
+
+        # all_fp16_grad_stash is only needed for fused optimizers.
+        stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+        # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+        stash.all_fp32_from_fp32_grad_stash = [None for _ in stash.all_fp32_from_fp32_params]
+
+        for param in stash.all_fp32_from_fp16_params:
+            param.grad = None
+
+        for param in stash.all_fp32_from_fp32_params:
+            param.grad = None
+
+        # Leverage state_dict() and load_state_dict() to recast preexisting per-param state tensors
+        self.load_state_dict(self.state_dict())
+
+
+def post_backward_models_are_masters(scaler, params, stashed_grads, scale_override=None):
+        grads_have_scale, stashed_have_scale, out_scale = scaler.loss_scale(), 1.0, 1.0
+
+        # not much to do if scale == 1.0 and static scaling
+        if scaler.loss_scale() == 1.0 and not scaler.dynamic:
+            # Clear the stash.
+            for i in range(len(stashed_grads)):
+                stashed_grads[i] = None
+            return
+        
+        if scale_override is not None:
+            grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+        # This is a lot of python overhead...
+        grads_needing_unscale = []
+        grads_needing_unscale_with_stash = []
+        stashed = []
+        for param, stashed_grad in zip(params, stashed_grads):
+            if param.grad is None and stashed_grad is not None:
+                param.grad = stashed_grad
+            elif param.grad is not None and stashed_grad is None:
+                grads_needing_unscale.append(param.grad)
+            elif param.grad is not None and stashed_grad is not None:
+                grads_needing_unscale_with_stash.append(param.grad)
+                stashed.append(stashed_grad)
+            else: # param.grad is None and stashed_grad is None
+                continue
+
+        # unscale() implements grads*(1/scale), so "scale" should be grads_have_scale/out_scale.
+        if len(grads_needing_unscale) > 0:
+            scaler.unscale(
+                grads_needing_unscale,
+                grads_needing_unscale,
+                None, # unused_scale, currently present to avoid API breakage elsewhere
+                models_are_masters=True,
+                scale_override=grads_have_scale/out_scale)
+
+        if len(grads_needing_unscale_with_stash) > 0:
+            scaler.unscale_with_stashed(
+                grads_needing_unscale_with_stash,
+                stashed,
+                grads_needing_unscale_with_stash,
+                scale_override=(grads_have_scale, stashed_have_scale, out_scale))
+
+        # Clear the stash.
+        for i in range(len(stashed_grads)):
+            stashed_grads[i] = None
+
+
+def prepare_backward_with_master_weights(self):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    for i, param in enumerate(stash.all_fp16_params):
+        # Set up to leverage grad copy elision.
+        # This may behave differently from an unpatched optimizer if zero_grad is used and the param is unused.
+        param.grad = None
+
+    # for i, param in enumerate(stash.all_fp32_from_fp16_params):
+    #     stash.all_fp32_from_fp16_grad_stash[i] = param.grad
+
+    for i, param in enumerate(stash.all_fp32_from_fp32_params):
+        stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+
+def post_backward_with_master_weights(self, scaler):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    # This is a lot of python overhead...
+    fp16_grads_needing_unscale = []
+    new_fp32_grads = []
+    fp16_grads_needing_unscale_with_stash = []
+    preexisting_fp32_grads = []
+    for fp16_param, fp32_param in zip(stash.all_fp16_params,
+                                      stash.all_fp32_from_fp16_params):
+        if fp16_param.grad is None and fp32_param.grad is not None:
+            continue
+        elif fp16_param.grad is not None and fp32_param.grad is None:
+            fp32_param.grad = torch.empty_like(fp32_param)
+            fp16_grads_needing_unscale.append(fp16_param.grad)
+            new_fp32_grads.append(fp32_param.grad)
+        elif fp16_param.grad is not None and fp32_param.grad is not None:
+            fp16_grads_needing_unscale_with_stash.append(fp16_param.grad)
+            preexisting_fp32_grads.append(fp32_param.grad)
+        else: # fp16_param.grad is None and fp32_param.grad is None:
+            continue
+
+    if len(fp16_grads_needing_unscale) > 0:
+        scaler.unscale(
+            fp16_grads_needing_unscale,
+            new_fp32_grads,
+            scaler.loss_scale(),
+            models_are_masters=False)
+
+    if len(fp16_grads_needing_unscale_with_stash) > 0:
+        scaler.unscale_with_stashed(
+            fp16_grads_needing_unscale_with_stash,
+            preexisting_fp32_grads,
+            preexisting_fp32_grads)
+
+    # fp32 params can be treated as they would be in the "no_master_weights" case.
+    post_backward_models_are_masters(
+        scaler,
+        stash.all_fp32_from_fp32_params,
+        stash.all_fp32_from_fp32_grad_stash)
+
+
+def lazy_init_no_master_weights(self):
+    stash = self._amp_stash
+    stash.all_fp16_params = []
+    stash.all_fp32_params = []
+    for i, param_group in enumerate(self.param_groups):
+        for i, param in enumerate(param_group['params']):
+            if param.type() == 'torch.cuda.HalfTensor':
+                stash.all_fp16_params.append(param)
+            elif param.type() == 'torch.cuda.FloatTensor':
+                stash.all_fp32_params.append(param)
+            else:
+                raise TypeError("Optimizer's parameters must be either "
+                                "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                "Received {}".format(param.type()))
+
+    stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+    stash.all_fp32_grad_stash = [None for _ in stash.all_fp32_params]
+
+
+def prepare_backward_no_master_weights(self):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    for i, param in enumerate(stash.all_fp16_params):
+        stash.all_fp16_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+    for i, param in enumerate(stash.all_fp32_params):
+        stash.all_fp32_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+
+def post_backward_no_master_weights(self, scaler):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+             (stash.all_fp32_params, stash.all_fp32_grad_stash))
+
+    for params, stashed_grads in split_types:
+        post_backward_models_are_masters(scaler, params, stashed_grads)
+
+
+#####################################################################################
+# FusedSGD versions
+#####################################################################################
+
+# FusedSGD never explicitly materializes the fp32 gradients for "fp32 from fp16" master params
+# outside the kernel, so we must accumulate directly into the model grads.
+def prepare_backward_with_master_weights_FusedSGD(self):
+    if self.materialize_master_grads:
+        prepare_backward_with_master_weights(self)
+    else:
+        stash = self._amp_stash
+
+        self._amp_lazy_init()
+
+        for i, param in enumerate(stash.all_fp16_params):
+            stash.all_fp16_grad_stash[i] = param.grad
+            # Set up to leverage grad copy elision:
+            param.grad = None
+
+        for i, param in enumerate(stash.all_fp32_from_fp32_params):
+            stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+            # Set up to leverage grad copy elision:
+            param.grad = None
+
+
+def post_backward_with_master_weights_FusedSGD(self, scaler):
+    if self.materialize_master_grads:
+        post_backward_with_master_weights(self, scaler)
+    else:
+        stash = self._amp_stash
+
+        self._amp_lazy_init()
+
+        grads_have_scale = scaler.loss_scale()
+        stashed_have_scale = self.most_recent_scale
+        out_scale = grads_have_scale
+        if self.scale_set_by_backward:
+            out_scale = min(grads_have_scale, self.most_recent_scale)
+
+        split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+                 (stash.all_fp32_from_fp32_params, stash.all_fp32_from_fp32_grad_stash))
+
+
+        # unscale_with_stashed() implements grads*1/scale + stashed_grads*1.
+        # stashed_grads are scaled by self.most_recent_scale.
+        for params, stashed_grads in split_types:
+            post_backward_models_are_masters(scaler, params, stashed_grads,
+                                             (grads_have_scale, stashed_have_scale, out_scale))
+
+        self.most_recent_scale = out_scale
+        self.scale_set_by_backward = True
+
+
+def prepare_backward_no_master_weights_FusedSGD(self):
+    prepare_backward_no_master_weights(self)
+
+
+def post_backward_no_master_weights_FusedSGD(self, scaler):
+    post_backward_no_master_weights(self, scaler)
+
+
+def _amp_lazy_init(self):
+    stash = self._amp_stash
+
+    if not stash.lazy_init_called:
+        self._lazy_init_maybe_master_weights()
+        stash.lazy_init_called = True
+
+
+def _process_optimizer(optimizer, properties):
+    if hasattr(optimizer, "_amp_stash"):
+        raise RuntimeError("A given optimizer should only be passed through amp.initialize once.")
+    else:
+        optimizer._amp_stash = AmpOptimizerState()
+
+    optimizer._amp_stash.lazy_init_called = False
+    optimizer._amp_stash.already_patched = False
+    optimizer._amp_stash.params_have_scaled_gradients = False
+
+    for name in ("_lazy_init_maybe_master_weights",
+                 "_master_params_to_model_params",
+                 "_prepare_amp_backward",
+                 "_post_amp_backward",
+                 "_amp_lazy_init"):
+        if hasattr(optimizer, name):
+            raise RuntimeError("Incoming optimizer already has {} defined.".format(name))
+
+    # TODO:  Centralize exposure and import error checking for the C backend.
+    if multi_tensor_applier.available:
+        import amp_C
+        optimizer._amp_stash.multi_tensor_scale = amp_C.multi_tensor_scale
+        optimizer._amp_stash.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
+        optimizer._amp_stash.dummy_overflow_buf = torch.cuda.IntTensor([0]);
+
+    if properties.master_weights:
+        optimizer._lazy_init_maybe_master_weights = types.MethodType(
+            lazy_init_with_master_weights, optimizer)
+
+        optimizer._master_params_to_model_params = types.MethodType(
+            _master_params_to_model_params, optimizer)
+
+        old_step = optimizer.step
+        def new_step(self, closure=None):
+            if closure is not None:
+                raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+            retval = old_step()
+            if not isinstance(self, FusedSGD):
+                self._master_params_to_model_params()
+            # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+            for param in self._amp_stash.all_fp32_from_fp16_params:
+                param.grad = None
+            return retval
+        optimizer.step = types.MethodType(new_step, optimizer)
+
+        old_zero_grad = optimizer.zero_grad
+        def new_zero_grad(self):
+            stash = self._amp_stash
+            self._amp_lazy_init()
+            # Zero the model grads.
+            for param in stash.all_fp16_params:
+                if param.grad is not None:
+                    param.grad.detach_()
+                    param.grad.zero_()
+            for param in stash.all_fp32_from_fp32_params:
+                if param.grad is not None:
+                    param.grad.detach_()
+                    param.grad.zero_()
+            # Clear the master grads that are independent of model grads
+            for param in self._amp_stash.all_fp32_from_fp16_params:
+                param.grad = None
+        optimizer.zero_grad = types.MethodType(new_zero_grad, optimizer)
+
+        if isinstance(optimizer, FusedSGD):
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_with_master_weights_FusedSGD, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_with_master_weights_FusedSGD, optimizer)
+        else:
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_with_master_weights, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_with_master_weights, optimizer)
+    else:
+        optimizer._lazy_init_maybe_master_weights = types.MethodType(
+            lazy_init_no_master_weights, optimizer)
+
+        if isinstance(optimizer, FusedSGD):
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_no_master_weights_FusedSGD, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_no_master_weights_FusedSGD, optimizer)
+        else:
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_no_master_weights, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_no_master_weights, optimizer)
+
+    optimizer._amp_lazy_init = types.MethodType(_amp_lazy_init, optimizer)
+
+    old_add_param_group = optimizer.add_param_group
+
+    def new_add_param_group(self, new_group):
+        stash = self._amp_stash
+
+        if not stash.lazy_init_called:
+            self._lazy_init_maybe_master_weights()
+            stash.lazy_init_called = True
+
+        assert isinstance(new_group, dict), "param group must be a dict"
+
+        new_params = new_group['params']
+        if isinstance(new_params, torch.Tensor):
+            new_group['params'] = [new_params]
+        elif isinstance(new_params, set):
+            raise TypeError('optimizer parameters need to be organized in ordered collections, but '
+                            'the ordering of tensors in sets will change between runs. Please use a list instead.')
+        else:
+            new_group['params'] = list(new_params)
+
+        if properties.master_weights:
+            # Mutate new_group in-place to use FP32 master params
+            fp16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_fp16_params_this_group = []
+            for i, param in enumerate(new_group['params']):
+                if param.requires_grad:
+                    if param.type() == 'torch.cuda.HalfTensor':
+                        fp16_params_this_group.append(param)
+                        master_param = param.detach().clone().float()
+                        master_param.requires_grad = True
+                        new_group['params'][i] = master_param
+                        fp32_from_fp16_params_this_group.append(master_param)
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        fp32_params_this_group.append(param)
+                        new_group['params'][i] = param
+                    else:
+                        raise TypeError("Optimizer's parameters must be either "
+                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                        "Received {}".format(param.type()))
+
+            stash.fp16_groups.append(fp16_params_this_group)
+            stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+            stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+            stash.all_fp16_params += fp16_params_this_group
+            stash.all_fp32_from_fp16_params += fp32_from_fp16_params_this_group
+            stash.all_fp32_from_fp32_params += fp32_params_this_group
+
+            # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+            stash.all_fp32_from_fp32_grad_stash += [None for _ in fp32_params_this_group]
+
+            # It should be ok to let params be added with existing .grad attributes.
+            # for param in fp16_params_this_group:
+            #     param.grad = None
+
+            # for param in fp32_from_fp16_params_this_group:
+            #     param.grad = None
+
+            # for param in stash.fp32_params_this_group:
+            #     param.grad = None
+        else:
+            for param in new_group['params']:
+                if param.type() == 'torch.cuda.HalfTensor':
+                    stash.all_fp16_params.append(param)
+                    stash.all_fp16_grad_stash.append(None)
+                elif param.type() == 'torch.cuda.FloatTensor':
+                    stash.all_fp32_params.append(param)
+                    stash.all_fp32_grad_stash.append(None)
+                else:
+                    raise TypeError("Optimizer's parameters must be either "
+                                    "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                    "Received {}".format(param.type()))
+
+        old_add_param_group(new_group)
+
+    optimizer.add_param_group = types.MethodType(new_add_param_group, optimizer)
+
+    return optimizer

apex/apex/amp/amp.py

@@ -0,0 +1,177 @@
+from . import compat, rnn_compat, utils, wrap
+from .handle import AmpHandle, NoOpHandle
+from .lists import functional_overrides, torch_overrides, tensor_overrides
+from ._amp_state import _amp_state
+from .frontend import *
+
+import functools
+import itertools
+
+import torch
+
+
+_DECORATOR_HANDLE = None
+_USER_CAST_REGISTRY = set()
+_USER_PROMOTE_REGISTRY = set()
+
+
+def _decorator_helper(orig_fn, cast_fn, wrap_fn):
+    def wrapper(*args, **kwargs):
+        handle = _DECORATOR_HANDLE
+        if handle is None or not handle.is_active():
+            return orig_fn(*args, **kwargs)
+        inner_cast_fn = utils.verbosify(cast_fn, orig_fn.__name__,
+                                  handle.verbose)
+        return wrap_fn(orig_fn, inner_cast_fn, handle)(*args, **kwargs)
+    return wrapper
+
+
+# Decorator form
+def half_function(fn):
+    wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=True)
+    return _decorator_helper(fn, utils.maybe_half, wrap_fn)
+
+
+def float_function(fn):
+    wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=False)
+    return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+def promote_function(fn):
+    wrap_fn = functools.partial(wrap.make_promote_wrapper)
+    return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+# Registry form
+def register_half_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_CAST_REGISTRY.add((module, name, utils.maybe_half))
+
+
+def register_float_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_CAST_REGISTRY.add((module, name, utils.maybe_float))
+
+
+def register_promote_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_PROMOTE_REGISTRY.add((module, name))
+
+
+# Top-level function to insert _all_ the hooks.
+def init(enabled=True, loss_scale="dynamic", enable_caching=True, verbose=False, allow_banned=False):
+    global _DECORATOR_HANDLE
+
+    if not enabled:
+        handle = NoOpHandle()
+        _DECORATOR_HANDLE = handle
+        return handle
+
+    handle = AmpHandle(loss_scale, enable_caching, verbose)
+
+    # 0) Force-{fp16, fp32} for user-annotated functions
+    for mod, fn, cast_fn in _USER_CAST_REGISTRY:
+        try_caching = (cast_fn == utils.maybe_half)
+        wrap.cached_cast(mod, fn, cast_fn, handle,
+                         try_caching, verbose)
+    _USER_CAST_REGISTRY.clear()
+
+    # 0.5) Force-promote for user-annotated functions
+    for mod, fn in _USER_PROMOTE_REGISTRY:
+        wrap.promote(mod, fn, handle, verbose)
+    _USER_PROMOTE_REGISTRY.clear()
+
+    # 1) Force-{fp16, fp32} on white- / black-list functions
+    override_modules = [functional_overrides,
+                        torch_overrides,
+                        tensor_overrides]
+    cast_table = [('FP16_FUNCS', utils.maybe_half),
+                  ('FP32_FUNCS', utils.maybe_float)]
+    for module, (list_name, cast_fn) in itertools.product(override_modules,
+                                                          cast_table):
+        for fn in getattr(module, list_name):
+            try_caching = (cast_fn == utils.maybe_half)
+            wrap.cached_cast(module.MODULE, fn, cast_fn, handle,
+                             try_caching, verbose)
+
+    # 1.5) Pre-0.4, put the blacklist methods on HalfTensor and whitelist
+    #      methods on FloatTensor, since they're distinct types.
+    if compat.tensor_is_float_tensor():
+        for fn in tensor_overrides.FP16_FUNCS:
+            wrap.cached_cast(torch.cuda.FloatTensor, fn, utils.maybe_half,
+                             handle, try_caching=True, verbose=verbose)
+        for fn in tensor_overrides.FP32_FUNCS:
+            wrap.cached_cast(torch.cuda.HalfTensor, fn, utils.maybe_float,
+                             handle, try_caching=False, verbose=verbose)
+
+    # 2) Enable type-promotion on multi-arg functions and methods.
+    #    NB: special handling for sequence fns (e.g. `torch.cat`).
+    promote_modules = [torch_overrides, tensor_overrides]
+    promote_table = [('CASTS', wrap.promote),
+                     ('SEQUENCE_CASTS', wrap.sequence_promote)]
+    for promote_mod, (list_name, promote_fn) in itertools.product(promote_modules,
+                                                                  promote_table):
+        for fn in getattr(promote_mod, list_name):
+            promote_fn(promote_mod.MODULE, fn, handle, verbose)
+
+    # 2.5) Pre-0.4, add blacklist methods directly to HalfTensor and FloatTensor types
+    if compat.tensor_is_float_tensor():
+        for cls, (list_name, promote_fn) in itertools.product([torch.cuda.FloatTensor,
+                                                               torch.cuda.HalfTensor],
+                                                              promote_table):
+            for fn in getattr(tensor_overrides, list_name):
+                promote_fn(cls, fn, handle, verbose)
+
+    # 3) For any in-place version of a blacklist function, error if any input is fp16.
+    #    NB: this is overly conservative.
+    for fn in utils.as_inplace(torch_overrides.FP32_FUNCS):
+        wrap.err_if_any_half(torch_overrides.MODULE, fn, handle)
+
+    # 3.5) For any in-place blacklist method, error if called on fp16 tensor
+    for fn in utils.as_inplace(tensor_overrides.FP32_FUNCS):
+        wrap.err_if_arg0_half(tensor_overrides.MODULE, fn, handle, verbose)
+        if compat.tensor_is_float_tensor():
+            wrap.err_if_arg0_half(torch.cuda.HalfTensor, fn, handle, verbose)
+
+    # 4) For other in-place methods, match the type of self tensor
+    for fn in utils.as_inplace(itertools.chain(
+            tensor_overrides.FP16_FUNCS,
+            tensor_overrides.CASTS)):
+        wrap.promote_match_arg0(tensor_overrides.MODULE, fn, handle, verbose)
+        if compat.tensor_is_float_tensor():
+            wrap.promote_match_arg0(torch.cuda.HalfTensor, fn, handle, verbose)
+            wrap.promote_match_arg0(torch.cuda.FloatTensor, fn, handle, verbose)
+
+    # 5) RNNs + RNN cells are whitelisted specially
+    if rnn_compat.has_old_rnns():
+        wrap.rnn_cast(torch.nn.backends.thnn.backend, 'RNN', handle, verbose)
+    if not rnn_compat.has_old_rnns():
+        # Patch in our own indirection of `_VF` in modules/rnn s.t. it is mutable.
+        torch.nn.modules.rnn._VF = rnn_compat.VariableFunctionsShim()
+        # Wrap all the rnns
+        for x in rnn_compat.RNN_NAMES:
+            wrap.new_rnn_cast(x.upper(), handle, verbose)
+
+    # Wrap all the RNN cells
+    rnn_compat.whitelist_rnn_cells(handle, verbose)
+
+    # 6) Place error+print message on banned functions.
+    #    Or, if allow_banned, then cast to FP32.
+    for fn, err_msg in functional_overrides.BANNED_FUNCS:
+        if allow_banned:
+            wrap.cached_cast(functional_overrides.MODULE, fn, utils.maybe_float,
+                             handle, try_caching=True, verbose=verbose)
+        else:
+            wrap.err_if_any_half(functional_overrides.MODULE, fn, handle, err_msg)
+
+    _DECORATOR_HANDLE = handle
+
+    _amp_state.handle = handle
+
+    return handle

apex/apex/amp/compat.py

@@ -0,0 +1,46 @@
+import torch
+
+# True for post-0.4, when Variables/Tensors merged.
+def variable_is_tensor():
+    v = torch.autograd.Variable()
+    return isinstance(v, torch.Tensor)
+
+def tensor_is_variable():
+    x = torch.Tensor()
+    return type(x) == torch.autograd.Variable
+
+# False for post-0.4
+def tensor_is_float_tensor():
+    x = torch.Tensor()
+    return type(x) == torch.FloatTensor
+
+# Akin to `torch.is_tensor`, but returns True for Variable
+# objects in pre-0.4.
+def is_tensor_like(x):
+    return torch.is_tensor(x) or isinstance(x, torch.autograd.Variable)
+
+# Wraps `torch.is_floating_point` if present, otherwise checks
+# the suffix of `x.type()`.
+def is_floating_point(x):
+    if hasattr(torch, 'is_floating_point'):
+        return torch.is_floating_point(x)
+    try:
+        torch_type = x.type()
+        return torch_type.endswith('FloatTensor') or \
+            torch_type.endswith('HalfTensor') or \
+            torch_type.endswith('DoubleTensor')
+    except AttributeError:
+        return False
+
+def scalar_python_val(x):
+    if hasattr(x, 'item'):
+        return x.item()
+    else:
+        if isinstance(x, torch.autograd.Variable):
+            return x.data[0]
+        else:
+            return x[0]
+
+# Accounts for the possibility that some ops may be removed from a namespace.
+def filter_attrs(module, attrs):
+    return list(attrname for attrname in attrs if hasattr(module, attrname))

apex/apex/amp/frontend.py

@@ -0,0 +1,442 @@
+import torch
+from ._initialize import _initialize
+from ._amp_state import _amp_state, warn_or_err, maybe_print
+from collections import OrderedDict
+
+
+class Properties(object):
+    """
+    This class has two purposes: to establish a set of default properties,
+    and to route setting of these attributes through __setattr__ so that (in theory)
+    they can be checked for consistency with other existing args.
+    """
+    def __init__(self):
+        self.options = {
+            "enabled" : False,
+            "opt_level" : None,
+            "cast_model_type" : None,
+            "patch_torch_functions" : False,
+            "keep_batchnorm_fp32" : None,
+            "master_weights" : None,
+            "loss_scale" : 1.0,
+            # Reserved for future functionality
+            # "fused_optimizer" : False,
+            # "enable_ddp_interop" : False,
+            }
+
+    """
+    This function allows updating several options at a time without routing through
+    __setattr__ checks, to avoid "you can't get there from here" scenarios.
+    Currently not intended to be exposed; users are expected to select an opt_level
+    and apply consistent modifications.
+    """
+    def _update_options_dict(self, new_options):
+        for k, v in new_options:
+            if k in self.options:
+                self.options[k] = v
+            else:
+                raise ValueError("Tried to set unexpected option {}".format(k))
+    """
+    The members of "options" are not direct attributes of self, so access attempts
+    will roll down to __getattr__.  This borrows from the logic in torch.nn.Module.
+    """
+    def __getattr__(self, name):
+        if "options" in self.__dict__:
+            options =  self.__dict__["options"]
+            if name in options:
+                return options[name]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, name))
+
+    def __setattr__(self, name, value):
+        if "options" in self.__dict__:
+            if name in self.options:
+                # print("setting {} {}".format(name, value))
+                if name == "cast_model_type":
+                    if self.opt_level == "O1" and value is not None:
+                        if value is not False:
+                            if value is not torch.float32:
+                                warn_or_err("O1 inserts casts around Torch functions rather than "
+                                            "model weights, so with O1, the model weights themselves "
+                                            "should remain FP32. If you wish to cast the model to a "
+                                            "different type, use opt_level='O2' or 'O3'. " +
+                                            "cast_model_type was {}".format(value))
+                    self.options[name] = value
+                elif name == "patch_torch_functions":
+                    if self.opt_level != "O1" and value:
+                        warn_or_err("Currently, patch_torch_functions=True should only be set by "
+                                    "selecting opt_level='O1'.")
+                    self.options[name] = value
+                elif name == "keep_batchnorm_fp32":
+                    if self.opt_level == "O1" and value is not None:
+                        warn_or_err("With opt_level O1, batchnorm functions are automatically patched "
+                                    "to run in FP32, so keep_batchnorm_fp32 should be None." +
+                                    " keep_batchnorm_fp32 was {}".format(value))
+                    if value == "False":
+                        self.options[name] = False
+                    elif value == "True":
+                        self.options[name] = True
+                    else:
+                        assert (value is True or value is False or value is None),\
+                            "keep_batchnorm_fp32 must be a boolean, the string 'True' or 'False', "\
+                            "or None, found keep_batchnorm_fp32={}".format(value)
+                        self.options[name] = value
+                elif name == "master_weights":
+                    if self.opt_level == "O1" and value is not None:
+                        warn_or_err("It doesn't make sense to use master_weights with O1. "
+                                    "With O1, your model weights themselves should be FP32.")
+                    self.options[name] = value
+                elif name == "loss_scale":
+                    if value == "dynamic":
+                        self.options[name] = value
+                    else:
+                        self.options[name] = float(value)
+                else:
+                    self.options[name] = value
+        else:
+            super(Properties, self).__setattr__(name, value)
+
+
+""" O0-O3 are convenience wrappers to establish defaults for typically used mixed precision options. """
+
+class O3:
+    brief = "O3:  Pure FP16 training."
+    more = "Calls .half() on your model, converting the entire model to FP16.\n"\
+        "A casting operation is also inserted to cast incoming Tensors to FP16,\n"\
+        "so you don't need to change your data pipeline.\n"\
+        "This mode is useful for establishing a performance ceiling.\n"\
+        "It's also possible training may 'just work' in this mode.\n"\
+        "If not, try other optimization levels."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O3"
+        properties.cast_model_type = torch.float16
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = False
+        properties.master_weights = False
+        properties.loss_scale = 1.0
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O2:
+    brief = "O2:  FP16 training with FP32 batchnorm and FP32 master weights.\n"
+    more = "Calls .half() on your model, converting the entire model (except for batchnorms)\n"\
+        "to FP16.  Batchnorms are retained in FP32 for additional stability.\n"\
+        "The forward pass is patched to cast incoming Tensors to FP16, so you don't need to change\n"\
+        "your data pipeline.\n"\
+        "O2 creates FP32 master weights outside the model and patches any optimizers to update\n"\
+        "these master weights, then copy the master weights into the FP16 model weights.\n"\
+        "Master weights can also improve convergence and stability."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O2"
+        properties.cast_model_type = torch.float16
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = True
+        properties.master_weights = True
+        properties.loss_scale = "dynamic"
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O1:
+    brief = "O1:  Insert automatic casts around Pytorch functions and Tensor methods.\n"
+    more = "The type of your model's weights is not altered.  However, internally,\n"\
+        "Pytorch functions are patched to cast any Tensor Core-friendly ops to FP16 for speed,\n"\
+        "while operations that might benefit from the additional stability of FP32 are patched\n"\
+        "to cast their inputs to fp32.\n"\
+        "O1 is the safest way to try mixed precision training, and is recommended when\n"\
+        "trying mixed precision training for the first time."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O1"
+        properties.cast_model_type = None
+        properties.patch_torch_functions = True
+        properties.keep_batchnorm_fp32 = None
+        properties.master_weights = None
+        properties.loss_scale = "dynamic"
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O0:
+    brief = "O0:  Pure FP32 training.\n"
+    more = "Your models are checked to make sure parameters are FP32, but otherwise the\n"\
+        "types of weights and internal Pytorch operations are not altered.  This mode disables any\n"\
+        "FP16 arithmetic, although other optimizations like DDP interop may still be requested.\n"
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O0"
+        properties.cast_model_type = torch.float32
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = None
+        properties.master_weights = False
+        properties.loss_scale = 1.0
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+opt_levels = {"O3": O3(),
+              "O2": O2(),
+              "O1": O1(),
+              "O0": O0()}
+
+
+# allow user to directly pass Properties struct as well?
+def initialize(
+    models,
+    optimizers=None,
+    enabled=True,
+    opt_level="O1",
+    cast_model_type=None,
+    patch_torch_functions=None,
+    keep_batchnorm_fp32=None,
+    master_weights=None,
+    loss_scale=None,
+    cast_model_outputs=None,
+    num_losses=1,
+    verbosity=1,
+    min_loss_scale=None,
+    max_loss_scale=2.**24
+    ):
+    """
+    Initialize your models, optimizers, and the Torch tensor and functional namespace according to the
+    chosen ``opt_level`` and overridden properties, if any.
+
+    ``amp.initialize`` should be called **after** you have finished
+    constructing your model(s) and
+    optimizer(s), but **before** you send your model through any DistributedDataParallel wrapper.
+    See `Distributed training`_ in the Imagenet example.
+
+    Currently, ``amp.initialize`` should only be called **once**,
+    although it can process an arbitrary number of
+    models and optimizers (see the corresponding `Advanced Amp Usage topic`_).
+    If you think your use case requires ``amp.initialize`` to be called more than once,
+    `let us know`_.
+
+    Any property keyword argument that is not ``None`` will be interpreted as a manual override.
+
+    To prevent having to rewrite anything else in your script, name the returned models/optimizers
+    to replace the passed models/optimizers, as in the code sample below.
+
+    Args:
+        models (torch.nn.Module or list of torch.nn.Modules):  Models to modify/cast.
+        optimizers (optional, torch.optim.Optimizer or list of torch.optim.Optimizers):  Optimizers to modify/cast.
+            REQUIRED for training, optional for inference.
+        enabled (bool, optional, default=True):  If False, renders all Amp calls no-ops, so your script
+            should run as if Amp were not present.
+        opt_level (str, optional, default="O1"):  Pure or mixed precision optimization level.  Accepted values are
+            "O0", "O1", "O2", and "O3", explained in detail above.
+        cast_model_type (``torch.dtype``, optional, default=None):  Optional property override, see
+            above.
+        patch_torch_functions (bool, optional, default=None):  Optional property override.
+        keep_batchnorm_fp32 (bool or str, optional, default=None):  Optional property override.  If
+            passed as a string, must be the string "True" or "False".
+        master_weights (bool, optional, default=None):  Optional property override.
+        loss_scale (float or str, optional, default=None):  Optional property override.  If passed as a string,
+            must be a string representing a number, e.g., "128.0", or the string "dynamic".
+        cast_model_outputs (torch.dtype, optional, default=None):  Option to ensure that the outputs
+            of your model(s) are always cast to a particular type regardless of ``opt_level``.
+        num_losses (int, optional, default=1):  Option to tell Amp in advance how many losses/backward
+            passes you plan to use.  When used in conjunction with the ``loss_id`` argument to
+            ``amp.scale_loss``, enables Amp to use a different loss scale per loss/backward pass,
+            which can improve stability.  See "Multiple models/optimizers/losses"
+            under `Advanced Amp Usage`_ for examples.  If ``num_losses`` is left to 1, Amp will still
+            support multiple losses/backward passes, but use a single global loss scale
+            for all of them.
+        verbosity (int, default=1):  Set to 0 to suppress Amp-related output.
+        min_loss_scale (float, default=None):  Sets a floor for the loss scale values that can be chosen by dynamic
+            loss scaling.  The default value of None means that no floor is imposed.
+            If dynamic loss scaling is not used, `min_loss_scale` is ignored.
+        max_loss_scale (float, default=2.**24):  Sets a ceiling for the loss scale values that can be chosen by
+            dynamic loss scaling.  If dynamic loss scaling is not used, `max_loss_scale` is ignored.
+
+    Returns:
+        Model(s) and optimizer(s) modified according to the ``opt_level``.
+        If either the ``models`` or ``optimizers`` args were lists, the corresponding return value will
+        also be a list.
+
+    Permissible invocations::
+
+        model, optim = amp.initialize(model, optim,...)
+        model, [optim1, optim2] = amp.initialize(model, [optim1, optim2],...)
+        [model1, model2], optim = amp.initialize([model1, model2], optim,...)
+        [model1, model2], [optim1, optim2] = amp.initialize([model1, model2], [optim1, optim2],...)
+
+        # This is not an exhaustive list of the cross product of options that are possible,
+        # just a set of examples.
+        model, optim = amp.initialize(model, optim, opt_level="O0")
+        model, optim = amp.initialize(model, optim, opt_level="O0", loss_scale="dynamic"|128.0|"128.0")
+
+        model, optim = amp.initialize(model, optim, opt_level="O1") # uses "loss_scale="dynamic" default
+        model, optim = amp.initialize(model, optim, opt_level="O1", loss_scale=128.0|"128.0")
+
+        model, optim = amp.initialize(model, optim, opt_level="O2") # uses "loss_scale="dynamic" default
+        model, optim = amp.initialize(model, optim, opt_level="O2", loss_scale=128.0|"128.0")
+        model, optim = amp.initialize(model, optim, opt_level="O2", keep_batchnorm_fp32=True|False|"True"|"False")
+
+        model, optim = amp.initialize(model, optim, opt_level="O3") # uses loss_scale=1.0 default
+        model, optim = amp.initialize(model, optim, opt_level="O3", loss_scale="dynamic"|128.0|"128.0")
+        model, optim = amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True|False|"True"|"False")
+
+    The `Imagenet example`_ demonstrates live use of various opt_levels and overrides.
+
+    .. _`Distributed training`:
+        https://github.com/NVIDIA/apex/tree/master/examples/imagenet#distributed-training
+
+    .. _`Imagenet example`:
+        https://github.com/NVIDIA/apex/tree/master/examples/imagenet
+
+    .. _`Advanced Amp Usage`:
+        https://nvidia.github.io/apex/advanced.html
+
+    .. _`Advanced Amp Usage topic`:
+        https://nvidia.github.io/apex/advanced.html#multiple-models-optimizers-losses
+
+    .. _`let us know`:
+        https://github.com/NVIDIA/apex/issues
+    """
+    _amp_state.opt_properties = Properties()
+    _amp_state.verbosity = verbosity
+
+    if not enabled:
+        if optimizers is None:
+            return models
+        else:
+            return models, optimizers
+
+    if not torch.backends.cudnn.enabled:
+        raise RuntimeError(
+            "Amp requires torch.backends.cudnn.enabled = True")
+
+    if opt_level not in opt_levels:
+        raise RuntimeError(
+            "Unexpected optimization level {}. ".format(opt_level) +
+            "Options are 'O0', 'O1', 'O2', 'O3'.  Note that in `O0`, `O1`, etc., the prefix O is the letter O, " +
+            "not the number zero.")
+    else:
+        _amp_state.opt_properties = opt_levels[opt_level](_amp_state.opt_properties)
+        maybe_print("Selected optimization level {}".format(opt_levels[opt_level].brief), True)
+        maybe_print("Defaults for this optimization level are:", True)
+        for k, v in _amp_state.opt_properties.options.items():
+            maybe_print("{:22} : {}".format(k, v), True)
+
+    _amp_state.min_loss_scale = min_loss_scale
+    _amp_state.max_loss_scale = max_loss_scale
+
+    maybe_print("Processing user overrides (additional kwargs that are not None)...", True)
+    # I chose to have the keyword arguments listed directly in the argument list,
+    # instead of **kwargs, so I can't use kwargs.items() here.
+    if enabled is not None:
+        _amp_state.opt_properties.enabled = enabled
+    if opt_level is not None:
+        _amp_state.opt_properties.opt_level = opt_level
+    if cast_model_type is not None:
+        _amp_state.opt_properties.cast_model_type = cast_model_type
+    if patch_torch_functions is not None:
+        _amp_state.opt_properties.patch_torch_functions = patch_torch_functions
+    if keep_batchnorm_fp32 is not None:
+        _amp_state.opt_properties.keep_batchnorm_fp32 = keep_batchnorm_fp32
+    if master_weights is not None:
+        _amp_state.opt_properties.master_weights = master_weights
+    if loss_scale is not None:
+        _amp_state.opt_properties.loss_scale = loss_scale
+
+    maybe_print("After processing overrides, optimization options are:", True)
+    for k, v in _amp_state.opt_properties.options.items():
+        maybe_print("{:22} : {}".format(k, v), True)
+
+    return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)
+
+
+def state_dict(destination=None):
+    if destination is None:
+        destination = OrderedDict()
+
+    for idx, loss_scaler in enumerate(_amp_state.loss_scalers):
+        destination['loss_scaler%d' % idx] = {
+            'loss_scale': loss_scaler.loss_scale(),
+            'unskipped': loss_scaler._unskipped,
+        }
+    return destination
+
+
+def load_state_dict(state_dict):
+    # Check if state_dict containes the same number of loss_scalers as current setup
+    if len(state_dict) != len(_amp_state.loss_scalers):
+        print('Warning: state_dict contains {} entries, while {} loss_scalers are used'.format(
+            len(state_dict), len(_amp_state.loss_scalers)))
+
+    state_dict = state_dict.copy()
+    
+    nb_loss_scalers = len(_amp_state.loss_scalers)
+    unexpected_keys = []
+    # Initialize idx outside, since unexpected_keys will increase it if enumerate is used
+    idx = 0
+    for key in state_dict:
+        if 'loss_scaler' not in key:
+            unexpected_keys.append(key)
+        else:
+            if idx > (nb_loss_scalers - 1):
+                print('Skipping loss_scaler[{}], since num_losses was set to {}'.format(
+                    idx, nb_loss_scalers))
+                break
+            _amp_state.loss_scalers[idx]._loss_scale = state_dict[key]['loss_scale']
+            _amp_state.loss_scalers[idx]._unskipped = state_dict[key]['unskipped']
+            idx += 1
+
+    if len(unexpected_keys) > 0:
+        raise RuntimeError(
+            'Error(s) in loading state_dict. Unexpected key(s) in state_dict: {}. '.format(
+                ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+
+
+# TODO:  is this necessary/useful?
+# def check_option_consistency(enabled=True,
+#                              opt_level=None,
+#                              cast_model_type=None,
+#                              patch_torch_functions=None,
+#                              keep_batchnorm_fp32=None,
+#                              master_weights=None,
+#                              loss_scale=None,
+#                              enable_ddp_interop=None,
+#                              hard_override=False):
+#     """
+#     Utility function that enables users to quickly check if the option combination they intend
+#     to use is permitted.  ``check_option_consistency`` does not require models or optimizers
+#     to be constructed, and can be called at any point in the script.  ``check_option_consistency``
+#     is totally self-contained; it does not set any amp global state or affect anything outside
+#     of itself.
+#     """
+#
+#     if not enabled:
+#         return
+#
+#     if opt_level not in opt_levels:
+#         raise RuntimeError("Unexpected optimization level.  Options are 'O0', 'O1', 'O2', 'O3'.")
+#     else:
+#         opt_properties = opt_levels[opt_level](Properties())
+#         print("Selected optimization level {}", opt_levels[opt_level].brief)
+#         print("Defaults for this optimization level are:")
+#         for k, v in opt_properties.options:
+#             print("{:22} : {}".format(k, v))
+#
+#     print("Processing user overrides (additional kwargs that are not None)...")
+#     for k, v in kwargs:
+#         if k not in _amp_state.opt_properties.options:
+#             raise RuntimeError("Unexpected kwarg {}".format(k))
+#         if v is not None:
+#             setattr(opt_properties, k, v)
+#
+#     print("After processing overrides, optimization options are:")
+#     for k, v in opt_properties.options:
+#         print("{:22} : {}".format(k, v))

apex/apex/amp/handle.py

@@ -0,0 +1,281 @@
+import contextlib
+import warnings
+import sys
+import torch
+
+from . import utils
+from .opt import OptimWrapper
+from .scaler import LossScaler
+from ._amp_state import _amp_state, master_params, maybe_print
+
+if torch.distributed.is_available():
+    from ..parallel.LARC import LARC
+
+
+# There's no reason to expose the notion of a "handle". Everything can happen through amp.* calls.
+@contextlib.contextmanager
+def scale_loss(loss,
+               optimizers,
+               loss_id=0,
+               model=None,
+               delay_unscale=False,
+               delay_overflow_check=False):
+    """
+    On context manager entrance, creates ``scaled_loss = (loss.float())*current loss scale``.
+    ``scaled_loss`` is yielded so that the user can call ``scaled_loss.backward()``::
+
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+
+    On context manager exit (if ``delay_unscale=False``), the gradients are checked for infs/NaNs
+    and unscaled, so that ``optimizer.step()`` can be called.
+
+    .. note::
+        If Amp is using explicit FP32 master params (which is the default for ``opt_level=O2``, and
+        can also be manually enabled by supplying ``master_weights=True`` to ``amp.initialize``)
+        any FP16 gradients are copied to FP32 master gradients before being unscaled.
+        ``optimizer.step()`` will then apply the unscaled master gradients to the master params.
+
+    .. warning::
+        If Amp is using explicit FP32 master params, only the FP32 master gradients will be
+        unscaled.  The direct ``.grad`` attributes of any FP16
+        model params will remain scaled after context manager exit.
+        This subtlety affects gradient clipping.  See "Gradient clipping" under
+        `Advanced Amp Usage`_ for best practices.
+
+    Args:
+        loss(Tensor):  Typically a scalar Tensor. The ``scaled_loss`` that the context
+            manager yields is simply ``loss.float()*loss_scale``, so in principle
+            ``loss`` could have more than one element, as long as you call
+            ``backward()`` on ``scaled_loss`` appropriately within the context manager body.
+        optimizers:  All optimizer(s) for which the current backward pass is creating gradients.
+            Must be an optimizer or list of optimizers returned from an earlier call
+            to ``amp.initialize``.  For example use with multiple optimizers, see
+            "Multiple models/optimizers/losses" under `Advanced Amp Usage`_.
+        loss_id(int, optional, default=0):  When used in conjunction with the ``num_losses`` argument
+            to ``amp.initialize``, enables Amp to use a different loss scale per loss.  ``loss_id``
+            must be an integer between 0 and ``num_losses`` that tells Amp which loss is
+            being used for the current backward pass.  See "Multiple models/optimizers/losses"
+            under `Advanced Amp Usage`_ for examples.  If ``loss_id`` is left unspecified, Amp
+            will use the default global loss scaler for this backward pass.
+        model(torch.nn.Module, optional, default=None):  Currently unused, reserved to enable future
+            optimizations.
+        delay_unscale(bool, optional, default=False):  ``delay_unscale`` is never necessary, and
+            the default value of ``False`` is strongly recommended.
+            If ``True``, Amp will not unscale the gradients or perform model->master
+            gradient copies on context manager exit.
+            ``delay_unscale=True`` is a minor ninja performance optimization and can result
+            in weird gotchas (especially with multiple models/optimizers/losses),
+            so only use it if you know what you're doing.
+            "Gradient accumulation across iterations" under `Advanced Amp Usage`_
+            illustrates a situation where this CAN (but does not need to) be used.
+
+    .. warning::
+        If ``delay_unscale`` is ``True`` for a given backward pass, ``optimizer.step()`` cannot be
+        called yet after context manager exit, and must wait for another, later backward context
+        manager invocation with ``delay_unscale`` left to False.
+
+    .. _`Advanced Amp Usage`:
+        https://nvidia.github.io/apex/advanced.html
+    """
+    if not hasattr(_amp_state, "opt_properties"):
+        raise RuntimeError("Invoked 'with amp.scale_loss`, but internal Amp state has not been initialized.  "
+                           "model, optimizer = amp.initialize(model, optimizer, opt_level=...) must be called "
+                           "before `with amp.scale_loss`.")
+
+    if not _amp_state.opt_properties.enabled:
+        yield loss
+        return
+
+    if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+        optimizers = [optimizers]
+
+    loss_scaler = _amp_state.loss_scalers[loss_id]
+    loss_scale = loss_scaler.loss_scale()
+
+    if ((not _amp_state.opt_properties.master_weights)
+        and (not loss_scaler.dynamic)
+        and loss_scale == 1.0):
+        yield loss.float()
+        # Needing to drop the cache here as well is an ugly gotcha.
+        # But for now I think it's necessary to short-circuit.
+        # Probably ok to skip this if not delay_unscale
+        if _amp_state.opt_properties.patch_torch_functions:
+            _amp_state.handle._clear_cache()
+        return
+
+    if not delay_unscale:
+        if isinstance(optimizers, list):
+            for optimizer in optimizers:
+                if not optimizer._amp_stash.params_have_scaled_gradients:
+                    optimizer._prepare_amp_backward()
+
+    yield (loss.float())*loss_scale
+
+    if delay_unscale:
+        for optimizer in optimizers:
+            optimizer._amp_stash.params_have_scaled_gradients = True
+    else:
+        # FusedSGD may take care of unscaling as part of their step() methods.
+        # if not isinstance(optimizers, FP16_Optimizer_for_fused):
+            loss_scaler.clear_overflow_state()
+            for optimizer in optimizers:
+                optimizer._post_amp_backward(loss_scaler)
+                optimizer._amp_stash.params_have_scaled_gradients = False
+            # For future fused optimizers that enable sync-free dynamic loss scaling,
+            # should_skip will always be False.
+            should_skip = False if delay_overflow_check else loss_scaler.update_scale()
+            if should_skip:
+                for optimizer in optimizers:
+                    if not optimizer._amp_stash.already_patched:
+                        # Close on loss_scaler and loss_id as well, to be safe.  Probably not
+                        # necessary because amp.scale_loss is already creating a temporary scope.
+                        def patch_step(opt, loss_scaler, loss_id):
+                            opt_step = opt.step
+                            def skip_step(closure=None):
+                                if closure is not None:
+                                    raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+                                maybe_print(("Gradient overflow.  Skipping step, loss scaler " +
+                                             "{} reducing loss scale to {}").format(loss_id,
+                                             loss_scaler.loss_scale()))
+                                # TODO:  I don't like the special casing for different optimizer implementations.
+                                # Maybe skip should delegate to a method owned by the optimizers themselves.
+                                if hasattr(opt._amp_stash, "all_fp32_from_fp16_params"):
+                                    # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+                                    for param in opt._amp_stash.all_fp32_from_fp16_params:
+                                        param.grad = None
+                                if hasattr(opt, "most_recent_scale"):
+                                    opt.most_recent_scale = 1.0
+                                    opt.scale_set_by_backward = False
+                                opt.step = opt_step
+                                opt._amp_stash.already_patched = False
+                            return skip_step
+                        optimizer.step = patch_step(optimizer, loss_scaler, loss_id)
+                        optimizer._amp_stash.already_patched = True
+
+    # Probably ok to skip this if not delay_unscale
+    if _amp_state.opt_properties.patch_torch_functions:
+        _amp_state.handle._clear_cache()
+
+
+# Free function version of AmpHandle.disable_casts, another step on the
+# path to removing the concept of "AmpHandle"
+@contextlib.contextmanager
+def disable_casts():
+    _amp_state.handle._is_active = False
+    yield
+    _amp_state.handle._is_active = True
+
+
+class AmpHandle(object):
+    def __init__(self, loss_scale="dynamic", enable_caching=True, verbose=False):
+        self._enable_caching = enable_caching
+        self._verbose = verbose
+        self._cache = dict()
+        self._default_scaler = LossScaler(loss_scale)
+        self._is_active = True
+        self._all_wrappers = []
+
+    def is_active(self):
+        return self._is_active
+
+    @contextlib.contextmanager
+    def _disable_casts(self):
+        self._is_active = False
+        yield
+        self._is_active = True
+
+    def wrap_optimizer(self, optimizer, num_loss=1):
+        self._default_scaler = None
+        return OptimWrapper(optimizer, self, num_loss)
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss, optimizer):
+        raise RuntimeError("The old Amp API is no longer supported.  Please move to the new API, "
+            "documented here:  https://nvidia.github.io/apex/amp.html.  Transition guide:  "
+            "https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users")
+
+        if not self.is_active():
+            yield loss
+            return
+
+        if self._default_scaler is None:
+            raise RuntimeError(
+                'After calling `handle.wrap_optimizer()`, you must explicitly ' +
+                'use `optimizer.scale_loss(loss)`.')
+
+        # TODO: this code block is duplicated here and `opt.py`. Unify.
+        loss_scale = self._default_scaler.loss_scale()
+        yield loss * loss_scale
+
+        self._default_scaler.clear_overflow_state()
+        self._default_scaler.unscale(
+            master_params(optimizer),
+            master_params(optimizer),
+            loss_scale)
+        should_skip = self._default_scaler.update_scale()
+        if should_skip:
+            optimizer_step = optimizer.step
+            def skip_step():
+                maybe_print('Gradient overflow, skipping update')
+                optimizer.step = optimizer_step
+            optimizer.step = skip_step
+
+        self._clear_cache()
+
+    def _clear_cache(self):
+        self._cache.clear()
+
+    # Experimental support for saving / restoring uncasted versions of functions
+    def _save_func(self, mod, fn, func):
+        self._all_wrappers.append((mod, fn, func))
+
+    def _deactivate(self):
+        for mod, fn, func in self._all_wrappers:
+            utils.set_func(mod, fn, func)
+        self._all_wrappers = []
+
+    @property
+    def has_cache(self):
+        return self._enable_caching
+
+    @property
+    def cache(self):
+        return self._cache
+
+    def remove_cache(self, param):
+        if self.has_cache and param in self.cache:
+            del self.cache[param]
+
+    @property
+    def verbose(self):
+        return self._verbose
+
+class NoOpHandle(object):
+    def is_active(self):
+        return False
+
+    @contextlib.contextmanager
+    def _disable_casts(self):
+        yield
+
+    def wrap_optimizer(self, optimizer, num_loss=1):
+        return OptimWrapper(optimizer, self, num_loss)
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss, optimizer):
+        yield loss
+
+    @property
+    def has_cache(self):
+        return False
+
+    @property
+    def verbose(self):
+        return False
+
+    def _clear_cache(self):
+        pass
+
+    def _deactivate(self):
+        pass

apex/apex/amp/lists/functional_overrides.py

@@ -0,0 +1,80 @@
+
+# TODO: think about the following two. They do weird things.
+# - torch.nn.utils.clip_grad (but it should always be fp32 anyway)
+# - torch.nn.utils.weight_norm
+
+# Notes:
+# F.instance_norm uses batch_norm internally. Which correctly handles
+#   fp16 in/out with fp32 weights. So we shouldn't do anything for
+#   either of these.
+# F.normalize calls `input.norm()` internally, so it's redundant, but
+#   kept here in case impl. changes.
+# F.cosine_similarity is same: calls `x.norm()` internally.
+
+import torch.nn.functional
+
+MODULE = torch.nn.functional
+
+FP16_FUNCS = [
+    'conv1d',
+    'conv2d',
+    'conv3d',
+    'conv_transpose1d',
+    'conv_transpose2d',
+    'conv_transpose3d',
+    'conv_tbc', # Undocumented / maybe new?
+    'linear',
+]
+
+FP32_FUNCS = [
+
+    # Interpolation/Upsampling TODO:  Remove for 1.2
+    'interpolate',
+    'grid_sample',
+
+    # Pointwise
+    'softplus',
+    'softmin',
+    'log_softmax',
+    'softmax',
+    'gelu',
+    
+    # Normalization
+    'layer_norm',
+    'group_norm',
+    'local_response_norm',
+    'normalize',
+    'cosine_similarity',
+
+    # Loss functions
+    # TODO: which of these can be fp16?
+    'poisson_nll_loss',
+    'cosine_embedding_loss',
+    'cross_entropy',
+    'hinge_embedding_loss',
+    'kl_div',
+    'l1_loss',
+    'mse_loss',
+    'margin_ranking_loss',
+    'multilabel_margin_loss',
+    'multilabel_soft_margin_loss',
+    'multi_margin_loss',
+    'nll_loss',
+    'binary_cross_entropy_with_logits',
+    'smooth_l1_loss',
+    'soft_margin_loss',
+    'triplet_margin_loss',
+    'ctc_loss'
+]
+
+BANNED_FUNCS = [
+    ('binary_cross_entropy',
+     ("\namp does not work out-of-the-box with `F.binary_cross_entropy` or `torch.nn.BCELoss.` "
+      "It requires that the output of the previous function be already a FloatTensor. \n\n"
+      "Most models have a Sigmoid right before BCELoss. In that case, you can use\n"
+      "    torch.nn.BCEWithLogitsLoss\nto combine Sigmoid+BCELoss into a single layer "
+      "that is compatible with amp.\nAnother option is to add\n"
+      "    amp.register_float_function(torch, 'sigmoid')\nbefore calling `amp.init()`.\n"
+      "If you _really_ know what you are doing, you can disable this warning by passing "
+      "allow_banned=True to `amp.init()`."))
+]

apex/apex/amp/lists/tensor_overrides.py

@@ -0,0 +1,63 @@
+from .. import compat
+from . import torch_overrides
+
+import importlib
+
+import torch
+
+# if compat.variable_is_tensor() and not compat.tensor_is_variable():
+MODULE = torch.Tensor
+# else:
+#     MODULE = torch.autograd.Variable
+
+
+FP16_FUNCS = compat.filter_attrs(MODULE, [
+    '__matmul__',
+])
+
+FP32_FUNCS = compat.filter_attrs(MODULE, [
+    '__ipow__',
+    '__pow__',
+    '__rpow__',
+
+    # Cast to fp32 before transfer to CPU
+    'cpu',
+])
+
+CASTS = compat.filter_attrs(MODULE, [
+    '__add__',
+    '__div__',
+    '__eq__',
+    '__ge__',
+    '__gt__',
+    '__iadd__',
+    '__idiv__',
+    '__imul__',
+    '__isub__',
+    '__itruediv__',
+    '__le__',
+    '__lt__',
+    '__mul__',
+    '__ne__',
+    '__radd__',
+    '__rdiv__',
+    '__rmul__',
+    '__rsub__',
+    '__rtruediv__',
+    '__sub__',
+    '__truediv__',
+])
+
+# None of these, but here to make code cleaner.
+SEQUENCE_CASTS = []
+
+# We need to grab all the methods from torch_overrides and add them to
+# the Tensor lists as well, as almost all methods are duplicated
+# between `torch` and `torch.Tensor` (and check with `hasattr`,
+# because a few random ones aren't defined on Tensor)
+_self_mod = importlib.import_module(__name__)
+for attrname in ['FP16_FUNCS', 'FP32_FUNCS', 'CASTS', 'SEQUENCE_CASTS']:
+    lst = getattr(_self_mod, attrname)
+    for fn in getattr(torch_overrides, attrname):
+        if hasattr(MODULE, fn):
+            lst.append(fn)

apex/apex/amp/lists/torch_overrides.py

@@ -0,0 +1,115 @@
+import torch
+
+from .. import utils
+
+MODULE = torch
+
+FP16_FUNCS = [
+    # Low level functions wrapped by torch.nn layers.
+    # The wrapper layers contain the weights which are then passed in as a parameter
+    # to these functions.
+    'conv1d',
+    'conv2d',
+    'conv3d',
+    'conv_transpose1d',
+    'conv_transpose2d',
+    'conv_transpose3d',
+    'conv_tbc',
+    'prelu',
+
+    # BLAS
+    'addmm',
+    'addmv',
+    'addr',
+    'matmul',
+    'mm',
+    'mv',
+]
+
+FP32_FUNCS = [
+    # Pointwise
+    'acos',
+    'asin',
+    'cosh',
+    'erfinv',
+    'exp',
+    'expm1',
+    'log',
+    'log10',
+    'log2',
+    'reciprocal',
+    'rsqrt',
+    'sinh',
+    'tan',
+
+    # Other math
+    'pow',
+
+    # Reduction
+    'cumprod',
+    'cumsum',
+    'dist',
+    # 'mean',
+    'norm',
+    'prod',
+    'std',
+    'sum',
+    'var',
+
+    # Misc
+    'renorm'
+]
+
+version_strings = torch.__version__.split('.')
+version_major = version_strings[0]
+version_minor = version_strings[1]
+version_num = float(version_major + "." + version_minor)
+# Before torch 1.1, mean must be blacklisted.
+if version_num < 1.1:
+    FP32_FUNCS.append('mean')
+
+# Before CUDA 9.1, batched matmul was missing fast FP16 kernels. We
+# check the CUDA version -- if at least 9.1, then put the bmm
+# functions on the fp16 list. Otherwise, put them on the fp32 list.
+_bmms = ['addbmm',
+         'baddbmm',
+         'bmm']
+
+if utils.is_cuda_enabled():
+  # workaround https://github.com/facebookresearch/maskrcnn-benchmark/issues/802
+  if utils.get_cuda_version() >= (9, 1, 0):
+      FP16_FUNCS.extend(_bmms)
+  else:
+      FP32_FUNCS.extend(_bmms)
+
+# Multi-tensor fns that may need type promotion
+CASTS = [
+    # Multi-tensor math
+    'addcdiv',
+    'addcmul',
+    'atan2',
+    'cross',
+    'bilinear',
+    'dot',
+
+    # Element-wise _or_ tensor-wise math
+    'add',
+    'div',
+    'mul',
+
+    # Comparison
+    'eq',
+    'equal',
+    'ge',
+    'gt',
+    'le',
+    'lt',
+    'ne'
+]
+
+# Functions that take sequence arguments. We need to inspect the whole
+# sequence and cast to the widest type.
+SEQUENCE_CASTS = [
+    'cat',
+    'stack'
+]

apex/apex/amp/opt.py

@@ -0,0 +1,103 @@
+import contextlib
+import warnings
+
+from .scaler import LossScaler, master_params
+from ._amp_state import maybe_print
+
+import numpy as np
+
+class OptimWrapper(object):
+    def __init__(self, optimizer, amp_handle, num_loss):
+        self._optimizer = optimizer
+        self._amp_handle = amp_handle
+        self._num_loss = num_loss
+        self._loss_idx = 0
+        self._skip_next = [False] * num_loss
+        self._loss_scaler = [LossScaler('dynamic') for _ in range(num_loss)]
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss):
+        if not self._amp_handle.is_active():
+            yield loss
+            return
+
+        # When there are multiple losses per-optimizer, we need
+        # to save out current grad accumulation, since we won't be
+        # able to unscale this particulare loss once the grads are
+        # all mixed together.
+        cached_grads = []
+        if self._loss_idx > 0:
+            for p in master_params(self._optimizer):
+                if p.grad is not None:
+                    cached_grads.append(p.grad.data.detach().clone())
+                else:
+                    cached_grads.append(None)
+            self._optimizer.zero_grad()
+
+        loss_scale = self._cur_loss_scaler().loss_scale()
+        yield loss * loss_scale
+
+        self._cur_loss_scaler().clear_overflow_state()
+        self._cur_loss_scaler().unscale(
+            master_params(self._optimizer),
+            master_params(self._optimizer),
+            loss_scale)
+        self._skip_next[self._loss_idx] = self._cur_loss_scaler().update_scale()
+        self._loss_idx += 1
+
+        if len(cached_grads) > 0:
+            for p, cached_grad in zip(master_params(self._optimizer),
+                                      cached_grads):
+                if cached_grad is not None:
+                    p.grad.data.add_(cached_grad)
+            cached_grads = []
+
+    def _cur_loss_scaler(self):
+        assert 0 <= self._loss_idx < self._num_loss
+        return self._loss_scaler[self._loss_idx]
+
+    def step(self, closure=None):
+        if not self._amp_handle.is_active():
+            return self._optimizer.step(closure=closure)
+
+        self._loss_idx = 0
+
+        for group in self._optimizer.param_groups:
+            for p in group['params']:
+                self._amp_handle.remove_cache(p)
+
+        if closure is not None:
+            raise NotImplementedError(
+                'The `closure` argument is unsupported by the amp ' +
+                'optimizer wrapper.')
+        if any(self._skip_next):
+            maybe_print('Gradient overflow, skipping update')
+            self._skip_next = [False] * self._num_loss
+        else:
+            return self._optimizer.step(closure=closure)
+
+    # Forward any attribute lookups
+    def __getattr__(self, attr):
+        return getattr(self._optimizer, attr)
+
+    # Forward all torch.optim.Optimizer methods
+    def __getstate__(self):
+        return self._optimizer.__getstate__()
+
+    def __setstate__(self):
+        return self._optimizer.__setstate__()
+
+    def __repr__(self):
+        return self._optimizer.__repr__()
+
+    def state_dict(self):
+        return self._optimizer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        return self._optimizer.load_state_dict(state_dict)
+
+    def zero_grad(self):
+        return self._optimizer.zero_grad()
+
+    def add_param_group(self, param_group):
+        return self._optimizer.add_param_group(param_group)

apex/apex/amp/rnn_compat.py

@@ -0,0 +1,53 @@
+from . import utils, wrap
+
+import torch
+_VF = torch._C._VariableFunctions
+RNN_NAMES = ['rnn_relu', 'rnn_tanh', 'gru', 'lstm']
+
+def _gen_VF_wrapper(name):
+    def wrapper(*args, **kwargs):
+        return getattr(_VF, name)(*args, **kwargs)
+    return wrapper
+
+# Some python magic to generate an object that has the rnn cell functions
+# defined on it, all of which call into corresponding _VF version.
+# Intended to patch torch.nn.modules.rnn._VF (aka, the ref named "_VF"
+# imported at module scope within torch.nn.modules.rnn).  This should
+# not affect third-party importers of _VF.py.
+class VariableFunctionsShim(object):
+    def __init__(self):
+        for name in RNN_NAMES:
+            for suffix in ['', '_cell']:
+               fn_name = name + suffix
+               setattr(self, fn_name, _gen_VF_wrapper(fn_name))
+
+def has_old_rnns():
+    try:
+        torch.nn.backends.thnn.backend.LSTMCell
+        return True
+    except:
+        return False
+
+def whitelist_rnn_cells(handle, verbose):
+    # Different module + function names in old/new RNN cases
+    if has_old_rnns():
+        fn_names = ['RNNReLUCell', 'RNNTanhCell', 'LSTMCell', 'GRUCell']
+        mod = torch.nn.backends.thnn.backend
+    else:
+        fn_names = [x + '_cell' for x in RNN_NAMES]
+        mod = torch.nn.modules.rnn._VF
+        assert isinstance(mod, VariableFunctionsShim)
+
+    # Insert casts on cell functions
+    for fn in fn_names:
+        wrap.cached_cast(mod, fn, utils.maybe_half, handle,
+                         try_caching=True, verbose=verbose)
+
+    if has_old_rnns():
+        # Special handling of `backward` for fused gru / lstm:
+        # The `backward` method calls Tensor.sum() (blacklist) internally,
+        # and then the resulting grad_input has the wrong type.
+        # TODO: where else is this a problem?
+        for rnn_type in ['GRUFused', 'LSTMFused']:
+            mod = getattr(torch.nn._functions.thnn.rnnFusedPointwise, rnn_type)
+            wrap.disable_casts(mod, 'backward', handle)

apex/apex/amp/scaler.py

@@ -0,0 +1,217 @@
+import torch
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import _amp_state, master_params, maybe_print
+from itertools import product
+
+def scale_check_overflow_python(model_grad, master_grad, scale, check_overflow=False):
+    # Exception handling for 18.04 compatibility
+    if check_overflow:
+        cpu_sum = float(model_grad.float().sum())
+        if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+            return True
+
+    if master_grad is not model_grad: # copy_ probably internally short-circuits this
+        master_grad.copy_(model_grad)
+    if scale != 1.0:
+        master_grad.mul_(scale)
+    return False
+
+def axpby_check_overflow_python(model_grad, stashed_grad, master_grad, a, b, check_overflow=False):
+    # Exception handling for 18.04 compatibility
+    if check_overflow:
+        cpu_sum = float(model_grad.float().sum())
+        if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+            return True
+
+    # if master_grad is not model_grad: # copy_ probably internally short-circuits this
+    #     master_grad.copy_(model_grad)
+    assert stashed_grad.dtype == master_grad.dtype
+    converted_model_grad = model_grad.data.to(master_grad.dtype)
+    master_grad.data = a*converted_model_grad.data + b*stashed_grad.data
+    return False
+
+class LossScaler(object):
+    warned_no_fused_kernel = False
+    warned_unscaling_non_fp32_grad = False
+    has_fused_kernel = False
+
+    def __init__(self,
+                 loss_scale,
+                 init_scale=2.**16,
+                 scale_factor=2.,
+                 scale_window=2000,
+                 min_loss_scale=None,
+                 max_loss_scale=2.**24):
+        if loss_scale == "dynamic":
+            self.dynamic = True
+            self._loss_scale = min(max_loss_scale, init_scale)
+        else:
+            self.dynamic = False
+            self._loss_scale = loss_scale
+        self._max_loss_scale = max_loss_scale
+        self._min_loss_scale = min_loss_scale
+        self._scale_seq_len = scale_window
+        self._unskipped = 0
+        self._has_overflow = False
+        self._overflow_buf = torch.cuda.IntTensor([0])
+        if multi_tensor_applier.available:
+            import amp_C
+            LossScaler.has_fused_kernel = multi_tensor_applier.available
+            LossScaler.multi_tensor_scale_cuda = amp_C.multi_tensor_scale
+            LossScaler.multi_tensor_axpby_cuda = amp_C.multi_tensor_axpby
+        else:
+            if not LossScaler.warned_no_fused_kernel:
+                maybe_print(
+                    "Warning:  multi_tensor_applier fused unscale kernel is unavailable, "
+                    "possibly because apex was installed without --cuda_ext --cpp_ext. "
+                    "Using Python fallback.  Original ImportError was: " +
+                    repr(multi_tensor_applier.import_err),
+                    True)
+            LossScaler.has_fused_kernel = False
+            LossScaler.warned_no_fused_kernel = True
+
+    def loss_scale(self):
+        return self._loss_scale
+
+    def unscale_python(self, model_grads, master_grads, scale):
+        for model, master in zip(model_grads, master_grads):
+            if model is not None:
+                if not LossScaler.warned_unscaling_non_fp32_grad:
+                    if master.dtype != torch.float32:
+                        maybe_print(
+                            "Attempting to unscale a grad with type {} ".format(master.type()) +
+                            "Unscaling non-fp32 grads may indicate an error. "
+                            "When using Amp, you don't need to call .half() on your model.")
+                        LossScaler.warned_unscaling_non_fp32_grad = True
+                self._has_overflow = scale_check_overflow_python(model,
+                                                                 master,
+                                                                 1./scale,
+                                                                 self.dynamic)
+                if self._has_overflow and self.dynamic:
+                    break
+
+    # unused_scale keeps some of the old API alive for hopefully a short time.
+    def unscale(self, model_grads, master_grads, unused_scale, models_are_masters=False, scale_override=None):
+        if self._has_overflow:
+            return
+
+        scale = self._loss_scale
+        if scale_override is not None:
+            scale = scale_override
+
+        if scale == 1.0 and models_are_masters and not self.dynamic:
+            return
+
+        if LossScaler.has_fused_kernel:
+            # if (not LossScaler.warned_unscaling_non_fp32_grad
+            #     and master_grads[0].dtype == torch.float16):
+            #     print("Warning:  unscaling grads that are not FP32. "
+            #           "Unscaling non-fp32 grads may indicate an error. "
+            #           "When using Amp, you don't need to call .half() on your model.")
+            #     # Setting this to True unconditionally allows the possibility of an escape
+            #     # if never-before-seen non-fp32 grads are created in some later iteration.
+            #     LossScaler.warned_unscaling_non_fp32_grad = True
+            multi_tensor_applier(LossScaler.multi_tensor_scale_cuda,
+                                 self._overflow_buf,
+                                 [model_grads, master_grads],
+                                 1./scale)
+        else:
+            self.unscale_python(model_grads, master_grads, scale)
+
+        # Defer to update_scale
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+        #     self._has_overflow = self._overflow_buf.item()
+
+    def unscale_with_stashed_python(self,
+                                    model_grads,
+                                    stashed_master_grads,
+                                    master_grads,
+                                    a,
+                                    b):
+        for model, stashed, master in zip(model_grads, stashed_master_grads, master_grads):
+            if model is None and stashed is None:
+                continue
+            else:
+                if not LossScaler.warned_unscaling_non_fp32_grad:
+                    if master.dtype != torch.float32:
+                        maybe_print(
+                            "Attempting to unscale a grad with type {} ".format(master.type()) +
+                            "Unscaling non-fp32 grads may indicate an error. "
+                            "When using Amp, you don't need to call .half() on your model.")
+                        LossScaler.warned_unscaling_non_fp32_grad = True
+                self._has_overflow = axpby_check_overflow_python(model,
+                                                                 stashed,
+                                                                 master,
+                                                                 a,
+                                                                 b,
+                                                                 self.dynamic)
+                if self._has_overflow and self.dynamic:
+                    break
+
+    def unscale_with_stashed(self,
+                             model_grads,
+                             stashed_master_grads,
+                             master_grads,
+                             scale_override=None):
+        if self._has_overflow:
+            return
+
+        grads_have_scale, stashed_have_scale, out_scale = self._loss_scale, 1.0, 1.0
+        if scale_override is not None:
+            grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+        if LossScaler.has_fused_kernel:
+            if (not LossScaler.warned_unscaling_non_fp32_grad
+                and master_grads[0].dtype == torch.float16):
+                print("Warning:  unscaling grads that are not FP32. "
+                      "Unscaling non-fp32 grads may indicate an error. "
+                      "When using Amp, you don't need to call .half() on your model.")
+                # Setting this to True unconditionally allows the possibility of an escape
+                # if never-before-seen non-fp32 grads are created in some later iteration.
+                LossScaler.warned_unscaling_non_fp32_grad = True
+            multi_tensor_applier(LossScaler.multi_tensor_axpby_cuda,
+                                 self._overflow_buf,
+                                 [model_grads, stashed_master_grads, master_grads],
+                                 out_scale/grads_have_scale,   # 1./scale,
+                                 out_scale/stashed_have_scale, # 1.0,
+                                 0) # check only arg 0, aka the incoming model grads, for infs
+        else:
+            self.unscale_with_stashed_python(model_grads,
+                                             stashed_master_grads,
+                                             master_grads,
+                                             out_scale/grads_have_scale,
+                                             out_scale/stashed_have_scale)
+
+        # Defer to update_scale
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+        #     self._has_overflow = self._overflow_buf.item()
+
+    def clear_overflow_state(self):
+        self._has_overflow = False
+        if self.has_fused_kernel:
+            self._overflow_buf.zero_()
+
+    # Separate so unscale() can be called more that once before updating.
+    def update_scale(self):
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+            self._has_overflow = self._overflow_buf.item()
+
+        if self._has_overflow and self.dynamic:
+            should_skip = True
+            if(self._min_loss_scale):
+                self._loss_scale = max(self._min_loss_scale, self._loss_scale/2.)
+            else:
+                self._loss_scale = self._loss_scale/2.
+            self._unskipped = 0
+        else:
+            should_skip = False
+            self._unskipped += 1
+
+        if self._unskipped == self._scale_seq_len and self.dynamic:
+            self._loss_scale = min(self._max_loss_scale, self._loss_scale*2.)
+            self._unskipped = 0
+
+        return should_skip

apex/apex/amp/utils.py

@@ -0,0 +1,210 @@
+from . import compat
+
+import functools
+import itertools
+
+import torch
+
+def is_cuda_enabled():
+    return torch.version.cuda is not None
+
+def get_cuda_version():
+    return tuple(int(x) for x in torch.version.cuda.split('.'))
+
+def is_fp_tensor(x):
+    if is_nested(x):
+        # Fast-fail version of all(is_fp_tensor)
+        for y in x:
+            if not is_fp_tensor(y):
+                return False
+        return True
+    return compat.is_tensor_like(x) and compat.is_floating_point(x)
+
+def is_nested(x):
+    return isinstance(x, tuple) or isinstance(x, list)
+
+def should_cache(x):
+    if is_nested(x):
+        # Fast-fail version of all(should_cache)
+        for y in x:
+            if not should_cache(y):
+                return False
+        return True
+    return isinstance(x, torch.nn.parameter.Parameter) and \
+        type_string(x) == 'FloatTensor'
+
+def collect_fp_tensor_types(args, kwargs):
+    def collect_types(x, types):
+        if is_nested(x):
+            for y in x:
+                collect_types(y, types)
+        else:
+            types.add(type_string(x))
+
+    all_args = itertools.chain(args, kwargs.values())
+    types = set()
+    for x in all_args:
+        if is_fp_tensor(x):
+            collect_types(x, types)
+    return types
+
+def type_string(x):
+    return x.type().split('.')[-1]
+
+def maybe_half(x, name='', verbose=False):
+    if is_nested(x):
+        return type(x)([maybe_half(y) for y in x])
+
+    if not x.is_cuda or type_string(x) == 'HalfTensor':
+        return x
+    else:
+        if verbose:
+            print('Float->Half ({})'.format(name))
+        return x.half()
+
+def maybe_float(x, name='', verbose=False):
+    if is_nested(x):
+        return type(x)([maybe_float(y) for y in x])
+
+    if not x.is_cuda or type_string(x) == 'FloatTensor':
+        return x
+    else:
+        if verbose:
+            print('Half->Float ({})'.format(name))
+        return x.float()
+
+# NB: returneds casted `args`, mutates `kwargs` in-place
+def casted_args(cast_fn, args, kwargs):
+    new_args = []
+    for x in args:
+        if is_fp_tensor(x):
+            new_args.append(cast_fn(x))
+        else:
+            new_args.append(x)
+    for k in kwargs:
+        val = kwargs[k]
+        if is_fp_tensor(val):
+            kwargs[k] = cast_fn(val)
+    return new_args
+
+def cached_cast(cast_fn, x, cache):
+    if is_nested(x):
+        return type(x)([cached_cast(y) for y in x])
+    if x in cache:
+        cached_x = cache[x]
+        if x.requires_grad and cached_x.requires_grad:
+            # Make sure x is actually cached_x's autograd parent.
+            if cached_x.grad_fn.next_functions[1][0].variable is not x:
+                raise RuntimeError("x and cache[x] both require grad, but x is not "
+                                   "cache[x]'s parent.  This is likely an error.")
+        # During eval, it's possible to end up caching casted weights with
+        # requires_grad=False.  On the next training iter, if cached_x is found
+        # and reused from the cache, it will not actually have x as its parent.
+        # Therefore, we choose to invalidate the cache (and force refreshing the cast)
+        # if x.requires_grad and cached_x.requires_grad do not match.
+        #
+        # During eval (i.e. running under with torch.no_grad()) the invalidation
+        # check would cause the cached value to be dropped every time, because
+        # cached_x would always be created with requires_grad=False, while x would
+        # still have requires_grad=True.  This would render the cache effectively
+        # useless during eval.  Therefore, if we are running under the no_grad()
+        # context manager (torch.is_grad_enabled=False) we elide the invalidation
+        # check, and use the cached value even though its requires_grad flag doesn't
+        # match.  During eval, we don't care that there's no autograd-graph
+        # connection between x and cached_x.
+        if torch.is_grad_enabled() and x.requires_grad != cached_x.requires_grad:
+            del cache[x]
+        else:
+            return cached_x
+
+    casted_x = cast_fn(x)
+    cache[x] = casted_x
+    return casted_x
+
+def verbosify(cast_fn, fn_name, verbose):
+    if verbose:
+        return functools.partial(cast_fn, name=fn_name, verbose=verbose)
+    else:
+        return cast_fn
+
+def as_inplace(fns):
+    for x in fns:
+        yield x + '_'
+
+def has_func(mod, fn):
+    if isinstance(mod, dict):
+        return fn in mod
+    else:
+        return hasattr(mod, fn)
+
+def get_func(mod, fn):
+    if isinstance(mod, dict):
+        return mod[fn]
+    else:
+        return getattr(mod, fn)
+
+def set_func(mod, fn, new_fn):
+    if isinstance(mod, dict):
+        mod[fn] = new_fn
+    else:
+        setattr(mod, fn, new_fn)
+
+def set_func_save(handle, mod, fn, new_fn):
+    cur_fn = get_func(mod, fn)
+    handle._save_func(mod, fn, cur_fn)
+    set_func(mod, fn, new_fn)
+
+# A couple problems get solved here:
+# - The flat_weight buffer is disconnected from autograd graph,
+#   so the fp16 weights need to be derived from the input weights
+#   to this forward call, not the flat buffer.
+# - The ordering of weights in the flat buffer is...idiosyncratic.
+# First problem is solved with combination of set_ (to set up
+# correct storage) and copy_ (so the fp16 weight derives from the
+# fp32 one in autograd.
+# Second is solved by doing ptr arithmetic on the fp32 weights
+# to derive the correct offset.
+#
+# TODO: maybe this should actually use
+# `torch._cudnn_rnn_flatten_weight`? But then I need to call
+# on first iter and cache the right offsets. Ugh.
+def synthesize_flattened_rnn_weights(fp32_weights,
+                                     fp16_flat_tensor,
+                                     rnn_fn='',
+                                     verbose=False):
+    fp16_weights = []
+    fp32_base_ptr = fp32_weights[0][0].data_ptr()
+    for layer_weights in fp32_weights:
+        fp16_layer_weights = []
+        for w_fp32 in layer_weights:
+            w_fp16 = w_fp32.new().half()
+            offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+            w_fp16.set_(fp16_flat_tensor.storage(),
+                        offset,
+                        w_fp32.shape)
+            w_fp16.copy_(w_fp32)
+            if verbose:
+                print('Float->Half ({})'.format(rnn_fn))
+            fp16_layer_weights.append(w_fp16)
+        fp16_weights.append(fp16_layer_weights)
+    return fp16_weights
+
+# Roughly same as above, just the `fp32_weights` aren't nested.
+# Code kept separate for readability.
+def new_synthesize_flattened_rnn_weights(fp32_weights,
+                                         fp16_flat_tensor,
+                                         rnn_fn='',
+                                         verbose=False):
+    fp16_weights = []
+    fp32_base_ptr = fp32_weights[0].data_ptr()
+    for w_fp32 in fp32_weights:
+        w_fp16 = w_fp32.new().half()
+        offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+        w_fp16.set_(fp16_flat_tensor.storage(),
+                    offset,
+                    w_fp32.shape)
+        w_fp16.copy_(w_fp32)
+        if verbose:
+            print('Float->Half ({})'.format(rnn_fn))
+        fp16_weights.append(w_fp16)
+    return fp16_weights

apex/apex/amp/wrap.py

@@ -0,0 +1,276 @@
+from . import compat
+from . import utils
+from ._amp_state import _amp_state
+from . import rnn_compat
+
+import functools
+
+import torch
+
+def make_cast_wrapper(orig_fn, cast_fn, handle,
+                      try_caching=False):
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        if not handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        if try_caching and handle.has_cache:
+            args = list(args)
+            for i in range(len(args)):
+                if utils.should_cache(args[i]):
+                    args[i] = utils.cached_cast(cast_fn, args[i], handle.cache)
+            for k in kwargs:
+                if utils.should_cache(kwargs[k]):
+                    kwargs[k] = utils.cached_cast(cast_fn, kwargs[k], handle.cache)
+        new_args = utils.casted_args(cast_fn,
+                                     args,
+                                     kwargs)
+        return orig_fn(*new_args, **kwargs)
+    return wrapper
+
+def cached_cast(mod, fn, cast_fn, handle,
+                try_caching=False, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    cast_fn = utils.verbosify(cast_fn, fn, verbose)
+    wrapper = make_cast_wrapper(orig_fn, cast_fn, handle, try_caching)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+# `handle` arg is unused, but simplifies API to make `make_cast_wrapper`
+# Annoyingly, make_promote_wrapper still uses the global handle.  Once everyone
+# is on the new API and I am free to get rid of handle, I can clean this up.
+def make_promote_wrapper(orig_fn, cast_fn, handle=None):
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        if not _amp_state.handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        types = utils.collect_fp_tensor_types(args, kwargs)
+
+        if len(types) <= 1:
+            return orig_fn(*args, **kwargs)
+        elif len(types) == 2 and types == set(['HalfTensor', 'FloatTensor']):
+            new_args = utils.casted_args(cast_fn,
+                                         args,
+                                         kwargs)
+            return orig_fn(*new_args, **kwargs)
+        else:
+            raise NotImplementedError('Do not know how to handle ' +
+                                      'these types to promote: {}'
+                                      .format(types))
+    return wrapper
+
+def promote(mod, fn, handle, verbose=False):
+    orig_fn = utils.get_func(mod, fn)
+    maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+    wrapper = make_promote_wrapper(orig_fn, maybe_float)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def sequence_promote(mod, fn, handle, verbose=False):
+    orig_fn = utils.get_func(mod, fn)
+    maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+    @functools.wraps(orig_fn)
+    def wrapper(seq, *args, **kwargs):
+        if not _amp_state.handle.is_active():
+            return orig_fn(seq, *args, **kwargs)
+
+        types = set([utils.type_string(x) for x in seq])
+        if len(types) <= 1:
+            return orig_fn(seq, *args, **kwargs)
+        elif types == set(['HalfTensor', 'FloatTensor']):
+            cast_seq = utils.casted_args(maybe_float,
+                                         seq, {})
+            return orig_fn(cast_seq, *args, **kwargs)
+        else:
+            # TODO: other mixed-type cases aren't due to amp.
+            #       Just pass through?
+            return orig_fn(seq, *args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def promote_match_arg0(mod, fn, handle, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(arg0, *args, **kwargs):
+        assert compat.is_tensor_like(arg0)
+        if not _amp_state.handle.is_active():
+            return orig_fn(arg0, *args, **kwargs)
+
+        if utils.type_string(arg0) == 'HalfTensor':
+            cast_fn = utils.maybe_half
+        elif utils.type_string(arg0) == 'FloatTensor':
+            cast_fn = utils.maybe_float
+        else:
+            return orig_fn(arg0, *args, **kwargs)
+        cast_fn = utils.verbosify(cast_fn, fn, verbose)
+        new_args = utils.casted_args(cast_fn, args, kwargs)
+        return orig_fn(arg0, *new_args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_any_half(mod, fn, handle, custom_err_msg=None):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        types = utils.collect_fp_tensor_types(args, kwargs)
+        if 'HalfTensor' in types:
+            if custom_err_msg:
+                raise NotImplementedError(custom_err_msg)
+            else:
+                raise NotImplementedError('Cannot call in-place function ' +
+                                          '{} with fp16 arguments.'.format(fn))
+        else:
+            return orig_fn(*args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_arg0_half(mod, fn, handle, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(arg0, *args, **kwargs):
+        assert compat.is_tensor_like(arg0)
+        if utils.type_string(arg0) == 'HalfTensor':
+            raise NotImplementedError('Cannot call in-place method ' +
+                                      '{} on fp16 Tensors.'.format(fn))
+        else:
+            cast_fn = utils.verbosify(utils.maybe_float, fn, verbose)
+            new_args = utils.casted_args(cast_fn, args, kwargs)
+            return orig_fn(arg0, *new_args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+# Current RNN approach:
+# - Wrap top-level `RNN` function in thnn backend
+# - Will call into either CudnnRNN or AutogradRNN
+#  - Each of these are factory functions that return a per-iter
+#    `forward` function
+# - We interpose on the factory function to:
+#   1) Interpose on the actual forward function and put in casts
+#   2) Insert an fp16 `flat_weight` if necessary
+def rnn_cast(backend, fn, handle, verbose=False):
+    orig_rnn = utils.get_func(backend, fn)
+    @functools.wraps(orig_rnn)
+    def rnn_wrapper(*args, **kwargs):
+        flat_weight = kwargs.get('flat_weight')
+        if flat_weight is not None:
+            # We replace `flat_weight` with an uninitialized fp16
+            # Tensor. The "actual" weight tensors (provided in `forward`),
+            # will then be set up as ptrs into the buffer and have the
+            # corresponding fp32 values copied in.
+            # We need to call `copy` on the "actual" weights so that the
+            # autograd graph correctly backprops from the wgrads computed
+            # inside cuDNN (on fp16 weights) into the fp32 weights.
+            assert utils.type_string(flat_weight) == 'FloatTensor'
+            if compat.tensor_is_float_tensor() or compat.tensor_is_variable():
+                # Pre-0.4. A little slower, since it zeros out memory.
+                flat_weight_fp16 = flat_weight.new().half().resize_(flat_weight.shape)
+            else:
+                flat_weight_fp16 = torch.empty_like(flat_weight,
+                                                    dtype=torch.float16)
+            kwargs['flat_weight'] = flat_weight_fp16
+        else:
+            flat_weight_fp16 = None
+
+        forward = orig_rnn(*args, **kwargs)
+        @functools.wraps(forward)
+        def fwd_wrapper(*fargs, **fkwargs):
+            assert len(fargs) == 3 or len(fargs) == 4
+            inputs, weights, hiddens = fargs[:3]
+            assert utils.is_fp_tensor(inputs)
+            assert isinstance(weights, list)
+            cast_fn = utils.verbosify(utils.maybe_half,
+                                      fn,
+                                      verbose)
+            new_args = []
+
+            # 0) Inputs
+            new_args.append(cast_fn(inputs))
+
+            # 1) Weights
+            if flat_weight_fp16 is not None:
+                fp16_weights = utils.synthesize_flattened_rnn_weights(
+                    weights, flat_weight_fp16, fn, verbose)
+            else:
+                fp16_weights = [[cast_fn(w) for w in layer]
+                                for layer in weights]
+            new_args.append(fp16_weights)
+
+            # 2) Inputs: either a tuple (for LSTM) or single tensor
+            if isinstance(hiddens, tuple):
+                new_args.append(tuple(cast_fn(x) for x in hiddens))
+            elif utils.is_fp_tensor(hiddens):
+                new_args.append(cast_fn(hiddens))
+            else:
+                # Hiddens can, in principle, be `None` -- pass through
+                new_args.append(hiddens)
+
+            # 3) Batch sizes (0.4 or later only)
+            if len(fargs) == 4:
+                new_args.append(fargs[3])
+
+            return forward(*new_args, **fkwargs)
+        return fwd_wrapper
+    utils.set_func_save(handle, backend, fn, rnn_wrapper)
+
+def new_rnn_cast(fn, handle, verbose=False):
+    # Forward+backward compatibility around https://github.com/pytorch/pytorch/pull/15744
+    # For rnn backend calls that route through _rnn_impls, we must patch the ref
+    # that _rnn_impls stashed.  For rnn backend calls that directly invoke
+    # _VF.<backend>, e.g. _VF.lstm, we can patch onto VariableFunctionsShim,
+    # which in turn has patched the ref named "_VF" in torch.nn.modules.rnn.
+    if utils.has_func(torch.nn.modules.rnn._rnn_impls, fn):
+        mod = torch.nn.modules.rnn._rnn_impls
+    else:
+        mod = torch.nn.modules.rnn._VF
+        assert isinstance(mod, rnn_compat.VariableFunctionsShim)
+        fn = fn.lower()
+    orig_fn = utils.get_func(mod, fn)
+    cast_fn = utils.verbosify(utils.maybe_half, fn, verbose)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        # Exact call signature from modules/rnn.py
+        assert len(args) == 9
+        assert len(kwargs) == 0
+
+        if not _amp_state.handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        if isinstance(args[6], bool):
+            params_idx = 2 # Not PackedSequence case
+        else:
+            params_idx = 3 # PackedSequence case
+
+        new_args = []
+        for i, arg in enumerate(args):
+            if i == params_idx:
+                num_params = sum([x.numel() for x in arg])
+                fp16_weight_buf = args[0].new_empty((num_params,),
+                                                    dtype=torch.half)
+                casted_weights = utils.new_synthesize_flattened_rnn_weights(
+                    arg, fp16_weight_buf, fn, verbose)
+                new_args.append(casted_weights)
+            elif utils.is_fp_tensor(arg):
+                new_args.append(cast_fn(arg))
+            else:
+                new_args.append(arg)
+
+        return orig_fn(*new_args)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def disable_casts(mod, fn, handle):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        with handle._disable_casts():
+            return orig_fn(*args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)

apex/apex/contrib/bottleneck/__init__.py

@@ -0,0 +1 @@
+from .bottleneck import Bottleneck

apex/apex/contrib/bottleneck/bottleneck.py

@@ -0,0 +1,214 @@
+import torch
+from torch import nn
+import fast_bottleneck
+
+def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    weight_tensor_nchw = tensor
+    nn.init.kaiming_uniform_(weight_tensor_nchw, a=a, mode=mode, nonlinearity=nonlinearity)
+
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed
+    """
+    def __init__(self, n):
+        super(FrozenBatchNorm2d, self).__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+
+    def get_scale_bias(self, nhwc=False):
+        scale = self.weight * self.running_var.rsqrt()
+        bias = self.bias - self.running_mean * scale
+        if nhwc:
+            scale = scale.reshape(1, 1, 1, -1)
+            bias = bias.reshape(1, 1, 1, -1)
+        else:
+            scale = scale.reshape(1, -1, 1, 1)
+            bias = bias.reshape(1, -1, 1, 1)
+        return scale, bias
+
+    def forward(self, x):
+        scale, bias = self.get_scale_bias()
+        return x * scale + bias
+
+
+@torch.jit.script
+def drelu_dscale1(grad_o, output, scale1):
+    relu_mask = (output>0).half()
+    dx_relu = relu_mask * grad_o
+    g1 = dx_relu * scale1
+    return g1, dx_relu
+
+@torch.jit.script
+def drelu_dscale2(grad_o, output, scale1, scale2):
+    relu_mask = (output>0).half()
+    dx_relu = relu_mask * grad_o
+    g1 = dx_relu * scale1
+    g2 = dx_relu * scale2
+    return g1, g2
+
+class BottleneckFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, nhwc, stride_1x1, scale, bias, x, *conv):
+        # TODO: clean up order of tensors
+        args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]
+        ctx.downsample = len(conv) > 3
+        if ctx.downsample:
+            args.append(conv[3])
+            args.append(scale[3])
+            args.append(bias[3])
+
+        # weight buffers are always in nhwc while shape can be nhwc or channels_last
+        # here we pass in flag and let c++ handle it
+        # alternatively, we can put all sizes into a fixed format and pass it in
+        outputs = fast_bottleneck.forward(nhwc, stride_1x1, args)
+        ctx.save_for_backward(*(args+outputs))
+        # save relu outputs for drelu
+        ctx.nhwc = nhwc
+        ctx.stride_1x1 = stride_1x1
+        return outputs[2]
+
+    # backward relu is not exposed, MUL with mask used now
+    # only support dgrad
+    @staticmethod
+    def backward(ctx, grad_o):
+        outputs = ctx.saved_tensors[-3:]
+
+        if ctx.downsample:
+            grad_conv3, grad_conv4 = drelu_dscale2(grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11])
+        else:
+            grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])
+
+        # create input vector for backward
+        t_list = [*ctx.saved_tensors[0:10]]
+        t_list.append(grad_conv3)
+        t_list.append(grad_conv4)
+
+        # outputs used for wgrad and generating drelu mask
+        t_list.append(outputs[0])
+        t_list.append(outputs[1])
+
+        # in case there is downsample
+        if ctx.downsample:
+            t_list.append(ctx.saved_tensors[10])
+
+        grads = fast_bottleneck.backward(ctx.nhwc, ctx.stride_1x1, t_list)
+
+        return (None, None, None, None, *grads)
+
+bottleneck_function = BottleneckFunction.apply
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class Bottleneck(torch.nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+    # here we put it at 1x1
+
+    def __init__(self, in_channels, bottleneck_channels, out_channels, stride=1, groups=1,
+                 dilation=1, norm_func=None, use_cudnn=False, explicit_nhwc=False):
+        super(Bottleneck, self).__init__()
+        if groups != 1:
+            raise RuntimeError('Only support groups == 1')
+        if dilation != 1:
+            raise RuntimeError('Only support dilation == 1')
+        if norm_func == None:
+            norm_func = FrozenBatchNorm2d
+        else:
+            raise RuntimeError('Only support frozen BN now.')
+
+        if stride != 1 or in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels, stride),
+                norm_func(out_channels),
+            )
+        else:
+            self.downsample = None
+
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)
+        self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)
+        self.conv3 = conv1x1(bottleneck_channels, out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+
+        self.bn1 = norm_func(bottleneck_channels)
+        self.bn2 = norm_func(bottleneck_channels)
+        self.bn3 = norm_func(out_channels)
+
+        self.use_cudnn = use_cudnn
+
+        # setup conv weights
+        self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]
+        if self.downsample is not None:
+            self.w_conv.append(self.downsample[0].weight)
+
+        # init weight in nchw format before possible transpose
+        for w in self.w_conv:
+            kaiming_uniform_(w, a=1)
+
+        # TODO: prevent unsupported case usage
+        # support cases
+        #                 native      cudnn
+        # normal             yes         no
+        # channel_last       yes        yes
+        # explicit_nhwc       no        yes
+        self.explicit_nhwc = explicit_nhwc
+        if self.explicit_nhwc:
+            for p in self.parameters():
+                with torch.no_grad():
+                    p.data = p.data.permute(0,2,3,1).contiguous()
+        return
+
+    def forward(self, x):
+        if self.use_cudnn:
+            # calculate scale/bias from registered buffers
+            # TODO: make this better
+            s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)
+            s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)
+            s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)
+            w_scale = [s1, s2, s3]
+            w_bias = [b1, b2, b3]
+            if self.downsample is not None:
+                s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)
+                w_scale.append(s4)
+                w_bias.append(b4)
+
+            out = bottleneck_function(self.explicit_nhwc, self.stride, w_scale, w_bias, x, *self.w_conv)
+            return out
+
+        if self.explicit_nhwc:
+            raise RuntimeError('explicit nhwc with native ops is not supported.')
+
+        # fallback to native ops
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out

apex/apex/contrib/bottleneck/test.py

@@ -0,0 +1,71 @@
+import torch
+from bottleneck import Bottleneck
+torch.manual_seed(23337)
+
+# use True to print layerwise sum for all outputs in reference code path
+DEBUG = False#True
+
+for stride, o_channel in [(1,32), (1,128), (2,32)]:
+    print("testing stride ==", stride, ", in_channel == 32 , out_channel ==", o_channel)
+    a_ = torch.randn(17,32,28,28)
+
+    a = a_.cuda().half().to(memory_format=torch.channels_last).requires_grad_()
+    model = Bottleneck(32,8,o_channel,stride=stride).cuda().half().to(memory_format=torch.channels_last)
+
+    # test model
+    b = model(a)
+    b.mean().backward()
+    d_grad = a.grad.float()
+    a.grad = None
+    torch.cuda.synchronize()
+
+    if DEBUG:
+        print("[DEBUG] ref dx :", d_grad.sum().item())
+        # print wgrad. we don't need to reset since later cpp print before accumulation
+        for i, w in enumerate(model.w_conv):
+            print("[DEBUG] ref wgrad{} :".format(i+1), w.grad.sum().item())
+
+    wgrads = []
+    for w in model.w_conv:
+        wgrads.append(w.grad.float())
+
+    model.use_cudnn = True
+    model.zero_grad()
+    c = model(a)
+    c.mean().backward()
+
+    torch.cuda.synchronize()
+    print("comparing native and channels_last:")
+    print("max error fprop:", (b-c).abs().max().item(), "max elem:", b.abs().max().item())
+    print("max error dgrad:", (d_grad-a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+    for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
+        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())
+
+    nhwc_a = a_.permute(0,2,3,1).contiguous().cuda().half().requires_grad_()
+    nhwc_model = Bottleneck(32,8,o_channel,stride=stride,explicit_nhwc=True, use_cudnn=True).cuda().half()
+    for p,q in zip(model.parameters(), nhwc_model.parameters()):
+        # model's storage is already in nhwc, we clone and assign to explicit nhwc model
+        q.data.copy_(p.data.permute(0,2,3,1).contiguous())
+    for p,q in zip(model.buffers(), nhwc_model.buffers()):
+        q.data.copy_(p.data)
+
+    d = nhwc_model(nhwc_a)
+    d.mean().backward()
+    torch.cuda.synchronize()
+
+    # reset reference to cudnn channels_last permute
+    #c_s = c.storage().tolist()
+    #d_s = d.storage().tolist()
+    #print(max([x-y for x,y in zip(c_s,d_s)]))
+    c = c.contiguous(memory_format=torch.contiguous_format).permute(0,2,3,1).contiguous()
+    d_grad = a.grad.float().permute(0,2,3,1).contiguous()
+    wgrads = []
+    for w in model.w_conv:
+        wgrads.append(w.grad.float().permute(0,2,3,1).contiguous())
+
+    torch.cuda.synchronize()
+    print("comparing nhwc and channels_last:")
+    print("max error fprop:", (d-c).abs().max().item(), "max elem:", c.abs().max().item())
+    print("max error dgrad:", (d_grad-nhwc_a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+    for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
+        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())

apex/apex/contrib/csrc/bottleneck/bottleneck.cpp

@@ -0,0 +1,1612 @@
+#include <ATen/ATen.h>
+#include <ATen/cudnn/Handle.h>  // for getcudnnhandle
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <vector>
+#include <cudnn_frontend.h>
+
+#include <iostream>
+
+#ifdef DEBUG
+#define DEBUG_MSG(str) do { std::cout << str << std::endl; } while( false )
+#else
+#define DEBUG_MSG(str) do { } while ( false )
+#endif
+
+#ifdef DEBUG_CUDNN
+#define DEBUG_CUDNN_MSG(buf, str) do { buf << str << std::endl; } while( false )
+#else
+#define DEBUG_CUDNN_MSG(buf, str) do { } while ( false )
+#endif
+
+#define checkCudnnErr(...)                                                        \
+    do {                                                                          \
+        int err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \
+        if (err) {                                                                \
+            return;                                                    \
+	}                                                                         \
+    } while (0)
+
+
+int checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {
+    if (code) {
+        printf("CUDNN error at %s:%d, code=%d (%s) in '%s'\n", file, line, (int)code, cudnnGetErrorString(code), expr);
+        return 1;
+    }
+    return 0;
+}
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort = true);
+#define checkCUDAError(val) { checkError((val), #val, __FILE__, __LINE__); }    // in-line regular function
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort)
+{
+  if (code != cudaSuccess)
+  {
+    const char * errorMessage = cudaGetErrorString(code);
+    fprintf(stderr, "CUDA error returned from \"%s\" at %s:%d, Error code: %d (%s)\n", func, file, line, code, errorMessage);
+    if (abort){
+      cudaDeviceReset();
+      exit(code);
+    }
+  }
+}
+
+void generateStrides(const int64_t* dimA, int64_t* strideA, int nbDims, cudnnTensorFormat_t filterFormat) {
+    // For INT8x4 and INT8x32 we still compute standard strides here to input
+    // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.
+    if (filterFormat == CUDNN_TENSOR_NCHW) {
+        strideA[nbDims - 1] = 1;
+        for (int64_t d = nbDims - 2; d >= 0; d--) {
+            strideA[d] = strideA[d + 1] * dimA[d + 1];
+        }
+    } else {
+        // Here we assume that the format is CUDNN_TENSOR_NHWC
+	strideA[1]          = 1;
+        strideA[nbDims - 1] = strideA[1] * dimA[1];
+        for (int64_t d = nbDims - 2; d >= 2; d--) {
+            strideA[d] = strideA[d + 1] * dimA[d + 1];
+        }
+        strideA[0] = strideA[2] * dimA[2];
+    }
+}
+
+
+int getFwdConvDilatedFilterDim(int filterDim, int dilation) {
+    return ((filterDim - 1) * dilation) + 1;
+}
+
+int getFwdConvPaddedImageDim(int tensorDim, int pad) {
+    return tensorDim + (2 * pad);
+}
+
+int getFwdConvOutputDim(
+    int tensorDim,
+    int pad,
+    int filterDim,
+    int stride,
+    int dilation)
+{
+    int p = (getFwdConvPaddedImageDim(tensorDim, pad) - getFwdConvDilatedFilterDim(filterDim, dilation)) / stride + 1;
+    return (p);
+}
+
+enum {
+    X_TENSOR,
+    Y_TENSOR,
+    W_TENSOR,
+    Z_TENSOR,
+    B_TENSOR,
+    AFTERADD_TENSOR,
+    AFTERBIAS_TENSOR,
+    AFTERCONV_TENSOR,
+    OPTIONAL,
+    AFTEROPT_TENSOR,
+};
+
+using common_conv_descriptors =
+    std::tuple<cudnn_frontend::Tensor, cudnn_frontend::Tensor, cudnn_frontend::Tensor, cudnn_frontend::ConvDesc>;
+
+
+common_conv_descriptors
+create_common_descriptors(int64_t* x_dim_padded,
+                          int64_t* padA,
+                          int64_t* convstrideA,
+                          int64_t* dilationA,
+                          int64_t* w_dim_padded,
+                          int64_t* y_dim_padded,
+                          cudnnDataType_t dataType,
+                          cudnnConvolutionMode_t mode) {
+    const int convDim = 2;
+
+    int64_t strideA_padded[4];
+    int64_t outstrideA_padded[4];
+    int64_t filterstrideA_padded[4];
+
+    generateStrides(w_dim_padded, filterstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, strideA_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, outstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return common_conv_descriptors(cudnn_frontend::TensorBuilder()
+                                       .setDim(4, x_dim_padded)
+                                       .setStrides(4, strideA_padded)
+                                       .setId('x')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::TensorBuilder()
+                                       .setDim(4, y_dim_padded)
+                                       .setStrides(4, outstrideA_padded)
+                                       .setId('y')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::TensorBuilder()
+                                       .setDim(4, w_dim_padded)
+                                       .setStrides(4, filterstrideA_padded)
+                                       .setId('w')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::ConvDescBuilder()
+                                       .setDataType(CUDNN_DATA_FLOAT)
+                                       .setMathMode(mode)
+                                       .setNDims(convDim)
+                                       .setStrides(convDim, convstrideA)
+                                       .setPrePadding(convDim, padA)
+                                       .setPostPadding(convDim, padA)
+                                       .setDilation(convDim, dilationA)
+                                       .build());
+}
+
+using common_convbias_descriptors = std::tuple<cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor>;
+
+common_convbias_descriptors
+create_conv_bias_add_act_descriptors(int64_t* x_dim_padded,
+                                     int64_t* padA,
+                                     int64_t* convstrideA,
+                                     int64_t* dilationA,
+                                     int64_t* w_dim_padded,
+                                     int64_t* y_dim_padded,
+                                     cudnnDataType_t dataType) {
+    const int convDim = 2;
+
+    int64_t b_dim_padded[4];
+    b_dim_padded[0] = 1;
+    b_dim_padded[1] = y_dim_padded[1];
+    b_dim_padded[2] = 1;
+    b_dim_padded[3] = 1;
+
+    int64_t x_stride_padded[4];
+    int64_t y_stride_padded[4];
+    int64_t w_stride_padded[4];
+    int64_t b_stride_padded[4];
+
+    generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return common_convbias_descriptors(cudnn_frontend::TensorBuilder()
+                                           .setDim(4, x_dim_padded)
+                                           .setStrides(4, x_stride_padded)
+                                           .setId('x')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('y')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, w_dim_padded)
+                                           .setStrides(4, w_stride_padded)
+                                           .setId('w')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, b_dim_padded)
+                                           .setStrides(4, b_stride_padded)
+                                           .setId('z')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, b_dim_padded)
+                                           .setStrides(4, b_stride_padded)
+                                           .setId('b')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setVirtual()
+                                           .setId('A')  // after add
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setVirtual()
+                                           .setId('B')  // after bias
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('C')  // after conv
+                                           .setAlignment(16)
+                                           .setVirtual()
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('i')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('D')  // after optional add
+                                           .setAlignment(16)
+                                           .setVirtual()
+                                           .setDataType(dataType)
+                                           .build());
+}
+
+// tensor descriptors used for dgrad
+enum {
+    X_OR_DX_TENSOR,
+    DY_TENSOR,
+    W_OR_DW_TENSOR,
+    SCALE_TENSOR,
+    RELU_TENSOR,
+    AFTER_DCONV_TENSOR,
+    AFTER_DRELU_TENSOR,
+};
+
+using dconv_descriptors = std::tuple<cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor>;
+
+dconv_descriptors
+create_dconv_descriptors(int64_t* x_dim_padded,
+                         int64_t* padA,
+                         int64_t* convstrideA,
+                         int64_t* dilationA,
+                         int64_t* w_dim_padded,
+                         int64_t* y_dim_padded,
+                         cudnnDataType_t dataType) {
+    const int convDim = 2;
+
+    int64_t b_dim_padded[4];
+    b_dim_padded[0] = 1;
+    b_dim_padded[1] = x_dim_padded[1];
+    b_dim_padded[2] = 1;
+    b_dim_padded[3] = 1;
+
+    int64_t x_stride_padded[4];
+    int64_t y_stride_padded[4];
+    int64_t w_stride_padded[4];
+    int64_t b_stride_padded[4];
+
+    generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return dconv_descriptors(cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setId('x')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, y_dim_padded)
+                             .setStrides(4, y_stride_padded)
+                             .setId('y')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, w_dim_padded)
+                             .setStrides(4, w_stride_padded)
+                             .setId('w')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, b_dim_padded)
+                             .setStrides(4, b_stride_padded)
+                             .setId('s')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setId('r')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setVirtual()
+                             .setId('A')  // after dconv
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setVirtual()
+                             .setId('B')  // after drelu
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build());
+}
+
+// create a cache for plan
+std::unordered_map<std::string, cudnn_frontend::ExecutionPlan> plan_cache;
+
+// TODO: better name
+std::string getConvFusionString(int64_t* x_dim_padded,
+                                int64_t* padA,
+                                int64_t* convstrideA,
+                                int64_t* dilationA,
+                                int64_t* w_dim_padded,
+                                cudnnDataType_t dataType,
+                                std::string fusion_string) {
+
+  for(int i=0;i<4;i++) {
+    fusion_string += 'X';
+    fusion_string += std::to_string(x_dim_padded[i]);
+  }
+  for(int i=0;i<4;i++) {
+    fusion_string += 'W';
+    fusion_string += std::to_string(w_dim_padded[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'P';
+    fusion_string += std::to_string(padA[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'S';
+    fusion_string += std::to_string(convstrideA[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'D';
+    fusion_string += std::to_string(dilationA[i]);
+  }
+  fusion_string += 'T';
+  fusion_string += std::to_string(dataType);
+  return fusion_string;
+}
+
+cudnn_frontend::ExecutionPlan& getOrCreatePlan(cudnnHandle_t handle_,
+                                               std::stringstream& log_buf,
+                                               cudnn_frontend::OperationGraph& opGraph,
+                                               std::string cache_string,
+                                               bool use_heuristic = true){
+  auto it = plan_cache.find(cache_string);
+  if (it != plan_cache.end()) {
+    DEBUG_CUDNN_MSG(log_buf, "Found plan in cache");
+    return it->second;
+  } else {
+    if (use_heuristic){
+      // TODO: confirm which mode to use
+      auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()
+        .setOperationGraph(opGraph)
+        .setHeurMode(CUDNN_HEUR_MODE_INSTANT)
+        .build();
+      // try 3 times for now as WAR for no heuristic training
+      int max_tries = 3, count = 0;
+      auto& engine_configs = heuristics.getEngineConfig(max_tries);
+      while(true) {
+        try {
+          plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder()
+                                                     .setHandle(handle_)
+                                                     .setEngineConfig(engine_configs[count], opGraph.getTag())
+                                                     .build()));
+          break;
+        } catch (cudnn_frontend::cudnnException e) {
+          if (++count == max_tries) throw e;
+        }
+      }
+    }else{
+    DEBUG_CUDNN_MSG(log_buf, "No plan in cache");
+    // How many engines support this operation graph ?
+    auto total_engines = opGraph.getEngineCount();
+    DEBUG_CUDNN_MSG(log_buf, opGraph.describe() << " has " << total_engines << " engines.");
+    // We have to randomly pick one engine from [0, total_engines)
+    // Selecting "0" by default
+    auto engine = cudnn_frontend::EngineBuilder().setGlobalEngineIdx(0).setOperationGraph(opGraph).build();
+    DEBUG_CUDNN_MSG(log_buf, engine.describe());
+    auto& knobs = engine.getSupportedKnobs();
+    for (auto it = std::begin(knobs); it != std::end(knobs); ++it) {
+      DEBUG_CUDNN_MSG(log_buf, it->describe());
+    }
+    if (knobs.begin() != knobs.end()) {
+      DEBUG_CUDNN_MSG(log_buf, "Updated knob choice");
+      knobs.begin()->setChoice(knobs.begin()->getMinValue() + 1);
+      DEBUG_CUDNN_MSG(log_buf, knobs.begin()->describe());
+    }
+
+    // Createmplacee the requisite engine config
+    auto engine_config = cudnn_frontend::EngineConfigBuilder().setEngine(engine).build();
+    DEBUG_CUDNN_MSG(log_buf, engine_config.describe());
+    plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder().setHandle(handle_).setEngineConfig(engine_config).build()));
+    }
+
+    return plan_cache.find(cache_string)->second;
+  }
+}
+
+void
+run_conv_scale_bias_add_activation(int64_t* x_dim_padded,
+                                   int64_t* pad,
+                                   int64_t* convstride,
+                                   int64_t* dilation,
+                                   int64_t* w_dim_padded,
+                                   int64_t* y_dim_padded,
+                                   cudnnDataType_t dataType,
+                                   at::Half* devPtrX,
+                                   at::Half* devPtrW,
+                                   at::Half* devPtrY,
+                                   at::Half* devPtrZ,
+                                   at::Half* devPtrB,
+                                   at::Half* devPtrI) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Y_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Z_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<B_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERADD_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERBIAS_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<OPTIONAL>(tensors).describe());
+
+        // Define the add operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+                           .setMode(CUDNN_POINTWISE_MUL)
+                           .setMathPrecision(CUDNN_DATA_FLOAT)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        // Define the bias operation
+        auto biasDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());
+
+        // optional add
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        // Define the activation operation
+        auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+                           .setMode(CUDNN_POINTWISE_RELU_FWD)
+                           .setMathPrecision(CUDNN_DATA_FLOAT)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+                           .setxDesc(std::get<X_TENSOR>(tensors))
+                           .setwDesc(std::get<W_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERCONV_TENSOR>(tensors))
+                           .setcDesc(convDesc)
+                           .setAlpha(alpha)
+                           .setBeta(beta)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create a Add Node with scaling parameters.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(conv_op.getOutputTensor())
+                           .setbDesc(std::get<Z_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERADD_TENSOR>(tensors))
+                           .setpwDesc(scaleDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create a Bias Node.
+        auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(scale_op.getOutputTensor())
+                           .setbDesc(std::get<B_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERBIAS_TENSOR>(tensors))
+                           .setpwDesc(biasDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, bias_op.describe());
+
+        // Create a optional add Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(bias_op.getOutputTensor())
+                           .setbDesc(std::get<OPTIONAL>(tensors))
+                           .setyDesc(std::get<AFTEROPT_TENSOR>(tensors))
+                           .setpwDesc(addDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+
+        // Create an Activation Node.
+        auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())
+                          .setyDesc(std::get<Y_TENSOR>(tensors))
+                          .setpwDesc(actDesc)
+                          .build();
+        DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 5> ops = {&conv_op, &scale_op, &bias_op, devPtrI ? &add_op : &act_op, &act_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(devPtrI ? ops.size() : 4, ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};
+        int64_t uids[]    = {'x', 'y', 'w', 'z', 'b', 'i'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+                               .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(devPtrI ? 6 : 5, data_ptrs)
+          .setUids(devPtrI ? 6 : 5, uids)
+                               .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_conv_scale_bias(int64_t* x_dim_padded,
+                    int64_t* pad,
+                    int64_t* convstride,
+                    int64_t* dilation,
+                    int64_t* w_dim_padded,
+                    int64_t* y_dim_padded,
+                    cudnnDataType_t dataType,
+                    at::Half* devPtrX,
+                    at::Half* devPtrW,
+                    at::Half* devPtrY,
+                    at::Half* devPtrZ,
+                    at::Half* devPtrB) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Y_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Z_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<B_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERADD_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERBIAS_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<OPTIONAL>(tensors).describe());
+
+        // Define the add operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_MUL)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        // Define the bias operation
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+                           .setxDesc(std::get<X_TENSOR>(tensors))
+                           .setwDesc(std::get<W_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERCONV_TENSOR>(tensors))
+                           .setcDesc(convDesc)
+                           .setAlpha(alpha)
+                           .setBeta(beta)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create a Add Node with scaling parameters.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(conv_op.getOutputTensor())
+          .setbDesc(std::get<Z_TENSOR>(tensors))
+          .setyDesc(std::get<AFTERADD_TENSOR>(tensors)) // TODO: change enum to aftermul
+          .setpwDesc(scaleDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create a Bias Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(scale_op.getOutputTensor())
+          .setbDesc(std::get<B_TENSOR>(tensors))
+          .setyDesc(std::get<Y_TENSOR>(tensors))
+          .setpwDesc(addDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 3> ops = {&conv_op, &scale_op, &add_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB};
+        int64_t uids[]    = {'x', 'y', 'w', 'z', 'b'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+                               .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(5, data_ptrs)
+          .setUids(5, uids)
+                               .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+
+void
+run_dconv_drelu_dscale(int64_t* x_dim_padded,
+                       int64_t* pad,
+                       int64_t* convstride,
+                       int64_t* dilation,
+                       int64_t* w_dim_padded,
+                       int64_t* y_dim_padded,
+                       cudnnDataType_t dataType,
+                       at::Half* devPtrX,
+                       at::Half* devPtrW,
+                       at::Half* devPtrY,
+                       at::Half* devPtrZ,
+                       at::Half* devPtrR) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        // Define the activation backward operation
+        auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_RELU_BWD)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+        // Define the scale backward operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_MUL)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+          .setdxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+          .setdyDesc(std::get<DY_TENSOR>(tensors))
+          .setcDesc(convDesc)
+          .setAlpha(alpha)
+          .setBeta(beta)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // TODO: do we need getOutputTensor(), and what it returns in backward case?
+        // Create an relu backward Node.
+        auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setdyDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setxDesc(std::get<RELU_TENSOR>(tensors))
+          .setdxDesc(std::get<AFTER_DRELU_TENSOR>(tensors))
+          .setpwDesc(actDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+        // Create a Scale Node.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(std::get<AFTER_DRELU_TENSOR>(tensors))
+          .setbDesc(std::get<SCALE_TENSOR>(tensors))
+          .setyDesc(std::get<X_OR_DX_TENSOR>(tensors))
+          .setpwDesc(scaleDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 3> ops = {&conv_op, &act_op, &scale_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR};
+        int64_t uids[]    = {'x', 'y', 'w', 's', 'r'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(5, data_ptrs)
+          .setUids(5, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_dconv(int64_t* x_dim_padded,
+          int64_t* pad,
+          int64_t* convstride,
+          int64_t* dilation,
+          int64_t* w_dim_padded,
+          int64_t* y_dim_padded,
+          cudnnDataType_t dataType,
+          at::Half* devPtrX,
+          at::Half* devPtrW,
+          at::Half* devPtrY,
+          cudnnBackendDescriptorType_t mode) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        // mode should be one of following
+        // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR
+        // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR
+        auto conv_op_builder = cudnn_frontend::OperationBuilder(mode);
+        if (mode == CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR) {
+          conv_op_builder.setdxDesc(std::get<X_OR_DX_TENSOR>(tensors))
+            .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+            .setdyDesc(std::get<DY_TENSOR>(tensors))
+            .setcDesc(convDesc)
+            .setAlpha(alpha)
+            .setBeta(beta);
+        }
+        else {
+          conv_op_builder.setxDesc(std::get<X_OR_DX_TENSOR>(tensors))
+            .setdwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+            .setdyDesc(std::get<DY_TENSOR>(tensors))
+            .setcDesc(convDesc)
+            .setAlpha(alpha)
+            .setBeta(beta);
+        }
+        auto conv_op = conv_op_builder.build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 1> ops = {&conv_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW};
+        int64_t uids[]    = {'x', 'y', 'w'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(3, data_ptrs)
+          .setUids(3, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_dconv_add(int64_t* x_dim_padded,
+              int64_t* pad,
+              int64_t* convstride,
+              int64_t* dilation,
+              int64_t* w_dim_padded,
+              int64_t* y_dim_padded,
+              cudnnDataType_t dataType,
+              at::Half* devPtrX,
+              at::Half* devPtrW,
+              at::Half* devPtrY,
+              at::Half* devPtrR) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        // Define the add backward operation
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_ADD)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+          .setdxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+          .setdyDesc(std::get<DY_TENSOR>(tensors))
+          .setcDesc(convDesc)
+          .setAlpha(alpha)
+          .setBeta(beta)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // TODO: do we need getOutputTensor(), and what it returns in backward case?
+        // Create add Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setbDesc(std::get<RELU_TENSOR>(tensors))
+          .setyDesc(std::get<X_OR_DX_TENSOR>(tensors))
+          .setpwDesc(addDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 2> ops = {&conv_op, &add_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrR};
+        int64_t uids[]    = {'x', 'y', 'w', 'r'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(4, data_ptrs)
+          .setUids(4, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+
+// inputs contains x,w,z,b,(i)
+std::vector<at::Tensor> bottleneck_forward(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+
+  std::cout << std::fixed;
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // setup dimensions
+  int64_t dimA[]         = {0, 0, 0, 0};
+  int64_t filterdimA1[]  = {0, 0, 0, 0};
+  int64_t filterdimA2[]  = {0, 0, 0, 0};
+  int64_t filterdimA3[]  = {0, 0, 0, 0};
+  int64_t filterdimA4[]  = {0, 0, 0, 0};
+
+  // All dim calculation after this order of n,c,h,w
+  int axis[] {0,1,2,3};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 3;
+    axis[2] = 1;
+    axis[3] = 2;
+  }
+  for (int dim=0;dim<4;dim++) {
+    dimA[dim] = inputs[0].size(axis[dim]);
+    filterdimA1[dim] = inputs[1].size(axis[dim]);
+    filterdimA2[dim] = inputs[2].size(axis[dim]);
+    filterdimA3[dim] = inputs[3].size(axis[dim]);
+  }
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    for (int dim=0;dim<4;dim++) {
+      filterdimA4[dim] = inputs[10].size(axis[dim]);
+    }
+  }
+
+  // output dim in n,c,h,w used by backend
+  int64_t outdimA1[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA2[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA3[]     = {0, 0, 0, 0}; // Computed Below
+
+  // use these fixed value for test run
+  int64_t padA[]        = {0, 0};
+  int64_t padA1[]        = {1, 1};
+  int64_t dilationA[] = {1, 1};
+  int64_t convstrideA[] = {1, 1};
+  int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+  // compute output from pad/stride/dilation
+  outdimA1[0] = dimA[0];
+  outdimA1[1] = filterdimA1[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+  }
+
+  outdimA2[0] = outdimA1[0];
+  outdimA2[1] = filterdimA2[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  outdimA3[0] = outdimA2[0];
+  outdimA3[1] = filterdimA3[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  // Create output tensor in the correct shape in pytorch's view
+  int64_t outdim1[]     = {0, 0, 0, 0};
+  int64_t outdim2[]     = {0, 0, 0, 0};
+  int64_t outdim3[]     = {0, 0, 0, 0};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 2;
+    axis[2] = 3;
+    axis[3] = 1;
+  }
+  for (int dim=0;dim<4;dim++) {
+    outdim1[dim] = outdimA1[axis[dim]];
+    outdim2[dim] = outdimA2[axis[dim]];
+    outdim3[dim] = outdimA3[axis[dim]];
+  }
+
+  // run
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+  at::Half* w = inputs[1].data_ptr<at::Half>();
+  at::Half* z = inputs[4].data_ptr<at::Half>();
+  at::Half* b = inputs[7].data_ptr<at::Half>();
+  auto out1 = at::empty(outdim1, inputs[0].type(), output_format);
+  at::Half* y1 = out1.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(dimA,
+                                     padA,
+                                     convstride1X1,
+                                     dilationA,
+                                     filterdimA1,
+                                     outdimA1,
+                                     CUDNN_DATA_HALF,
+                                     x,
+                                     w,
+                                     y1,
+                                     z,
+                                     b,
+                                     nullptr);
+
+  DEBUG_MSG("[DEBUG] new relu1 : " << out1.to(at::kFloat).sum().item<float>());
+
+  w = inputs[2].data_ptr<at::Half>();
+  z = inputs[5].data_ptr<at::Half>();
+  b = inputs[8].data_ptr<at::Half>();
+  auto out2 = at::empty(outdim2, inputs[0].type(), output_format);
+  at::Half* y2 = out2.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(outdimA1,
+                                     padA1,
+                                     convstrideA,
+                                     dilationA,
+                                     filterdimA2,
+                                     outdimA2,
+                                     CUDNN_DATA_HALF,
+                                     y1,
+                                     w,
+                                     y2,
+                                     z,
+                                     b,
+                                     nullptr);
+  DEBUG_MSG("[DEBUG] new relu2 : " << out2.to(at::kFloat).sum().item<float>());
+
+  // create output of conv3
+  auto out3 = at::empty(outdim3, inputs[0].type(), output_format);
+  at::Half* y3 = out3.data_ptr<at::Half>();
+
+  // create output of conv4 that may exist
+  auto identity = at::empty_like(out3);
+  at::Half* yi = identity.data_ptr<at::Half>();
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+
+    w = inputs[10].data_ptr<at::Half>();
+    z = inputs[11].data_ptr<at::Half>();
+    b = inputs[12].data_ptr<at::Half>();
+    run_conv_scale_bias(dimA,
+                        padA,
+                        convstride1X1,
+                        dilationA,
+                        filterdimA4,
+                        outdimA3,
+                        CUDNN_DATA_HALF,
+                        x,
+                        w,
+                        yi,
+                        z,
+                        b);
+    DEBUG_MSG("[DEBUG] new downsample : " << identity.to(at::kFloat).sum().item<float>());
+  }
+  else {
+    yi = x;
+  }
+
+  w = inputs[3].data_ptr<at::Half>();
+  z = inputs[6].data_ptr<at::Half>();
+  b = inputs[9].data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(outdimA2,
+                                     padA,
+                                     convstrideA,
+                                     dilationA,
+                                     filterdimA3,
+                                     outdimA3,
+                                     CUDNN_DATA_HALF,
+                                     y2,
+                                     w,
+                                     y3,
+                                     z,
+                                     b,
+                                     yi);
+  DEBUG_MSG("[DEBUG] new relu3 : " << out3.to(at::kFloat).sum().item<float>());
+
+  outputs.push_back(out1);
+  outputs.push_back(out2);
+  outputs.push_back(out3);
+
+  return outputs;
+}
+
+std::vector<at::Tensor> bottleneck_backward(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // setup dimensions
+  int64_t dimA[]         = {0, 0, 0, 0};
+  int64_t filterdimA1[]  = {0, 0, 0, 0};
+  int64_t filterdimA2[]  = {0, 0, 0, 0};
+  int64_t filterdimA3[]  = {0, 0, 0, 0};
+  int64_t filterdimA4[]  = {0, 0, 0, 0};
+
+  // All dim calculation after this order of n,c,h,w
+  int axis[] {0,1,2,3};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 3;
+    axis[2] = 1;
+    axis[3] = 2;
+  }
+  for (int dim=0;dim<4;dim++) {
+    dimA[dim] = inputs[0].size(axis[dim]);
+    filterdimA1[dim] = inputs[1].size(axis[dim]);
+    filterdimA2[dim] = inputs[2].size(axis[dim]);
+    filterdimA3[dim] = inputs[3].size(axis[dim]);
+  }
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    for (int dim=0;dim<4;dim++) {
+      filterdimA4[dim] = inputs[14].size(axis[dim]);
+    }
+  }
+
+  // output dim in n,c,h,w used by backend
+  int64_t outdimA1[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA2[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA3[]     = {0, 0, 0, 0}; // Computed Below
+
+  // use these fixed value for test run
+  int64_t padA[]        = {0, 0};
+  int64_t padA1[]        = {1, 1};
+  int64_t dilationA[] = {1, 1};
+  int64_t convstrideA[] = {1, 1};
+  int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+  // compute output from pad/stride/dilation
+  outdimA1[0] = dimA[0];
+  outdimA1[1] = filterdimA1[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+  }
+
+  outdimA2[0] = outdimA1[0];
+  outdimA2[1] = filterdimA2[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  outdimA3[0] = outdimA2[0];
+  outdimA3[1] = filterdimA3[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  // Create output tensor in the correct shape in pytorch's view
+  int64_t outdim1[]     = {0, 0, 0, 0};
+  int64_t outdim2[]     = {0, 0, 0, 0};
+  int64_t outdim3[]     = {0, 0, 0, 0};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 2;
+    axis[2] = 3;
+    axis[3] = 1;
+  }
+  for (int dim=0;dim<4;dim++) {
+    outdim1[dim] = outdimA1[axis[dim]];
+    outdim2[dim] = outdimA2[axis[dim]];
+    outdim3[dim] = outdimA3[axis[dim]];
+  }
+
+  // dconv3+drelu2+dscale2
+  at::Half* conv_in = inputs[13].data_ptr<at::Half>();
+  at::Half* dy3 = inputs[10].data_ptr<at::Half>();
+
+  DEBUG_MSG("[DEBUG] new dconv3 : " << inputs[10].to(at::kFloat).sum().item<float>());
+
+  // wgrad
+  auto wgrad3 = at::empty_like(inputs[3]);
+  at::Half* dw3 = wgrad3.data_ptr<at::Half>();
+  run_dconv(outdimA2,
+            padA,
+            convstrideA,
+            dilationA,
+            filterdimA3,
+            outdimA3,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw3,
+            dy3,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  auto grad_out2 = at::empty(outdim2, inputs[0].type(), output_format);
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+  at::Half* w = inputs[3].data_ptr<at::Half>();
+  at::Half* z = inputs[5].data_ptr<at::Half>();
+
+  at::Half* relu2 = inputs[13].data_ptr<at::Half>();
+
+  run_dconv_drelu_dscale(outdimA2,
+                         padA,
+                         convstrideA,
+                         dilationA,
+                         filterdimA3,
+                         outdimA3,
+                         CUDNN_DATA_HALF,
+                         dy2,
+                         w,
+                         dy3,
+                         z,
+                         relu2);
+
+  DEBUG_MSG("[DEBUG] new dconv2 : " << grad_out2.to(at::kFloat).sum().item<float>());
+
+  // dconv2+drelu1+dscale1
+  conv_in = inputs[12].data_ptr<at::Half>();
+
+  // wgrad
+  auto wgrad2 = at::empty_like(inputs[2]);
+  at::Half* dw2 = wgrad2.data_ptr<at::Half>();
+  run_dconv(outdimA1,
+            padA1,
+            convstrideA,
+            dilationA,
+            filterdimA2,
+            outdimA2,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw2,
+            dy2,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  auto grad_out1 = at::empty(outdim1, inputs[0].type(), output_format);
+  at::Half* dy1 = grad_out1.data_ptr<at::Half>();
+  w = inputs[2].data_ptr<at::Half>();
+  z = inputs[4].data_ptr<at::Half>();
+
+  at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+  // fused dgrad
+  run_dconv_drelu_dscale(outdimA1,
+                         padA1,
+                         convstrideA,
+                         dilationA,
+                         filterdimA2,
+                         outdimA2,
+                         CUDNN_DATA_HALF,
+                         dy1,
+                         w,
+                         dy2,
+                         z,
+                         relu1);
+
+/*
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (stride_1X1 != 1){
+    // dgrad
+    run_dconv(outdimA1,
+              padA1,
+              convstride1X1,
+              dilationA,
+              filterdimA2,
+              outdimA2,
+              CUDNN_DATA_HALF,
+              dy1,
+              w,
+              dy2,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+
+    // mul fused mask
+    grad_out1.mul_(inputs[15]);
+  }
+  else {
+    at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+    // fused dgrad
+    run_dconv_drelu_dscale(outdimA1,
+                           padA1,
+                           convstride1X1,
+                           dilationA,
+                           filterdimA2,
+                           outdimA2,
+                           CUDNN_DATA_HALF,
+                           dy1,
+                           w,
+                           dy2,
+                           z,
+                           relu1);
+  }
+*/
+  DEBUG_MSG("[DEBUG] new dconv1 : " << grad_out1.to(at::kFloat).sum().item<float>());
+
+  // create grads of conv4 that may exist
+  auto grad_x_conv4 = at::empty_like(inputs[0]);
+  at::Half* dx_conv4 = grad_x_conv4.data_ptr<at::Half>();
+  at::Tensor wgrad4;
+
+  // x used for dconv1 and dconv4 wgrad
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+    w = inputs[14].data_ptr<at::Half>();
+    at::Half* dy_conv4 = inputs[11].data_ptr<at::Half>();
+    if (requires_grad) {
+      run_dconv(dimA,
+                padA,
+                convstride1X1,
+                dilationA,
+                filterdimA4,
+                outdimA3,
+                CUDNN_DATA_HALF,
+                dx_conv4,
+                w,
+                dy_conv4,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx
+      // DEBUG_MSG("[DEBUG] new dx_identity : " << grad_x_conv4.to(at::kFloat).sum().item<float>());
+    }
+    // wgrad
+    wgrad4 = at::empty_like(inputs[14]);
+    at::Half* dw4 = wgrad4.data_ptr<at::Half>();
+    run_dconv(dimA,
+              padA,
+              convstride1X1,
+              dilationA,
+              filterdimA4,
+              outdimA3,
+              CUDNN_DATA_HALF,
+              x,
+              dw4,
+              dy_conv4,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+  }
+  else {
+    // if there is no downsample, dx_conv4 is fork of drelu3
+    dx_conv4 = inputs[11].data_ptr<at::Half>();
+  }
+
+  // dconv1+add
+  // wgrad
+  auto wgrad1 = at::empty_like(inputs[1]);
+  at::Half* dw1 = wgrad1.data_ptr<at::Half>();
+  run_dconv(dimA,
+            padA,
+            convstride1X1,
+            dilationA,
+            filterdimA1,
+            outdimA1,
+            CUDNN_DATA_HALF,
+            x,
+            dw1,
+            dy1,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  w = inputs[1].data_ptr<at::Half>();
+  auto grad_x = at::empty_like(inputs[0]);
+  at::Half* dx = grad_x.data_ptr<at::Half>();
+
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (requires_grad){
+    if (stride_1X1 != 1){
+      run_dconv(dimA,
+                padA,
+                convstride1X1,
+                dilationA,
+                filterdimA1,
+                outdimA1,
+                CUDNN_DATA_HALF,
+                dx,
+                w,
+                dy1,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // add 2 together
+      grad_x.add_(grad_x_conv4);
+    }
+    else {
+      run_dconv_add(dimA,
+                    padA,
+                    convstride1X1,
+                    dilationA,
+                    filterdimA1,
+                    outdimA1,
+                    CUDNN_DATA_HALF,
+                    dx,
+                    w,
+                    dy1,
+                    dx_conv4);
+    }
+  }
+
+  DEBUG_MSG("[DEBUG] new dx : " << grad_x.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad1 : " << wgrad1.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad2 : " << wgrad2.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad3 : " << wgrad3.to(at::kFloat).sum().item<float>());
+  outputs.push_back(grad_x);
+  outputs.push_back(wgrad1);
+  outputs.push_back(wgrad2);
+  outputs.push_back(wgrad3);
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    DEBUG_MSG("[DEBUG] new wgrad4 : " << wgrad4.to(at::kFloat).sum().item<float>());
+    outputs.push_back(wgrad4);
+  }
+
+  return outputs;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &bottleneck_forward, "Bottleneck block forward");
+  m.def("backward", &bottleneck_backward, "Bottleneck block backward");
+}

apex/apex/contrib/csrc/fmha/fmha_api.cpp

@@ -0,0 +1,305 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include "fmha.h"
+
+void run_fmha_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                  bool is_training,
+                                  cudaStream_t stream);
+void run_fmha_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                  bool is_training,
+                                  cudaStream_t stream);
+void run_fmha_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                  bool is_training,
+                                  cudaStream_t stream);
+void run_fmha_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                  bool is_training,
+                                  cudaStream_t stream);
+
+void run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                        cudaStream_t stream);
+void run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                        cudaStream_t stream);
+void run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                        cudaStream_t stream);
+void run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params,
+                                        cudaStream_t stream);
+
+void set_params(Fused_multihead_attention_fprop_params &params,
+                // sizes
+                const size_t b,
+                const size_t s,
+                const size_t h,
+                const size_t d,
+                // device pointers
+                void *qkv_packed_d,
+                void *cu_seqlens_d,
+                void *seqlens_d,
+                void *o_packed_d,
+                void *s_d,
+                float p_dropout) {
+
+    Data_type acc_type = DATA_TYPE_FP32;
+    Data_type data_type = DATA_TYPE_FP16;
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    // Set the pointers and strides.
+    params.qkv_ptr = qkv_packed_d;
+    params.qkv_stride_in_bytes = get_size_in_bytes(h * 3 * d, data_type);
+    params.o_ptr = o_packed_d;
+    params.o_stride_in_bytes = get_size_in_bytes(h * d, data_type);
+
+    params.cu_seqlens = static_cast<int *>(cu_seqlens_d);
+    params.seqlens = static_cast<int *>(seqlens_d);
+
+    // S = softmax(P)
+    params.s_ptr = s_d;
+    params.s_stride_in_bytes = get_size_in_bytes(b * h * s, data_type);
+
+    // Set the dimensions.
+    params.b = b;
+    params.h = h;
+    params.s = s;
+    params.d = d;
+
+    // Set the different scale values.
+    const float scale_bmm1 = 1.f / sqrtf(d);
+    constexpr float scale_softmax = 1.f;
+    constexpr float scale_bmm2 = 1.f;
+
+    set_alpha(params.scale_bmm1, scale_bmm1, acc_type);
+    set_alpha(params.scale_softmax, scale_softmax, acc_type);
+    set_alpha(params.scale_bmm2, scale_bmm2, data_type);
+
+    // Set this to probability of keeping an element to simplify things.
+    params.p_dropout = 1.f - p_dropout;
+    params.rp_dropout = 1.f / params.p_dropout;
+    TORCH_CHECK(p_dropout < 1.f);
+    set_alpha(params.scale_dropout, params.rp_dropout, data_type);
+}
+
+constexpr uint32_t NUM_HEADS_DIM = 2;
+constexpr uint32_t THREE_DIM = 1;
+
+std::vector<at::Tensor>
+mha_fwd(const at::Tensor &qkv,  // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+        const at::Tensor &cu_seqlens,  // b+1
+        const at::Tensor &seqlens,     // b
+        const float p_dropout,
+        const int max_seq_len,
+        const bool is_training,
+        c10::optional<at::Generator> gen_) {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    TORCH_CHECK(dprops->major == 8 && dprops->minor == 0);
+    int seq_len = 512;
+    auto launch = &run_fmha_fp16_512_64_sm80;
+    if( max_seq_len <= 128 ) {
+        seq_len = 128;
+        launch = &run_fmha_fp16_128_64_sm80;
+    } else if( max_seq_len <= 256 ) {
+        seq_len = 256;
+        launch = &run_fmha_fp16_256_64_sm80;
+    } else if( max_seq_len <= 384 ) {
+        seq_len = 384;
+        launch = &run_fmha_fp16_384_64_sm80;
+    } else if( max_seq_len <= 512 ) {
+        seq_len = 512;
+        launch = &run_fmha_fp16_512_64_sm80;
+    } else {
+        TORCH_CHECK(false);
+    }
+
+    constexpr int warps_m = 1;
+    constexpr int warps_n = 4;  // this leads to an upper bound
+    const int mmas_m = seq_len / 16 / warps_m;
+    const int mmas_n = seq_len / 16 / warps_n;
+    
+    const int elts_per_thread = 8 * mmas_m * mmas_n;
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.dtype() == torch::kFloat16);
+    TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32);
+    TORCH_CHECK(seqlens.dtype() == torch::kInt32);
+
+    TORCH_CHECK(qkv.is_cuda())
+    TORCH_CHECK(cu_seqlens.is_cuda())
+
+    TORCH_CHECK(qkv.is_contiguous())
+    TORCH_CHECK(cu_seqlens.is_contiguous())
+    TORCH_CHECK(seqlens.is_contiguous())
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+    TORCH_CHECK(seqlens.dim() == 1);
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    TORCH_CHECK(seqlens.numel() == batch_size);
+    const int total = sizes[0];
+    const int num_heads = sizes[NUM_HEADS_DIM];
+    const int head_size = sizes[3];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+    auto opts = qkv.options();
+
+    auto ctx = torch::empty({ total, num_heads, head_size }, opts);
+
+    auto s = torch::empty({ batch_size, num_heads, seq_len, seq_len }, opts);
+
+    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
+        gen_, at::cuda::detail::getDefaultCUDAGenerator());
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               seqlens.data_ptr(),
+               ctx.data_ptr(),
+               s.data_ptr(),
+               p_dropout);
+
+    // number of times random will be generated per thread, to offset philox counter in thc random
+    // state
+    int64_t counter_offset = elts_per_thread;
+    at::PhiloxCudaState rng_engine_inputs;
+
+    if( is_training ) {
+        // See Note [Acquire lock when using random generators]
+        std::lock_guard<std::mutex> lock(gen->mutex_);
+        params.philox_args = gen->philox_cuda_state(counter_offset);
+    }
+
+    launch(params, is_training, stream);
+
+    return { ctx, s };
+}
+
+std::vector<at::Tensor>
+mha_bwd(const at::Tensor &dout,  // total x num_heads, x head_size
+        const at::Tensor &qkv,   // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+        at::Tensor &softmax,     // b x h x s x s softmax and dmask - will be overwritten with dP
+        const at::Tensor &cu_seqlens,  // b+1
+        const at::Tensor &seqlens,     // b
+        const float p_dropout,         // probability to drop
+        const int max_seq_len          // max sequence length to choose the kernel
+) {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    TORCH_CHECK(dprops->major == 8 && dprops->minor == 0);
+    int seq_len = 512;
+    auto launch = &run_fmha_dgrad_fp16_512_64_sm80;
+    if( max_seq_len <= 128 ) {
+        seq_len = 128;
+        launch = &run_fmha_dgrad_fp16_128_64_sm80;
+    } else if( max_seq_len <= 256 ) {
+        seq_len = 256;
+        launch = &run_fmha_dgrad_fp16_256_64_sm80;
+    } else if( max_seq_len <= 384 ) {
+        seq_len = 384;
+        launch = &run_fmha_dgrad_fp16_384_64_sm80;
+    } else if( max_seq_len <= 512 ) {
+        seq_len = 512;
+        launch = &run_fmha_dgrad_fp16_512_64_sm80;
+    } else {
+        TORCH_CHECK(false);
+    }
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.dtype() == torch::kFloat16);
+    TORCH_CHECK(dout.dtype() == torch::kFloat16);
+    TORCH_CHECK(softmax.dtype() == torch::kFloat16);
+    TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32);
+    TORCH_CHECK(seqlens.dtype() == torch::kInt32);
+
+    TORCH_CHECK(qkv.is_cuda());
+    TORCH_CHECK(cu_seqlens.is_cuda());
+
+    TORCH_CHECK(qkv.is_contiguous());
+    TORCH_CHECK(cu_seqlens.is_contiguous());
+    TORCH_CHECK(seqlens.is_contiguous());
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+    TORCH_CHECK(seqlens.dim() == 1);
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    TORCH_CHECK(seqlens.numel() == batch_size);
+    const int num_heads = sizes[NUM_HEADS_DIM];
+    const int head_size = sizes[3];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+
+    auto dqkv = torch::empty_like(qkv);
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               seqlens.data_ptr(),
+               dout.data_ptr(),     // we set o_ptr to dout
+               softmax.data_ptr(),  // softmax gets overwritten by dP!
+               p_dropout);
+
+    // we're re-using these scales scales
+    Data_type acc_type = DATA_TYPE_FP32;
+    set_alpha(params.scale_bmm1, 1.f, acc_type);
+    set_alpha(params.scale_softmax, 1.f / sqrtf(head_size), acc_type);
+    set_alpha(params.scale_bmm2, 1.f, DATA_TYPE_FP16);
+    params.dqkv_ptr = dqkv.data_ptr();
+
+    launch(params, stream);
+    return { dqkv, softmax };
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "Fused Multi-head Self-attention for BERT";  
+    m.def("fwd", &mha_fwd, "Forward pass");
+    m.def("bwd", &mha_bwd, "Backward pass");
+}

apex/apex/contrib/csrc/fmha/src/fmha.h

@@ -0,0 +1,92 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda.h>
+#include <vector>
+
+#include <ATen/CUDAGeneratorImpl.h>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+
+#include <fmha_utils.h>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Qkv_params {
+    // The QKV matrices.
+    void *qkv_ptr;
+
+    // The stride between rows of the Q, K and V matrices.
+    size_t qkv_stride_in_bytes;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Fused_multihead_attention_fprop_params : public Qkv_params {
+
+    // The dQKV matrices.
+    void *dqkv_ptr;
+
+    // The O matrix (output).
+    void *o_ptr;
+
+    // The stride between rows of O.
+    int64_t o_stride_in_bytes;
+
+    // The pointer to the S matrix, overwritten by the dP matrix (bwd).
+    void *s_ptr;
+    // The stride between rows of the S matrix.
+    int64_t s_stride_in_bytes;
+
+    // The dimensions.
+    int b, h, s, d;
+
+    // The scaling factors for the kernel.
+    uint32_t scale_bmm1, scale_softmax, scale_bmm2;
+
+    // array of length b+1 holding starting offset of each sequence.
+    int *cu_seqlens;
+
+    // array of length b holding the actual sequence lenghts.
+    int *seqlens;
+
+    // The dropout probability (probability of keeping an activation).
+    float p_dropout;
+
+    // Scale factor of 1 / (1 - p_dropout).
+    float rp_dropout;
+
+    // Scale factor of 1 / (1 - p_dropout), in half2.
+    uint32_t scale_dropout;
+
+    // Random state.
+    at::PhiloxCudaState philox_args;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+

apex/apex/contrib/csrc/fmha/src/fmha/gemm.h

@@ -0,0 +1,317 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <fmha/utils.h>
+
+#define FMHA_DIV_UP(m, n) (((m) + (n)-1) / (n))
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Data_type_, int NUM_ELTS_, int BITS_PER_ELT_, int ALIGNMENT_ >
+struct Fragment_base_ {
+
+    // The data type.
+    using Data_type = Data_type_;
+    // default input type
+    using Input_type_ = Data_type_;
+    // Does it store the array of elements.
+    enum { HAS_ELTS = BITS_PER_ELT_ >= 8 };
+    // The number of elements.
+    enum { NUM_ELTS = NUM_ELTS_ };
+    // The size of element in bits.
+    enum { BITS_PER_ELT = BITS_PER_ELT_ };
+    // The size of byte of a single register.
+    enum { BYTES_PER_REG = 4 };
+    // The size in bits.
+    enum { BITS_PER_REG = BYTES_PER_REG * 8 };
+    // The number of registers needed to store the fragment.
+    enum { NUM_REGS = Div_up<NUM_ELTS * BITS_PER_ELT, BITS_PER_REG>::VALUE };
+    // The size in bytes (as returned by sizeof(Fragment_base<>).
+    enum { SIZE_IN_BYTES = NUM_REGS * BYTES_PER_REG };
+    // The alignment.
+    enum { ALIGNMENT = ALIGNMENT_ > 0 ? ALIGNMENT_ : Min<NUM_REGS * BYTES_PER_REG, 16>::VALUE };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The type of the elements.
+    typename Data_type_,
+    // The number of elements.
+    int NUM_ELTS_,
+    // The alignment if you want to force a value -- use 0 otherwise.
+    int ALIGNMENT_ = 0,
+    // The base class.
+    typename Base_ = Fragment_base_<Data_type_, NUM_ELTS_, 8 * sizeof(Data_type_), ALIGNMENT_>
+>
+struct alignas(static_cast<int>(Base_::ALIGNMENT)) Fragment : public Base_ {
+
+    // The size of a load/store.
+    enum { BYTES_PER_LOAD_STORE = Base_::NUM_REGS * sizeof(uint32_t) };
+
+    // Clear the fragment. Using PTX in that code seems to produce better SASS...
+    inline __device__ void clear() {
+        #pragma unroll
+        for( int ii = 0; ii < Base_::NUM_REGS; ++ii ) {
+            asm volatile("mov.u32 %0, 0; \n" : "=r"(this->reg(ii)) : );
+        }
+    }
+
+    // Immutable access to a register.
+    inline __device__ const uint32_t& reg(int ii) const {
+        return this->regs_[ii];
+    }
+
+    // Mutable access to a register.
+    inline __device__ uint32_t& reg(int ii) {
+        return this->regs_[ii];
+    }
+
+    uint32_t regs_[Base_::NUM_REGS];
+
+    // Immutable access to the elements.
+    inline __device__ const Data_type_& elt(int ii) const {
+        return reinterpret_cast<const Data_type_*>(&this->regs_[0])[ii];
+    }
+
+    // Mutable access to the elements.
+    inline __device__ Data_type_& elt(int ii) {
+        return reinterpret_cast<Data_type_*>(&this->regs_[0])[ii];
+    }
+
+    // Immutable access to the elements with a cast.
+    template< typename Cast_type >
+    inline __device__ const Cast_type& elt_as(int ii) const {
+        return reinterpret_cast<const Cast_type*>(&this->regs_[0])[ii];
+    }
+
+    // Mutable access to the elements.
+    template< typename Cast_type >
+    inline __device__ Cast_type& elt_as(int ii) {
+        return reinterpret_cast<Cast_type*>(&this->regs_[0])[ii];
+    }
+
+    // Add another fragment.
+    inline __device__ void add(const Fragment &other) {
+        #pragma unroll
+        for( int ii = 0; ii < NUM_ELTS_; ++ii ) {
+            this->elt(ii) += other.elt(ii);
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Layout >
+struct Fragment_a : public Fragment<uint16_t, 8> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Layout >
+struct Fragment_b : public Fragment<uint16_t, 8> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Fragment_accumulator : public Fragment<float, 8> {
+
+    // The base class.
+    using Base = Fragment<float, 8>;
+
+    // Add two fragments.
+    template< typename Other_fragment_ >
+    inline __device__ void add(const Other_fragment_ &other) {
+        for( int ii = 0; ii < Base::NUM_ELTS; ++ii ) {
+            this->elt(ii) = this->elt(ii) + other.elt(ii);
+        }
+    }
+
+    // Do the HMMA.
+    template< typename Layout_a, typename Layout_b >
+    inline __device__ void mma(const Fragment_a<Layout_a> &a,
+                               const Fragment_b<Layout_b> &b) {
+        asm volatile( \
+            "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" \
+            "    {%0, %1, %2, %3}, \n" \
+            "    {%4, %5, %6, %7}, \n" \
+            "    {%8, %9}, \n" \
+            "    {%0, %1, %2, %3}; \n" \
+                    : "+f"(  elt(0)), "+f"(  elt(1)), "+f"(  elt(2)), "+f"(  elt(3))
+                    :  "r"(a.reg(0)),  "r"(a.reg(1)),  "r"(a.reg(2)),  "r"(a.reg(3))
+                    ,  "r"(b.reg(0)),  "r"(b.reg(1)));
+        asm volatile( \
+            "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" \
+            "    {%0, %1, %2, %3}, \n" \
+            "    {%4, %5, %6, %7}, \n" \
+            "    {%8, %9}, \n" \
+            "    {%0, %1, %2, %3}; \n" \
+                    : "+f"(  elt(4)), "+f"(  elt(5)), "+f"(  elt(6)), "+f"(  elt(7))
+                    :  "r"(a.reg(0)),  "r"(a.reg(1)),  "r"(a.reg(2)),  "r"(a.reg(3))
+                    ,  "r"(b.reg(2)),  "r"(b.reg(3)));
+    }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Fragment, int M, int N >
+inline __device__ void clear(Fragment (&frag)[M][N]) {
+    #pragma unroll
+    for( int mi = 0; mi < M; ++mi ) {
+        #pragma unroll
+        for( int ni = 0; ni < N; ++ni ) {
+            frag[mi][ni].clear();
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Accumulator_type, int WARPS_K >
+struct Clear_accumulator {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int WARPS_K >
+struct Clear_accumulator<float, WARPS_K> {
+  template< typename Acc, int M, int N >
+  static inline __device__ void apply(Acc (&acc)[M][N], bool = false) {
+    fmha::clear(acc);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Acc, typename A, typename B, int M, int N>
+inline __device__ void gemm(Acc (&acc)[M][N], const A (&a)[M], const B (&b)[N]) {
+
+    #pragma unroll
+    for( int mi = 0; mi < M; ++mi ) {
+        #pragma unroll
+        for( int ni = 0; ni < N; ++ni ) {
+            acc[mi][ni].mma(a[mi], b[ni]);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The number of rows in the CTA tile.
+    int M_,
+    // The number of cols in the CTA tile.
+    int N_,
+    // The number of elements in the the K dimension of the GEMM loop.
+    int K_,
+    // The number of rows of warps.
+    int WARPS_M_,
+    // The number of cols of warps.
+    int WARPS_N_,
+    // The number of warps in the K dimension of the GEMM loop.
+    int WARPS_K_>
+struct Cta_tile_ {
+
+    enum { M = M_, N = N_, K = K_ };
+    // The number of warps.
+    enum { WARPS_M = WARPS_M_, WARPS_N = WARPS_N_, WARPS_K = WARPS_K_ };
+    // The number of warps per CTA.
+    enum { WARPS_PER_CTA = WARPS_M * WARPS_N * WARPS_K };
+    // The number of threads per warp.
+    enum { THREADS_PER_WARP = 32 };
+    // The number of threads per CTA.
+    enum { THREADS_PER_CTA = WARPS_PER_CTA * THREADS_PER_WARP };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile>
+struct Hmma_tile {
+    // The number of elements computed with a single warp-MMA.
+    enum { M_PER_MMA = 16, N_PER_MMA = 16, K_PER_MMA = 16 };
+
+    // The number of elements computed with a single CTA-MMA.
+    enum {
+        M_PER_MMA_PER_CTA = M_PER_MMA * Cta_tile::WARPS_M,
+        N_PER_MMA_PER_CTA = N_PER_MMA * Cta_tile::WARPS_N,
+        K_PER_MMA_PER_CTA = K_PER_MMA * Cta_tile::WARPS_K
+    };
+
+    // The number of MMAs needed to compute the GEMM.
+    enum {
+        MMAS_M = Div_up<Cta_tile::M, M_PER_MMA_PER_CTA>::VALUE,
+        MMAS_N = Div_up<Cta_tile::N, N_PER_MMA_PER_CTA>::VALUE,
+        MMAS_K = Div_up<Cta_tile::K, K_PER_MMA_PER_CTA>::VALUE,
+    };
+
+    // The number of elements computed per warp.
+    enum {
+        M_PER_WARP = MMAS_M * M_PER_MMA,
+        N_PER_WARP = MMAS_N * N_PER_MMA,
+        K_PER_WARP = MMAS_K * K_PER_MMA,
+    };
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using A_type = uint16_t;
+using B_type = uint16_t;
+using C_type = uint16_t;
+using Accumulator_type = float;
+using Epilogue_type = float;
+
+constexpr int BITS_PER_ELEMENT_A = sizeof(A_type) * 8;
+constexpr int BITS_PER_ELEMENT_B = sizeof(B_type) * 8;
+constexpr int BITS_PER_ELEMENT_C = sizeof(C_type) * 8;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int M, int N, int K, int WARPS_M, int WARPS_N, int WARPS_K>
+using Cta_tile_extd = Cta_tile_<M, N, K, WARPS_M, WARPS_N, WARPS_K>;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile_>
+using Cta_tile_with_k_with_padding = Cta_tile_extd<Cta_tile_::M,
+                                                   Cta_tile_::N,
+                                                   Next_power_of_two<Cta_tile_::K>::VALUE,
+                                                   Cta_tile_::WARPS_M,
+                                                   Cta_tile_::WARPS_N,
+                                                   Cta_tile_::WARPS_K>;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha