I'm doing the online Computer Vision course by UMich and am new to PyTorch. One of the assignment questions is on batch matrix multiplication, where we have to find the batch matrix product with and without the bmm function. Here is the code.
def batched_matrix_multiply(x, y, use_loop=True):
"""
Perform batched matrix multiplication between the tensor x of shape (B, N, M)
and the tensor y of shape (B, M, P).
If use_loop=True, then you should use an explicit loop over the batch
dimension B. If loop=False, then you should instead compute the batched
matrix multiply without an explicit loop using a single PyTorch operator.
Inputs:
- x: Tensor of shape (B, N, M)
- y: Tensor of shape (B, M, P)
- use_loop: Whether to use an explicit Python loop.
Hint: torch.stack, bmm
Returns:
- z: Tensor of shape (B, N, P) where z[i] of shape (N, P) is the result of
matrix multiplication between x[i] of shape (N, M) and y[i] of shape
(M, P). It should have the same dtype as x.
"""
z = None
#############################################################################
# TODO: Implement this function #
#############################################################################
# Replace "pass" statement with your code
z = torch.zeros(x.shape[0], x.shape[1], y.shape[2])
if use_loop == True:
for i in range(x.shape[0]):
z[i] = torch.mm(x[i], y[i])
else:
z = torch.bmm(x,y)
#############################################################################
# END OF YOUR CODE #
#############################################################################
return z
I've managed to do it without bmm, but without using the torch.stack hint. I initialized a zeroes matrix 'z' with the dimensions of the output matrix and performed normal matrix multiplication for each batch using the for loop.
I'd like to know what the more efficient answer using torch.stack is.
