Pytoch implements tf Functions of the gather() function

Purpose ：

Large size medical image [3, 897, 1196] In a batch in

Divide by loop [12, 3, 299, 299], After superposition and segmentation, it is [24, 3, 299, 299]

 

There is a problem with the order on the left side of the above figure , Now we need some operations to convert it to the above figure , Write code .

        for i, (inputs, labels,img_path) in enumerate(train_loader):
            D_L = len(inputs)
            # labels = labels.to(device)
            img_tensor = Variable(inputs.to(device), volatile=True)
            crop_img = torch.zeros((4, 3))
            print(" The first ",i," individual ")

            # Big picture [3, 897, 1196] Divided into [12, 3, 299, 299]
            for j in range(3):
                for ii in range(4):
                    if ii == 0 and j == 0:
                        crop_img = img_tensor[:, :, 299 * j:299 * (j + 1), 299 * ii:299 * (ii + 1)]
                    else:
                        crop_img = torch.cat(
                            (crop_img, img_tensor[:, :, 299 * j:299 * (j + 1), 299 * ii:299 * (ii + 1)]), 0)
            print("crop_img:", crop_img.shape)

            #  Rearrange by index , take [24*batch_size,3,299,299], Arrange the pictures in order of large size 
            trans_image = torch.zeros([12, 3, 299, 299]).to(device)
            for ind in range(D_L):
                index = np.arange(ind, 12 * D_L, D_L)
                index = torch.from_numpy(index).long().to(device)
                index = index.unsqueeze(1).unsqueeze(1).unsqueeze(1)
                index = index.expand(12, 3, 299, 299)
                t = torch.gather(crop_img, 0, index)
                if ind == 0:
                    trans_image = t
                else:
                    trans_image = torch.cat([trans_image, t], 0)
            print("233", trans_image.shape)

 

Tensorflow Medium gather And pytorch Medium gather The functions of are different , If pytorch There are and tf. Gather Functions with the same function can be used directly .

 

tf.gather()

  A one-dimensional

Multi dimension by row

Multidimensional by column

tf.gather_nd()

torch.gather()

 torch.gather(input, dim, index, out=None)

effect ： Collect the value of the specified position of the specific dimension entered

Parameters ：

input(tensor):   Number to be manipulated . Let's set its dimension as （x1, x2, …, xn）
dim(int):   Dimensions to be operated .
index(LongTensor):   How to input To operate . Its dimensions are limited , For example, when dim=i when ,index The dimensions are （x1, x2, …y, …,xn）, As a general input Of the i The size of the dimension is changed to y, And to satisfy y>=1（ Except for i Dimensions other than dimensions , Size and input bring into correspondence with ）.
out:   Pay attention to the output and index The dimensions of are consistent

Use torch.gather() Construction and tf.gather() Functions with the same function

According to the first 0 dimension

 

According to the first 1 dimension

According to the first 2 dimension

 

 

torch.gather() Realization tf.gather()