Manual Implementation of Unrolled 3D Convolutions

username: cll
cooked: <p>Given that <em>torch.nn.Unfold</em> can be used to unroll 2D convolutions, so that they can be computed using Vector Matrix Multiplication (VMMs), and that the same unrolling approach can be used to compute 3D convolutions as VMMs (described in <a href="https://www.mdpi.com/2079-9292/8/1/65/pdf" rel="nofollow noopener">https://www.mdpi.com/2079-9292/8/1/65/pdf</a>),  how can PyTorch be used to unroll 3D convolutions?</p>
<p>There is currently no native implementation of <em>torch.nn.Unfold</em> for 5D tensors, however, a response on <a href="https://github.com/pytorch/pytorch/issues/30798" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/30798</a> indicates that <em>x.unfold(2, k, s).unfold(3, k, s).unfold(4, k, s)</em> can be used. Unfortunately, using this approach, I can’t seem to get the same result as using <em>torch.nn.Conv3d</em> when the same filter weights are used and biases are disabled.</p>
<p>Any help would be greatly appreciated!</p>

username: ptrblck
cooked: <p>Could you post your code so that we could have a look, why it might give you a different result? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: cll
cooked: <p>Sure! Please see the code below. The 2D Convolution block appears to work well. I have since managed to implement the 3D Convolution block using a similar approach, however, it only appears to function when padding=(0, 0, 0). I’m using <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.unfold" rel="nofollow noopener">torch.Tensor.unfold</a> to unfold 5D input tensors. Unfortunately, it does not have a padding argument.</p>
<p>Do you have any idea on how I could extrapolate the 3D Convolution block below to support arbitrary padding lengths?</p>
<p><strong>UPDATE</strong> I have managed to use <a href="https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.pad" rel="nofollow noopener">torch.nn.functional.pad</a> to manually zero-pad the input before unfolding it, and it appears to be working for arbitrary padding lengths. I have updated the code below.</p>
<p>I believe this has now been resolved <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<pre><code class="lang-auto">import torch
from torch import nn
import torch.nn.functional as F


# 2D Convolution
input_dim = [4, 4]
in_channels = 2
out_channels = 4
kernel_size = (3, 3)
padding = (1, 1)
stride = (1, 1)
input_tensor = torch.zeros(1, in_channels, input_dim[0], input_dim[1]).uniform_(-1, 1)
conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding, stride=stride, bias=False)
output_tensor = conv(input_tensor)
output_dim = [0, 0]
output_dim[0] = int((input_tensor.shape[2] - kernel_size[0] + 2 * padding[0]) / stride[0]) + 1
output_dim[1] = int((input_tensor.shape[3] - kernel_size[1] + 2 * padding[1]) / stride[1]) + 1
unfolded_input_tensor = F.unfold(input_tensor, kernel_size=kernel_size, padding=padding, stride=stride)
kernels_flat = conv.weight.detach().clone().view(out_channels, -1)
alt_output_tensor = torch.matmul(kernels_flat, unfolded_input_tensor).view(1, out_channels, output_dim[0], output_dim[1])
print(torch.all(torch.isclose(output_tensor, alt_output_tensor)))

# 3D Convolution
input_dim = [22, 59, 114]
in_channels = 1
out_channels = 16
kernel_size = (22, 5, 5)
padding = (1, 2, 2)
stride = (1, 2, 2)
input_tensor = torch.zeros(1, in_channels, input_dim[0], input_dim[1], input_dim[2]).uniform_(-1, 1)
conv = nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding, stride=stride, bias=False)
output_tensor = conv(input_tensor)
output_dim = [0, 0, 0]
output_dim[0] = int((input_tensor.shape[2] - kernel_size[0] + 2 * padding[0]) / stride[0]) + 1
output_dim[1] = int((input_tensor.shape[3] - kernel_size[1] + 2 * padding[1]) / stride[1]) + 1
output_dim[2] = int((input_tensor.shape[4] - kernel_size[2] + 2 * padding[2]) / stride[2]) + 1
if not all(item == 0 for item in padding):
    input_tensor = F.pad(input_tensor, pad=(padding[2], padding[2], padding[1], padding[1], padding[0], padding[0]), mode="constant", value=0)

unfolded_input_tensor = input_tensor.unfold(2, kernel_size[0], stride[0]).unfold(3, kernel_size[1], stride[1]).unfold(4, kernel_size[2], stride[2])
unfolded_input_tensor = unfolded_input_tensor.reshape(-1, kernel_size[0] * kernel_size[1] * kernel_size[2])
kernels_flat = conv.weight.detach().clone().view(out_channels, -1)
alt_output_tensor = torch.matmul(kernels_flat, unfolded_input_tensor.T).view(1, out_channels, output_dim[0], output_dim[1], output_dim[2])
print(torch.all(torch.isclose(output_tensor, alt_output_tensor)))
</code></pre>

