Python: Vectorisation de Calcul mis en œuvre à l'aide de Démarche Itérative

Car vous avez affaire à trois dimensions n, met c (en minuscules pour éviter toute ambiguïté), il peut être utile de changer la forme de tous vos tenseurs de (n, m, c), en reproduisant leurs valeurs sur la dimension manquante (par ex. M(m, c) devient M(n, m, c)).

Cependant, vous pouvez ignorer l'explicite la réplication et l'utilisation de la radiodiffusion, de sorte qu'il est suffisant pour unsqueeze la dimension manquante (par ex. M(m, c) devient M(1, m, c).

Compte tenu de ces considérations, la vectorisation de votre code devient comme suit

cond = y.unsqueeze(2) == torch.arange(M.size(dim=1)).unsqueeze(0)
pos = ds.unsqueeze(2) * M.unsqueeze(1) * cond
neg = ds.unsqueeze(2) * M.unsqueeze(1).sum(dim=0, keepdim=True) * ~cond
grad += (pos - neg).sum(dim=0)

Voici un petit test pour vérifier la validité de la solution

import torch

n, m, c = 11, 5, 7

y = torch.randint(c, size=(n, 1))
ds = torch.rand(n, m)
M = torch.randint(2, size=(n, c))
grad = torch.rand(m, c)


def slow_grad(y, ds, M, grad):
    for i in range(M.size(dim=1)):
        for j in range(ds.size(dim=0)):
            if y[j] == i:
                grad[:,i] = grad[:,i] - (ds[j,:].T * sum(M[j,:]))
            else:
                if M[j,i] > 0:
                    grad[:,i] = grad[:,i] + ds[j,:].T
    return grad


def fast_grad(y, ds, M, grad):
    cond = y.unsqueeze(2) == torch.arange(M.size(dim=1)).unsqueeze(0)
    pos = ds.unsqueeze(2) * M.unsqueeze(1) * cond
    neg = ds.unsqueeze(2) * M.unsqueeze(1).sum(dim=0, keepdim=True) * ~cond
    grad += (pos - neg).sum(dim=0)
    return grad
  
# Assert equality of all elements function outputs, throws an exception if false
assert torch.all(slow_grad(y, ds, M, grad) == fast_grad(y, ds, M, grad))

N'hésitez pas à tester sur d'autres affaires aussi bien!