Torch tensor reshape by removing a channel

vision
1

so i’m training a 3D Unet where my input size is (1, 1, 64, 64, 64) and my output is (1, 2, 64, 64, 64). How do I get rid of one of the channels for the output so that i can compare the two?

I want my tensor to go from (1, 2, 64, 64, 64) to (1, 1, 64, 64, 64) by removing one of the channels.

2

Hi,

I think it is not good idea to just get rid of a channel. If your model generates outputs of size [1, 2, ...] then simply you can change last layer to generate 1 channel outputs. Otherwise, you have to find a function that maps 2 channel to 1 channel, otherwise you might lose a lot of information.
Doing it on UNet model is really easy.

Bests

3

so i tried this by implementing my own function and it doesn’t seem to work:

def combine_channel(output):

  first, second = output[0][0], output[0][1]

  first = torch.reshape(first, shape=(1, 262144))
  second = torch.reshape(second, shape=(1, 262144))
  arr = []

  for fir, sec in zip(first[0], second[0]):
    if(fir > sec):
      val = 0
      arr.append(val)
    else:
      val = 1
      arr.append(val)

  final = torch.reshape(torch.from_numpy(np.asarray(arr)), shape=(1, 262144))
  image = torch.reshape(final, shape=(1, 1, 64, 64, 64)).type(torch.float)
  return image

this just compares the two channels (since they’re probabilities of class 1 or 2) and then creates a new tensor with the larger one.

4

made a new issue for the new error i’m getting :slight_smile:

5

You can do this by torch.max

output = torch.randn(1, 2, 64, 64, 64)
output,_ = torch.max(output, 1, keepdim=True)
print(output.shape) # [1, 1, 64, 64, 64]