Вот одна из возможностей использования индексации для вычисления всех различий и норм одновременно. Вам нужно проверить, работает ли оно быстрее на вашем оборудовании.
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
на графическом процессоре.