CODEWITHSUNDEEP
pythonpytorchtensortorchtext
chico0913
chico0913

Reputation: 647

How to get rid of every column that are filled with zero from a Pytorch tensor?

I have a pytorch tensor A like below:

A = 
tensor([[  4,   3,   3,  ...,   0,   0,   0],
        [ 13,   4,  13,  ...,   0,   0,   0],
        [707, 707,   4,  ...,   0,   0,   0],
        ...,
        [  7,   7,   7,  ...,   0,   0,   0],
        [  0,   0,   0,  ...,   0,   0,   0],
        [195, 195, 195,  ...,   0,   0,   0]], dtype=torch.int32)

I would like to:

I can imagine doing:

zero_list = []
for j in range(A.size()[1]):
    if torch.sum(A[:,j]) == 0:
         zero_list = zero_list.append(j)

to identify the columns that only has 0 for its elements but I am not sure how to delete such columns filled with 0 from the original tensor.

How can I delete the columns with zero from a pytorch tensor based on the index number?

Thank you,

Upvotes: 4

Views: 6041

Answers (2)

Aydo
Aydo

Reputation: 416

Identify all the columns whose all of its entries are equal to 0

non_empty_mask = A.abs().sum(dim=0).bool()

This sums up over the absolute values of each column and then converts the result to a boolean, i.e. False if the sum is zero and True otherwise.

Delete only those columns that has all of their entries equal to 0

A[:,non_empty_mask]

This simply applies the mask to the original tensor, i.e. it keeps the rows where non_empty_mask is True.

Upvotes: 7

jodag
jodag

Reputation: 22274

It makes more sense to index the columns you want to keep instead of what you want to delete.

valid_cols = []
for col_idx in range(A.size(1)):
    if not torch.all(A[:, col_idx] == 0):
        valid_cols.append(col_idx)
A = A[:, valid_cols]

Or a little more cryptically

valid_cols = [col_idx for col_idx, col in enumerate(torch.split(A, 1, dim=1)) if not torch.all(col == 0)]
A = A[:, valid_cols]

Upvotes: 2

Related Questions