Modules chunk

[wanghongsheng01]

[Flowingsun007]

这下面的空行都删了吧

[wanghongsheng01]

嗯嗯，好的

[Flowingsun007]

这个操作torch也是没有反向吗

[doombeaker]

有反向的，hongsheng 还是加入反向的测试吧（因为是用 slice 拼的，所以不用自己实现反向，但是反向的测试还是要有的）

[wanghongsheng01]

嗯嗯，有反向的，马上添加

[Flowingsun007]

test case多加两个吧，如：4-d input；dim=2/3；不同的chunks ...

[wanghongsheng01]

好的，马上添加

[Flowingsun007]

需要导出下Tensor.chunk

[wanghongsheng01]

嗯嗯，好的

[Flowingsun007]

可以用functional api改一下（flow.F.slice(xxx)）

[wanghongsheng01]

好的，这个有没有其他已 merge 的 module，可供参考呢？

[wanghongsheng01]

已改好👌

[wanghongsheng01]

flow.chunk 思路

1. 求出每块的大小，除不尽时，求出最后一块的大小

channel = input.dim()
dim_size = input.shape[dim]
chunk_size = dim_size / chunks if dim_size % chunks == 0 else (int)(dim_size / chunks)
last_chunk_size = dim_size / chunks if dim_size % chunks == 0 else dim_size - (chunk_size * (chunks - 1))

2. 因为利用 slice 拼，所以对于每块 chunk，需要求用 slice 提取该 chunk 的 tensor 所需的整张图的 tuple_list
3. 如何求每块 chunk对应的整张图的 tuple_list？
   对于切的维度，求该维度对应的 tuple_list = [start, stop, step]; 对于不切的维度，对应 tuple_list = [None, None, None]

• 保存切的维度的 tuple_list：一键多值 dict = {
  '切的维度' : [ chunk1 对应的 [start, stop, step]，chunk2 对应的 [start, stop, step]，... ]
  }

chunk_dim_dict = {}
tup_ndim = []
splits = []

for chunk in range(0, chunks): 
    if dim_size % chunks == 0:
        start = chunk * chunk_size
        stop = (chunk + 1) * chunk_size
    else:
        start = chunk * chunk_size if chunk < chunks - 1 else chunk_size * (chunks - 1)
        stop = (chunk + 1) * chunk_size if chunk < chunks - 1 else dim_size
    step = 1
    chunk_dim_dict.setdefault(dim, []).append([int(start), int(stop), step])

• 整张图的 tuple_list
  遍历切的维度的 value，对于每个 chunk，为该维度的 tuple_list 添加其他不切维度的 tuple_list = [None, None, None]
• 对于每个 chunk，利用 slice 和该 chunk 对应的整张图的 tuple_list，提取该 chunk 的 tensor 分量

for k, v in chunk_dim_dict.items():
      for v_chunk in v:
          tup_list = []
          for i in range(0, channel):
              if i != dim:
                  tup_list.append([None, None, None])
              else:
                  tup_list.append(v_chunk)
          splits.append(
              flow.experimental.slice(input, slice_tup_list=tup_list)
          )