CODEWITHSUNDEEP
pytorchrepeat
River
River

Reputation: 682

How to repeat tensor in a specific new dimension in PyTorch

If I have a tensor A which has shape [M, N], I want to repeat the tensor K times so that the result B has shape [M, K, N] and each slice B[:, k, :] should has the same data as A. Which is the best practice without a for loop. K might be in other dimension.

torch.repeat_interleave() and tensor.repeat() does not seem to work. Or I am using it in a wrong way.

Upvotes: 50

Views: 75534

Answers (5)

0-_-0
0-_-0

Reputation: 1473

Explicitly repeating values can quickly create huge memory cost. In most cases, you can keep the values implicit by utilizing broadcasting instead. So can use A[:, None, :] and get as new shape A.shape==(M, 1, N).

One case where I would agree to repeating the values, is that of in-place operations in the following steps. As numpy and torch differ in their implementations I like the agnostic (A * torch.ones(K, 1, 1))) followed by a transpose.

Upvotes: 1

fatpanda2049
fatpanda2049

Reputation: 573

tensor.expand might be a better choice than tensor.repeat because according to this: "Expanding a tensor does not allocate new memory, but only creates a new view on the existing tensor where a dimension of size one is expanded to a larger size by setting the stride to 0."

However, be aware that: "More than one element of an expanded tensor may refer to a single memory location. As a result, in-place operations (especially ones that are vectorized) may result in incorrect behavior. If you need to write to the tensors, please clone them first."

M = N = K = 3
A = torch.arange(0, M * N).reshape((M, N))
B = A.unsqueeze(1).expand(M, K, N)
B

'''
tensor([[[0, 1, 2],
         [0, 1, 2],
         [0, 1, 2]],

        [[3, 4, 5],
         [3, 4, 5],
         [3, 4, 5]],

        [[6, 7, 8],
         [6, 7, 8],
         [6, 7, 8]]])
'''

Upvotes: 4

jodag
jodag

Reputation: 22224

tensor.repeat should suit your needs but you need to insert a unitary dimension first. For this we could use either tensor.unsqueeze or tensor.reshape. Since unsqueeze is specifically defined to insert a unitary dimension we will use that.

B = A.unsqueeze(1).repeat(1, K, 1)

Code Description A.unsqueeze(1) turns A from an [M, N] to [M, 1, N] and .repeat(1, K, 1) repeats the tensor K times along the second dimension.

Upvotes: 89

Ashish Alex
Ashish Alex

Reputation: 11

Adding to the answer provided by @Alleo. You can use following Einops function.

einops.repeat(example_tensor, 'b h w -> (repeat b) h w', repeat=b)

Where b is the number of times you want your tensor to be repeated and h, w the additional dimensions to the tensor.

Example -

example_tensor.shape -> torch.Size([1, 40, 50]) 
repeated_tensor = einops.repeat(example_tensor, 'b h w -> (repeat b) h w', repeat=8)
repeated_tensor.shape -> torch.Size([8, 40, 50])

More examples here - https://einops.rocks/api/repeat/

Upvotes: 1

Alleo
Alleo

Reputation: 8528

Einops provides repeat function

import einops
einops.repeat(x, 'm n -> m k n', k=K)

repeat can add arbitrary number of axes in any order and reshuffle existing axes at the same time.

Upvotes: 12

Related Questions