username: johnson111788
cooked: <p>Thanks for sharing your code which helps a lot!<br>
However, I met some problems when I trying to modifying the in_channel.<br>
For example, I set<br>
input_dim = [5, 7, 7]<br>
in_channels = 2<br>
out_channels = 4<br>
kernel_size = (4, 6, 6)<br>
padding = (1, 1, 1)<br>
stride = (1, 1, 1)<br>
then it shows</p>
<p><code>RuntimeError: size mismatch, m1: [4 x 288], m2: [144 x 128] at /Users/distiller/project/conda/conda-bld/pytorch_1570710797334/work/aten/src/TH/generic/THTensorMath.cpp:197</code></p>
<p>then I set<br>
<code>unfolded_input_tensor = unfolded_input_tensor.reshape(-1, in_channels * kernel_size[0] * kernel_size[1] * kernel_size[2])</code><br>
since the kernel shape is <code>out_c*in_c*k1*k2*k3</code><br>
Eventhough I pass the error, the final result is tensor(False)</p>
<p>I would like to know how to unroll Conv3d in multi-input channels.<br>
Thanks for your time to answer the question!</p>

username: Youngseok0001
cooked: <p>The above code does not seem to work when the in_channel is bigger than 1.</p>
<p>Here is my approach which seems to work well.</p>
<pre><code class="lang-python">
    import numpy as np
    import torch
    from torch import nn
    import torch.nn.functional as F

    inputShape = [128, 128, 128]    
    batchSioze = 2
    CIn        = 4
    COut       = 8    
    kernelSize = (10,5,3)
    pad        = (2,3,1)
    stride     = (1,2,1)

    # normal conv
    conv = nn.Conv3d(CIn, COut, kernelSize, stride, pad, bias=False).cuda()
                      
    # alternativeConv
    def alternativeConv(X, K, 
                        COut       = None,
                        kernelSize = (3,3,3),
                        pad        = (1,1,1),
                        stride     = (1,1,1) ):

        def unfold3d(tensor, kernelSize, pad, stride): 

            B, C, _, _, _ = tensor.shape

            # Input shape: (B, C, D, H, W)
            tensor = F.pad(tensor,
                           (pad[2], pad[2],
                            pad[1], pad[1],
                            pad[0], pad[0])
                          )

            tensor = (tensor
                      .unfold(2, size=kernelSize[0], step=stride[0])
                      .unfold(3, size=kernelSize[1], step=stride[1])
                      .unfold(4, size=kernelSize[2], step=stride[2])
                      .permute(0, 2, 3, 4, 1, 5, 6, 7)
                      .reshape(B, -1, C * np.prod(kernelSize))
                      .transpose(1, 2)
                     )
            
            return tensor
    
        B,CIn,H,W,D = X.shape
        outShape = ( (np.array([H,W,D]) - np.array(kernelSize) + 2 * np.array(pad)) / np.array(stride) ) + 1
        outShape = outShape.astype(np.int32)
        
        X = unfold3d(X, kernelSize, pad, stride)
  
        K = K.view(COut, -1)
        #K = torch.randn(COut, CIn, *kernelSize).cuda() 
        #K = K.view(COut, -1)
                    
        Y = torch.matmul(K, X).view(B, COut, *outShape)
        
        return Y
    
    X = torch.randn(batchSioze, CIn, *inputShape).cuda()
    
    Y1 = conv(X)
    
    Y2 = alternativeConv(X, conv.weight, 
                         COut       = COut,
                         kernelSize = kernelSize,
                         pad        = pad,
                         stride     = stride
                         )
    
    print(torch.all(torch.isclose(Y1, Y2)))   
</code></pre>

username: f-dangel
cooked: <p>Hi,</p>
<p>not sure if you’re still looking for a solution, but I have recently written a small package named <a href="https://github.com/f-dangel/unfoldNd" rel="noopener nofollow ugc">unfoldNd</a> that should do what you want.</p>
<p><code>unfoldNd.UnfoldNd</code> works exactly like <code>torch.nn.Unfold</code>, also for 5d inputs. <a href="https://github.com/f-dangel/unfoldNd/blob/main/examples/example.py" rel="noopener nofollow ugc">Here</a> is a small demo.</p>
<p>Hope it’s useful.</p>