Apply a multi-head self-attention block across the pixels of a feature map and add its output to the input.
ZOj=tokens(GN(x)),Q=ZWq⊤+bq,K=ZWk⊤+bk,V=ZWv⊤+bv=softmax(dQjKj⊤)Vj,out=x+image([O1,…,On]Wo⊤+bo)Here x has shape (B, C, H, W) and GN is a group norm with num_groups groups, eps 1e-5, weight norm_weight and bias norm_bias. The tokens step reshapes each sample to (HW, C) with pixel (i, j) at row iW + j. The weights W_q, W_k, W_v and W_o are q_weight, k_weight, v_weight and out_weight, used as in a linear layer, and the biases b_q, b_k, b_v and b_o are q_bias, k_bias, v_bias and out_bias. With n equal to num_heads and d = C / n, head j uses channels j*d to (j+1)*d minus 1 of Q, K and V. The softmax runs over the key axis, the head outputs are concatenated back along channels in head order, and the image step undoes the tokens step.
| key | shape |
|---|---|
| norm_weight, norm_bias | (C,) |
| q_weight, k_weight, v_weight, out_weight | (C, C) |
| q_bias, k_bias, v_bias, out_bias | (C,) |
Return a PyTorch tensor of shape (B, C, H, W) with the dtype of x.
Input: x = [[[[0.47,1.061],[-0.77,-0.715]],[[0.833,1.842],[2.28,-2.226]],[[0.193,0.666],[0.586,-0.091]],[[-1.114,-1.565],[-0.727,2.123]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (4,); q_weight, k_weight, v_weight, out_weight (4, 4); values in the test case}, num_groups = 2, num_heads = 2
Output: [[[[0.637623,1.464757],[-0.454166,-1.642022]],[[0.763857,1.862223],[2.309599,-2.682708]],[[-0.078845,0.325763],[0.296335,-0.116245]],[[-0.969084,-1.4404],[-0.620036,2.427084]]]]
Explanation: The 4 channels form 2 heads of 2 channels, so the scores are divided by sqrt(2). In head 0 the first pixel spreads its attention as about 0.240, 0.257, 0.305 and 0.197 over the four pixels, and the projected mix is added back to x.
Input: x = [[[[1.212,2.034,-1.923],[1.651,-0.567,0.778]],[[-1.268,-2.842,0.056],[0.71,-0.708,-0.232]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (2,); q_weight, k_weight, v_weight, out_weight (2, 2); values in the test case}, num_groups = 1, num_heads = 1
Output: [[[[1.069589,1.877189,-2.020708],[1.498128,-0.67997,0.642647]],[[-1.2719,-2.888039,0.151091],[0.683075,-0.641081,-0.214674]]]]
Input: x = [[[[-0.748]],[[-0.549]],[[-0.535]],[[1.884]]],[[[0.514]],[[-0.94]],[[-0.825]],[[0.399]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (4,); q_weight, k_weight, v_weight, out_weight (4, 4); values in the test case}, num_groups = 4, num_heads = 2
Output: [[[[-0.866129]],[[-0.714209]],[[-0.67807]],[[2.055774]]],[[[0.395871]],[[-1.105209]],[[-0.96807]],[[0.570774]]]]
x.flatten(2).transpose(1, 2) turns (B, C, H, W) into (B, H*W, C) tokens.
q.view(B, N, num_heads, d).transpose(1, 2) splits heads.
torch.softmax(scores, dim=-1) normalizes over the keys.
Accepts: array
Accepts: object
Accepts: integer
Accepts: integer
Apply a multi-head self-attention block across the pixels of a feature map and add its output to the input.
ZOj=tokens(GN(x)),Q=ZWq⊤+bq,K=ZWk⊤+bk,V=ZWv⊤+bv=softmax(dQjKj⊤)Vj,out=x+image([O1,…,On]Wo⊤+bo)Here x has shape (B, C, H, W) and GN is a group norm with num_groups groups, eps 1e-5, weight norm_weight and bias norm_bias. The tokens step reshapes each sample to (HW, C) with pixel (i, j) at row iW + j. The weights W_q, W_k, W_v and W_o are q_weight, k_weight, v_weight and out_weight, used as in a linear layer, and the biases b_q, b_k, b_v and b_o are q_bias, k_bias, v_bias and out_bias. With n equal to num_heads and d = C / n, head j uses channels j*d to (j+1)*d minus 1 of Q, K and V. The softmax runs over the key axis, the head outputs are concatenated back along channels in head order, and the image step undoes the tokens step.
| key | shape |
|---|---|
| norm_weight, norm_bias | (C,) |
| q_weight, k_weight, v_weight, out_weight | (C, C) |
| q_bias, k_bias, v_bias, out_bias | (C,) |
Return a PyTorch tensor of shape (B, C, H, W) with the dtype of x.
Input: x = [[[[0.47,1.061],[-0.77,-0.715]],[[0.833,1.842],[2.28,-2.226]],[[0.193,0.666],[0.586,-0.091]],[[-1.114,-1.565],[-0.727,2.123]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (4,); q_weight, k_weight, v_weight, out_weight (4, 4); values in the test case}, num_groups = 2, num_heads = 2
Output: [[[[0.637623,1.464757],[-0.454166,-1.642022]],[[0.763857,1.862223],[2.309599,-2.682708]],[[-0.078845,0.325763],[0.296335,-0.116245]],[[-0.969084,-1.4404],[-0.620036,2.427084]]]]
Explanation: The 4 channels form 2 heads of 2 channels, so the scores are divided by sqrt(2). In head 0 the first pixel spreads its attention as about 0.240, 0.257, 0.305 and 0.197 over the four pixels, and the projected mix is added back to x.
Input: x = [[[[1.212,2.034,-1.923],[1.651,-0.567,0.778]],[[-1.268,-2.842,0.056],[0.71,-0.708,-0.232]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (2,); q_weight, k_weight, v_weight, out_weight (2, 2); values in the test case}, num_groups = 1, num_heads = 1
Output: [[[[1.069589,1.877189,-2.020708],[1.498128,-0.67997,0.642647]],[[-1.2719,-2.888039,0.151091],[0.683075,-0.641081,-0.214674]]]]
Input: x = [[[[-0.748]],[[-0.549]],[[-0.535]],[[1.884]]],[[[0.514]],[[-0.94]],[[-0.825]],[[0.399]]]], params = {10 tensors: norm_weight, norm_bias and the four biases (4,); q_weight, k_weight, v_weight, out_weight (4, 4); values in the test case}, num_groups = 4, num_heads = 2
Output: [[[[-0.866129]],[[-0.714209]],[[-0.67807]],[[2.055774]]],[[[0.395871]],[[-1.105209]],[[-0.96807]],[[0.570774]]]]
x.flatten(2).transpose(1, 2) turns (B, C, H, W) into (B, H*W, C) tokens.
q.view(B, N, num_heads, d).transpose(1, 2) splits heads.
torch.softmax(scores, dim=-1) normalizes over the keys.
Accepts: array
Accepts: object
Accepts: integer
Accepts: integer
