Understanding an Einsum usage for graph convolution

user_6159217 0
Viewed 95

I am reading the code for the spatial-temporal graph convolution operation here: https://github.com/yysijie/st-gcn/blob/master/net/utils/tgcn.py and I'm having some difficulty understanding what is happening with an einsum operation. In particular

For x a tensor of shape (N, kernel_size, kc//kernel_size, t, v), where kernel_size is typically 3, and lets say kc=64*kernel_size, t is the number of frames, say 64, and v the number of vertices, say 25. N is the batch size.

Now for a tensor A of shape (3, 25, 25) where each dimension is a filtering op on the graph vertices, an einsum is computed as:

x = torch.einsum('nkctv,kvw->nctw', (x, A))

I'm not sure how to interpret this expression. What I think it's saying is that for each batch element, for each channel c_i out of 64, sum each of the three matrices obtained by matrix multiplication of the (64, 25) feature map at that channel with the value of A[i]. Do I have this correct? The expression is a bit of a mouthful, and notation wise there seems to be a bit of a strange usage of kc as one variable name, but then decomposition of k as kernel size and c as the number of channels (192//3 = 64) in the expression for the einsum. Any insights appreciated.

pytorch
2 Answers
Y = torch.einsum('nkctv,kvw->nctw', (x, A)) means:

einsum interpretation on graph

For better understanding, I have replaced the x in left hand side with Y

user_50798800

Helps when you look closely at the notation:

  • nkctv for left side
  • kvw on the right side
  • nctw being the result

Missing from the result are:

  • k
  • v

These elements are summed together into a single value and squeezed, leaving us the resulting shape.

Something along the lines of (expanded shapes (added 1s) are broadcasted and sum per element):

  • left: (n, k, c, t, v, 1)
  • right: (1, k, 1, 1, v, w)

Now it goes (l, r for left and right):

  • torch.mul(l, r)
  • torch.sum(l, r, dim=(1, 4))
  • squeeze any singular dimensions

It is pretty hard to get, hence Einstein's summation helps in terms of thinking about resulting shapes “mixed” with each other, at least for me.

user_108864200
Related