how to add list of arrays (tensors)

user_3899975 0
Viewed 4959

I am defining a simple conv2d function to calculate the cross-correlation between input and kernel (both 2D tensor) as below:

import torch 

def conv2D(X, K):
    h = K.shape[0]
    w = K.shape[1]
    ĥ = X.shape[0] - h + 1
    ŵ = X.shape[1] - w + 1
    Y = torch.zeros((ĥ, ŵ)) 
    for i in range (ĥ):
        for j in range (ŵ):
            Y[i, j] = (X[i: i+h, j: j+w]*K).sum()

    return Y

When X and K are of rank-3 tensor, I calculate the conv2d for each channel and then add them together as below:

def conv2D_multiple(X, K):
    cross = []
    result = 0
    for x, k in zip(X, K):
        cross.append(conv2D(x,k))

    for t in cross:
        result += t

    return result

To test my function:

X_2 = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]], 
                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]]], dtype=torch.float32)
K_2 = torch.tensor([[[0, 1], [2, 3]], [[1, 2], [3, 4]]], dtype=torch.float32)

conv2D_multiple(X_2, K_2)

The results is:

tensor([[ 56.,  72.],
        [104., 120.]])

The result is as expected, however, I believe my second for loop inside conv2D_multiple(X, K) function is redundant. My question is how to sum (element wise) tensors (arrays) in the list so I omit the second for loop.

pytorch tensorflow numpy python
1 Answers

Since your conv2D operates on a per slice behaviour, what you can do is allocate a 3D tensor so that when you use the first for loop, you store the results by taking each result and populating each slice. You can then sum along the dimension of the slices using PyTorch's built-in torch.sum operator on the tensor to get the same result. To make it palatable, I'll make the slice dimension dim=0. Therefore, replace cross from being an initial empty list to a Torch tensor that is 3D to allow you to store the intermediate results, then compress along the slice dimension by summing. We can get away with doing this as your initial implementation stored the intermediate results as a list of 2D tensors. To make it easier, go to 3D and allow PyTorch to sum along the slice axis.

This will require that you define the correct dimensions for this 3D tensor first prior to looping:

def conv2D_multiple(X, K):
    h = K.shape[1]
    w = K.shape[2]
    ĥ = X.shape[1] - h + 1
    ŵ = X.shape[2] - w + 1
    c = X.shape[0]
    cross = torch.zeros((c, ĥ, ŵ), dtype=torch.float32)
    for i, (x, k) in enumerate(zip(X, K)):
        cross[i] = conv2D(x,k)

    result = cross.sum(dim=0)
    return result

Notice that for each slice you're iterating over between the input and kernel, instead of appending to a new list we directly place this into a slice in the intermediate tensor. Once you store these results, sum along the slice axis to finally compress it into what you expect. Running the new function above with your example inputs generates the same result.

If this isn't a desired result for you, another way is to simply take the list of tensors you created, build the intermediate tensor out of that by stacking them all together using torch.stack and sum. By default it stacks along the first axis (dim=0):

def conv2D_multiple(X, K):
    cross = []
    result = 0
    for x, k in zip(X, K):
        cross.append(conv2D(x,k))

    cross = torch.stack(cross)
    result = cross.sum(dim=0)
    return result
user_32508290
Related