Is it possible to write this function in more compact/efficient way?

Hi everyone

I am swithcing from programin in matlab to pytorch and yet there are alot to learn for me.
I have a function as following, it takes an input in 4D, bachsize x channels x dimention1 x dimention2 and
do some operations on it and gives the output in the same size.
I will really appreciate it if you guys can help me and tell me if there is a more efficient and simple way to write this function in pytorch (so i can get rid of the loops) :slight_smile:

Thanks in advance …

def EdgeDetection(Input_Temp):

K1 = torch.FloatTensor([[1,0,0],[0,-1,0],[0,0,0]]).to(device)
K2 = torch.FloatTensor([[0,1,0],[-1,0,0],[0,0,0]]).to(device)

# print "The input size is" , Input_Temp.size() ,  "\n" #The input size is for example (128, 3, 32, 32)

D = 0 # D is responsible for number of Batches  or the 0st dimention
for BatchNum in range(Input_Temp.size(0)):
    C = 0 # C is responsible for number of channels or the 1st dimention 
    for ChannelNum in range(Input_Temp.size(1)):

        TempImg = Input_Temp[BatchNum,ChannelNum,:,:].clone()
        TempImgNormalize = TempImg

        # Makeing the input to have size of d x c x m x m so it can be used in conv2d
        A1 = torch.unsqueeze(torch.unsqueeze(TempImgNormalize,0),0).to(device)
        Gx = F.conv2d(A1, torch.unsqueeze(torch.unsqueeze(K1,0),0),padding= K1.size(0)/2 )
        Gy = F.conv2d(A1, torch.unsqueeze(torch.unsqueeze(K2,0),0),padding= K2.size(0)/2 )
        GM = (Gx**2 + Gy**2)**0.5


        # concatinating the results, concatenating along the 1st dimention:
        if C == 0:
            APP = GM 
        else:
            APP = torch.cat((APP,GM),1)
        C += 1

    # concatinating the results, concatenating along the 0st dimention:
    if D == 0:
        Final = APP
        Final = Final.to(device)
    else:
        Final = torch.cat((Final,APP),0)
    D = D + 1

return Final

Hi,

Please find below code that does it with 2 convs.
I did not do timings as they might depend on if you use a GPU or not. But that should be significantly faster, especially during the backward pass.


import torch
from torch.nn import functional as F

device = torch.device("cpu")

# Original code
def EdgeDetection(Input_Temp):
    K1 = torch.FloatTensor([[1,0,0],[0,-1,0],[0,0,0]]).to(device)
    K2 = torch.FloatTensor([[0,1,0],[-1,0,0],[0,0,0]]).to(device)

    # print "The input size is" , Input_Temp.size() ,  "\n" #The input size is for example (128, 3, 32, 32)

    D = 0 # D is responsible for number of Batches  or the 0st dimention
    for BatchNum in range(Input_Temp.size(0)):
        C = 0 # C is responsible for number of channels or the 1st dimention 
        for ChannelNum in range(Input_Temp.size(1)):

            TempImg = Input_Temp[BatchNum,ChannelNum,:,:].clone()
            TempImgNormalize = TempImg

            # Makeing the input to have size of d x c x m x m so it can be used in conv2d
            A1 = torch.unsqueeze(torch.unsqueeze(TempImgNormalize,0),0).to(device)
            Gx = F.conv2d(A1, torch.unsqueeze(torch.unsqueeze(K1,0),0),padding= K1.size(0)/2 )
            Gy = F.conv2d(A1, torch.unsqueeze(torch.unsqueeze(K2,0),0),padding= K2.size(0)/2 )
            GM = (Gx**2 + Gy**2)**0.5


            # concatinating the results, concatenating along the 1st dimention:
            if C == 0:
                APP = GM 
            else:
                APP = torch.cat((APP,GM),1)
            C += 1

        # concatinating the results, concatenating along the 0st dimention:
        if D == 0:
            Final = APP
            Final = Final.to(device)
        else:
            Final = torch.cat((Final,APP),0)
        D = D + 1

    return Final

# New one
# Pre-allocate these so that this is not done every time
# Either in the current file or in the __init__ or the nn.Module
K1 = torch.FloatTensor([[1,0,0],[0,-1,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k1 = 1
K2 = torch.FloatTensor([[0,1,0],[-1,0,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k2 = 1
def EdgeDetectionNew(Input_Temp):
    # Move the input to the right device (remove this line if this is not needed)
    Input_Temp = Input_Temp.to(device)
    # For the conv, make the channel look like batch
    # Compute the size for the conv
    inp_size = Input_Temp.size()
    conv_size = list(inp_size)
    conv_size[0] = inp_size[0] * inp_size[1]
    conv_size[1] = 1
    # View the input acordingly
    conv_input = Input_Temp.view(conv_size)
    Gx = F.conv2d(conv_input, K1, padding=pad_k1)
    Gy = F.conv2d(conv_input, K2, padding=pad_k2)
    GM = (Gx**2 + Gy**2)**0.5
    # View the output back to get the channel
    GM = GM.view(inp_size)

    return GM

inp = torch.rand(10, 3, 20, 20)
out = EdgeDetection(inp)
out2 = EdgeDetectionNew(inp)

print((out - out2).max())
1 Like

@albanD man, that was waaaay faster…
Thank you!

I have a follow up question.
If in a tensor T with size AxBxDxD i want to subtract the mean of each layer, do you know a fast way to do that too 8-?
Lets say for simplicity B is 1.

the classical way is to have two for loop for range(A) and range(B) and then subtract the mean from each layer, i was wondering if there is a fast way to do that too?
I cannot do T-torch.mean(T) because torch.mean(T) will give me the average of whole layers in T not each of them separately…

Edit: I figured it out
M = torch.mean(torch.mean(T,2),2)

a = (a.unsqueeze_(-1))
print a.size()
a = a.expand(A,D,D)
print a.size()
a.unsqueeze_(1)

@albanD I have a follow up question.
If i want to keep K1 as it is, but change K2 during training (with backpropagation), how should i do that?

K1 = torch.FloatTensor([[1,0,0],[0,-1,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k1 = 1
K2 = torch.FloatTensor([[0,1,0],[-1,0,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k2 = 1
def EdgeDetectionNew(Input_Temp):
   # Move the input to the right device (remove this line if this is not needed)
   Input_Temp = Input_Temp.to(device)
   # For the conv, make the channel look like batch
   # Compute the size for the conv
   inp_size = Input_Temp.size()
   conv_size = list(inp_size)
   conv_size[0] = inp_size[0] * inp_size[1]
   conv_size[1] = 1
   # View the input acordingly
   conv_input = Input_Temp.view(conv_size)
   Gx = F.conv2d(conv_input, K1, padding=pad_k1)
   Gy = F.conv2d(conv_input, K2, padding=pad_k2)
   GM = (Gx**2 + Gy**2)**0.5
   # View the output back to get the channel
   GM = GM.view(inp_size)

   return GM

Make K2 an nn.Parameter and store it in the nn.Module. Like it is done for all nn modules.
If you don’t use nn.Module, juste make sure that K2 requires gradients and that it is passed to your optimizer.