У меня есть два тензора формы a(16,8,8,64) и b(64,64). Предположим, я извлекаю последнее измерение a в другой вектор-столбец c, я хочу вычислить matmul(matmul(c.T, b), c). Я хочу, чтобы это было сделано в каждом из первых трех измерений a. То есть конечный продукт должен иметь форму (16,8,8,1). Как я могу добиться этого в pytorch?

2
Nagabhushan S N 20 Дек 2019 в 17:56

1 ответ

Лучший ответ

Можно сделать следующим образом:

row_vec = a[:, :, :, None, :].float()
col_vec = a[:, :, :, :, None].float()
b = (b[None, None, None, :, :]).float()
prod = torch.matmul(torch.matmul(row_vec, b), col_vec)
0
Nagabhushan S N 21 Дек 2019 в 06:55
59427249