Есть ли способ векторизации, применяющий функцию к каждому элементу тензора? - PullRequest
1 голос
/ 18 октября 2019

Я хочу реализовать правило обновления в процедуре оптимизации. Я хотел векторизовать его, используя функции pytorch, чтобы улучшить сложность времени. Возможно, запускается на графическом процессоре.

Правило обновления похоже на

u[i][j] = multiplier*NORM(Q.column(i)-Q.column(j))

Я пытался написать простой вложенный цикл для его выполнения, но он кажется очень медленным.

def u_vector(Q,parameters):
    u=torch.zeros(Q.shape[0],Q.shape[0],requires_grad=False)
    for i in range(Q.shape[0]):
        for j in range(i+1,Q.shape[1]):
            u[i][j]=0.5*(torch.norm(Q[:,i]-Q[:,j]))*parameters['gamma']
            u[j][i]=u[i][j]
    return u

Может кто-нибудь помочь, пожалуйста?

1 Ответ

0 голосов
/ 18 октября 2019

Вот одна из возможностей использования индексации для вычисления всех различий и норм одновременно. Вам нужно проверить, работает ли оно быстрее на вашем оборудовании.

def u_vector(Q,parameters):
    cols = Q.shape[1]
    # equivalent to generating a meshgrid
    c1 = [i for i in range(cols) for _ in range(cols)]
    c2 = [i for _ in range(cols) for i in range(cols)]
    return 0.5 * torch.norm(Q[:,c1] - Q[:,c2], dim=0).reshape(cols, cols) * parameters['gamma']

Отказ от ответственности . Довольно странно, что вы используете Q.shape[0] для определения формы u и диапазона индексов, поскольку вся ваша индексация используется для доступа к столбцам Q. Приведенный выше код соответствует вашей функции только в том случае, если Q является квадратом (также ваш произойдет сбой из-за неверного индекса, если Q имеет больше столбцов, чем строк). Если вы хотите что-то такое же, что и ваш, то сработает следующее. Единственная разница между этой и предыдущей версией состоит в том, что она дополняется нулями, когда Q.shape[0] > Q.shape[1].

def u_vector(Q,parameters):
    rows, cols = Q.shape[0], Q.shape[1]
    c1 = [i for i in range(cols) for _ in range(cols)]
    c2 = [i for _ in range(cols) for i in range(cols)]
    u = 0.5 * torch.norm(Q[:,c1] - Q[:,c2], dim=0).reshape(cols, cols) * parameters['gamma']
    return torch.nn.functional.pad(u, (0, rows-cols, 0, rows-cols), 'constant', 0)

Также обратите внимание, что она будет работать на графическом процессоре, если Q на графическом процессоре.

...