Заполните тензор другим тензором, где маска истинна

#python #pytorch #tensor

Вопрос:

Мне нужно вставить элементы нового тензора в старый тензор с определенной вероятностью, скажем, для простоты 0,8. По существу, это то, что сделал бы masked_fill, но он работает только с одномерным тензором. На самом деле я делаю

     prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    dim1, dim2, dim3, dim4 = new.shape
    for a in range(dim1):
        for b in range(dim2):
            for c in range(dim3):
                for d in range(dim4):
                    old[a][b][c][d] = old[a][b][c][d] if mask[a][b][c][d] else new[a][b][c][d]
 

и это ужасно. Я бы хотел что-то вроде

     prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    old = trgs.multidimensional_masked_fill(mask, new)
 

Ответ №1:

Я не уверен, каковы некоторые из ваших объектов, но это должно заставить вас сделать то, что вам нужно, в кратчайшие сроки:

old это ваши существующие данные.

mask является ли маска, которую вы сгенерировали с вероятностью p

new это новый тензор, содержащий элементы, которые вы хотите вставить.

 # torch.where
result = old.where(mask, new)