0
$\begingroup$

I would like to produce the following tensor of size (N*N) where the ones (D) appear as follows:

create_mask(N=10, D=5)

tensor([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]])

or

create_mask(N=5, D=3)

tensor([[1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0],
        [0, 1, 1, 1, 0],
        [0, 0, 1, 1, 1],
        [0, 0, 1, 1, 1]])

I need this function to be very efficient/fast because I will be generating large tensors (usually N~[4000,10000]).

Here is my current approach (too slow) :

import torch
import numpy as np


def create_mask(num_words, num_around):
    
    assert num_around & 1
    num_neighbor = num_around//2

    output = np.array([np.array(range(i-num_neighbor, i+num_neighbor+1)) for i in range(num_words)])
    output[:num_neighbor, :] = output[num_neighbor]
    output[-num_neighbor:, :] = output[-num_neighbor-1]
    output = torch.tensor(output)

    mask = torch.nn.functional.one_hot(output, num_words).sum(1)
    
    return mask

Does anyone have a better approach?

$\endgroup$
1
  • $\begingroup$ This question probably belongs on Stackoverflow. $\endgroup$ Apr 14 at 13:38

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.