Tf.extract_image_patches in pytorch

xiao1228: <p>Is there a function like tf.extract_image_patches in pytorch?<br>
Thank you</p>

ptrblck: <p>I think <code>tensor.unfold</code> should correspond to this method.</p>

xiao1228: <p>is it support for overlapping as well?<br>
Because i have a tensor size (128, 32, 32, 16)<br>
with ksize = [1,3,3,1] stride [1,1,1,1] and rate [1,1,1,1]<br>
In tensorflow it can gives me (128,32,32,144)<br>
However, I cant figure out a way in Pytorch using unfold to give the same.<br>
Thank you</p>

ptrblck: <p>Could you explain how the output is calculated in tensorflow?<br>
If looks like you somehow reshaped the patches to the original input size (<code>128x32x32</code>)?</p>

xiao1228: <p>This is the output tensor from a convolution layer, 128 is batch size of 32*32 and output channel of that layer is 16.  it is cifar10 with resnet56.</p>

ptrblck: <p><code>tensor.unfold</code> will give you the image patches, so I assume you are applying some operation on each patch and reshape it to this activation shape afterwards?<br>
What operation would you like to apply? Maybe a standard <code>nn.Conv2d</code> module will also do the trick?</p>

xiao1228: <p>In tensorflow, it passes this input tensor (128,32,32,16) into the tf.extract_image_patches and output a tensor with size (128,32,32,144) with kernel size (1,3,3,1) and stride (1,1,1,1) as the parameter for the tf.extract_image_patches. so i assume in the tf.extract_image_patches it do overlap sampling. And i tried tensor.unfold cannot perform this. Maybe it can and I dont know how it will output the same as tensorflow…<br>
Using a nn.Conv2d is an idea, by setting the weights to 1? but it misses the sampling part…<br>
Thank you again…</p>

ptrblck: <p>I’m not sure why the method is called <code>extract_image_patches</code> if you won’t get the patches, but apparently a view of <code>[batch_size, height, width, channels*kernel_height*kernel_width]</code>.</p>
<p>However, this code should yield the same result in PyTorch:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as F

batch_size = 128
channels = 16
height, width = 32, 32
x = torch.randn(batch_size, channels, height, width)

kh, kw = 3, 3
dh, dw = 1, 1

# Pad tensor to get the same output
x = F.pad(x, (1, 1, 1, 1))

# get all image windows of size (kh, kw) and stride (dh, dw)
patches = x.unfold(2, kh, dh).unfold(3, kw, dw)
print(patches.shape)  # [128, 16, 32, 32, 3, 3]
# Permute so that channels are next to patch dimension
patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()  # [128, 32, 32, 16, 3, 3]
# View as [batch_size, height, width, channels*kh*kw]
patches = patches.view(*patches.size()[:3], -1)
print(patches.shape)
&gt; torch.Size([128, 32, 32, 144])
</code></pre>
<p>Note that in PyTorch the channel dimension is in dim1, so I changed your input shape to match the PyTorch conversion. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>
<pre><code class="lang-auto"></code></pre>

xiao1228: <p>Thank you so much for the reply! Really appreciate it! I will give a try! Thanks again!!</p>

FloCF: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> for your solution.<br>
Put it into a function with usual pytorch format BxCxHxW which sould mimic tf function. Since I am a beginner, there is probably a much more efficient way. Also I haven’t tried backprop yet…</p>
<pre><code class="lang-auto">import math
import torch.nn.functional as F

def extract_image_patches(x, kernel, stride=1, dilation=1):
    # Do TF 'SAME' Padding
    b,c,h,w = x.shape
    h2 = math.ceil(h / stride)
    w2 = math.ceil(w / stride)
    pad_row = (h2 - 1) * stride + (kernel - 1) * dilation + 1 - h
    pad_col = (w2 - 1) * stride + (kernel - 1) * dilation + 1 - w
    x = F.pad(x, (pad_row//2, pad_row - pad_row//2, pad_col//2, pad_col - pad_col//2))
    
    # Extract patches
    patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
    patches = patches.permute(0,4,5,1,2,3).contiguous()
    
    return patches.view(b,-1,patches.shape[-2], patches.shape[-1])
</code></pre>

xuef: <aside class="quote no-group" data-username="FloCF" data-post="10" data-topic="43837">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/flocf/48/13467_2.png" class="avatar"> FloCF:</div>
<blockquote>
<p><code>x = F.pad(x, (pad_row//2, pad_row - pad_row//2, pad_col//2, pad_col - pad_col//2))</code></p>
</blockquote>
</aside>
<p>Functional pad takes different order, should be:<br>
<code>x = F.pad(x, (pad_col//2, pad_col - pad_col//2, pad_row//2, pad_row - pad_row//2, ))</code></p>

Debt: <p>Hello, I used the above code, but during training, the train loss and val loss did not change</p>