散布してnumpyで操作を収集する方法は？

Question

NumpyでTensorflowまたはPyTorchの分散および収集操作を実装したいと考えています。私はしばらく頭を掻いてきました。すべてのポインタは大いに感謝しています！

• PyTorch Scatter

• PyTorch Gather

Answer 1

の方法は、私が予想していた以上に多くの作業であることが判明しました。私はそれのためにNumPyの準備が整った機能を見つけませんでした。私はNumPyで実装する必要があるかもしれない人のためにここで共有しています。 （PS selfは、法の宛先または出力されます。）

def scatter_numpy(self, dim, index, src): 
    """ 
    Writes all values from the Tensor src into self at the indices specified in the index Tensor. 

    :param dim: The axis along which to index 
    :param index: The indices of elements to scatter 
    :param src: The source element(s) to scatter 
    :return: self 
    """ 
    if index.dtype != np.dtype('int_'): 
     raise TypeError("The values of index must be integers") 
    if self.ndim != index.ndim: 
     raise ValueError("Index should have the same number of dimensions as output") 
    if dim >= self.ndim or dim < -self.ndim: 
     raise IndexError("dim is out of range") 
    if dim < 0: 
     # Not sure why scatter should accept dim < 0, but that is the behavior in PyTorch's scatter 
     dim = self.ndim + dim 
    idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1:] 
    self_xsection_shape = self.shape[:dim] + self.shape[dim + 1:] 
    if idx_xsection_shape != self_xsection_shape: 
     raise ValueError("Except for dimension " + str(dim) + 
         ", all dimensions of index and output should be the same size") 
    if (index >= self.shape[dim]).any() or (index < 0).any(): 
     raise IndexError("The values of index must be between 0 and (self.shape[dim] -1)") 

    def make_slice(arr, dim, i): 
     slc = [slice(None)] * arr.ndim 
     slc[dim] = i 
     return slc 

    # We use index and dim parameters to create idx 
    # idx is in a form that can be used as a NumPy advanced index for scattering of src param. in self 
    idx = [[*np.indices(idx_xsection_shape).reshape(index.ndim - 1, -1), 
      index[make_slice(index, dim, i)].reshape(1, -1)[0]] for i in range(index.shape[dim])] 
    idx = list(np.concatenate(idx, axis=1)) 
    idx.insert(dim, idx.pop()) 

    if not np.isscalar(src): 
     if index.shape[dim] > src.shape[dim]: 
      raise IndexError("Dimension " + str(dim) + "of index can not be bigger than that of src ") 
     src_xsection_shape = src.shape[:dim] + src.shape[dim + 1:] 
     if idx_xsection_shape != src_xsection_shape: 
      raise ValueError("Except for dimension " + 
          str(dim) + ", all dimensions of index and src should be the same size") 
     # src_idx is a NumPy advanced index for indexing of elements in the src 
     src_idx = list(idx) 
     src_idx.pop(dim) 
     src_idx.insert(dim, np.repeat(np.arange(index.shape[dim]), np.prod(idx_xsection_shape))) 
     self[idx] = src[src_idx] 

    else: 
     self[idx] = src 

    return self 

ありgatherのためのシンプルな解決策になるが、これは私が上で定住するものである可能性があり：
（ここself値が収集されているということですndarray ）

def gather_numpy(self, dim, index): 
    """ 
    Gathers values along an axis specified by dim. 
    For a 3-D tensor the output is specified by: 
     out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 
     out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 
     out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 

    :param dim: The axis along which to index 
    :param index: A tensor of indices of elements to gather 
    :return: tensor of gathered values 
    """ 
    idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1:] 
    self_xsection_shape = self.shape[:dim] + self.shape[dim + 1:] 
    if idx_xsection_shape != self_xsection_shape: 
     raise ValueError("Except for dimension " + str(dim) + 
         ", all dimensions of index and self should be the same size") 
    if index.dtype != np.dtype('int_'): 
     raise TypeError("The values of index must be integers") 
    data_swaped = np.swapaxes(self, 0, dim) 
    index_swaped = np.swapaxes(index, 0, dim) 
    gathered = np.choose(index_swaped, data_swaped) 
    return np.swapaxes(gathered, 0, dim) 

Answer 2

フォアrefとindices numpyの配列である：

散布更新：

は収集：

ref[indices]   # tf.gather(ref, indices) 
ref[:, indices]  # tf.gather(ref, indices, axis=1) 
ref[..., indices, :] # tf.gather(ref, indices, axis=-2) 
ref[..., indices]  # tf.gather(ref, indices, axis=-1) 

は、より多くのためにnumpy docs on indexingを参照してください。

Answer 3

@DomJackで提案されているように、スライスの割り当てを使用するのではなく、散乱の場合は、np.add.at;これはスライスの割り当てとは異なり、重複したインデックスの存在下では明確な動作をしているためです。