Repository: pytorch
Issue: 3D grouped & depthwise convolution very slow on backward pass
Labels: module: cudnn,module: cuda,triaged

Hi,

It seems that 3D grouped & depthwise convolution is very slow on the **backward** pass. Backward pass on depthwise convolution takes about 10 times the time of a standard 3D convolution's forward pass.

I understand cudnn does not yet support fp32 depthwise convolution (ie forward pass is slower than standard convolution by about 2 times or so for 2D convolution). However, this problem seems to be significantly worse for 3D convolution's _backward_ pass. Reading cudnn docs, it seems 3D depthwise isn't supported. 

Any suggestions?

```python
import torch
import torch.nn as nn
import time

# Timing code taken from https://github.com/facebookresearch/pycls/blob/master/pycls/utils/benchmark.py

class Timer(object):
    
    def __init__(self):
        self.total_time = 0.0
        self.calls = 0
        self.start_time = 0.0
        self.diff = 0.0
        self.average_time = 0.0

    def tic(self):
        self.start_time = time.time()

    def toc(self):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        
        
def compute_precise_time(model, input_size, loss_fun, batch_size, num_iter=3, label_dtype=torch.int64):
    """Computes precise time."""
    # Generate a dummy mini-batch
    inputs = torch.rand(batch_size, *input_size)
    labels = torch.zeros(batch_size, dtype=label_dtype)
    # Copy the data to the GPU
    inputs = inputs.cuda(non_blocking=False)
    labels = labels.cuda(non_blocking=False)
    
    # Compute precise time
    fw_timer = Timer()
    bw_timer = Timer()
    model.train()
    for _cur_iter in range(num_iter):
        # Forward
        fw_timer.tic()
        preds = model(inputs)
        loss = loss_fun(preds, labels)
        torch.cuda.synchronize()
        fw_timer.toc()
        # Backward
        bw_timer.tic()
        loss.backward()
        torch.cuda.synchronize()
        bw_timer.toc()
    torch.cuda.synchronize()
    
    print({"prec_train_fw_time": fw_timer.average_time, "prec_train_bw_time": bw_timer.average_time})


class GroupedNet3D(nn.Module):
    
    def __init__(self, data_size, layer_channel, groups):
        
        super(GroupedNet3D, self).__init__()
        
        if groups is None:
            groups = layer_channel # depth wise
            
        self.f = nn.Sequential(
            nn.Conv3d(data_size[0], layer_channel, 3, bias=False),
            nn.Conv3d(layer_channel, layer_channel, 3, groups=groups, bias=False),
            nn.Conv3d(layer_channel, layer_channel, 3, groups=groups, bias=False),
            nn.Conv3d(layer_channel, layer_channel, 3, groups=groups, bias=False),
            nn.Conv3d(layer_channel, layer_channel, 3, groups=groups, bias=False),
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(layer_channel, 1)
        )
    
    def forward(self, x):
        return self.f(x)
    
loss_fun = nn.CrossEntropyLoss().cuda()
data_size = (1, 50, 50, 50)
batch_size = 16
channels = 32
Model = GroupedNet3D


# standard 3D convolution
model = Model(data_size, channels, 1).cuda()
compute_precise_time(model, data_size, loss_fun, batch_size)
# {'prec_train_fw_time': 0.1468369960784912, 'prec_train_bw_time': 0.3601357142130534}

# grouped convolution
model = Model(data_size, channels, 8).cuda()
compute_precise_time(model, data_size, loss_fun, batch_size)
# {'prec_train_fw_time': 0.3155516783396403, 'prec_train_bw_time': 1.0136785507202148}

# depth wise convolution
model = Model(data_size, channels, None).cuda()
compute_precise_time(model, data_size, loss_fun, batch_size)
# {'prec_train_fw_time': 0.41037146250406903, 'prec_train_bw_time': 1.7994076410929363}

torch.backends.cudnn.version() # 7603
torch.version.cuda # '10.1'
torch.__version__ # '1.4.0'
```

cc @ngimel @csarofeen @ptrblck