代码片段记录11

需求:对于两个向量$a$、$b$,$a,b \in R^d$,定义一种减法,有:

其中$M \in R^{d\times d}$,$M_{ij}=a_i-b_j$

在代码中实际的维度:

1
2
a=torch.rand(batch_size,sequence_len,dim)
b=torch.rand(batch_size,sequence_len,dim)

方法①:for循环

1
2
3
4
5
6
M=torch.zeros(bz,seq_len,seq_len)
for b_i in range(bz):
for i in range(seq_len):
for j in range(seq_len):
M_ij=torch.norm(a[b_i][i]-b[b_i][j])
M[b][i][j]=M_ij

方法②:矩阵运算

1
2
3
a=a.unsqueeze(2)  # bz,seq_len,1,dim
b=b.unsqueeze(1) # bz,1,seq_lens,dim
M=torch.norm(a-b,dim=-1) # will broadcast

需求,生成一个mask矩阵,每一行有一段连续的位置填充1,其中每一行填充1的开始位置和结束位置都不同。具体来说,先生成一个中心位置center,则开始位置为center-window;结束位置为center+window。其中开始位置和结束位置不能越界,也即不小于0和大于行的总长度。
如:

思路:
①先生成n行每行对应的随机中心位置,然后再获得左和右边界

1
2
3
4
5
6
7
8

centers=torch.randint(low=0,high=query_len,size=(query_len,),dtype=torch.long)

left=centers-self.window
left=torch.max(left,torch.LongTensor([0])).unsqueeze(1) # query_len,1

right=centers+self.window
right=torch.min(right,torch.LongTensor([query_len-1])).unsqueeze(1) # query_len,1

②生成一个每行都用[0,n-1]填充的矩阵,[0,n-1]表示的是该元素的index,亦即:

1
range_matrix=torch.range(0,query_len-1,dtype=torch.long).unsqueeze(0).expand(query_len,-1)  # query_len,query_len

③利用<=和>=获得一个左边界和右边界矩阵,左边界矩阵表示在该左边界的左边都是填充的1;右边界矩阵表示在该右边界右边都是填充的1。再进行异或操作。

1
2
3
4
range_matrix=torch.range(0,query_len-1,dtype=torch.long).unsqueeze(0).expand(query_len,-1)  # query_len,query_len
left_matrix=range_matrix<=left
right_matrix=range_matrix<=right
final_matrix=left_matrix^right_matrix