how I create the final result after the SOFTMAX function of the CNN network in binary image form (0,1)?. With the entry of this network is a single image and from this image we reserve (80% train, 20% test).
Help me plzzzz
You can use torch.argmax and either pass the logits or the probabilities to it in case you are working on a multi-class classification.
Assuming your model outputs logits in the shape [batch_size, nb_classes ,*], where * denotes additional dimensions, this would work:
output = model(input)
pred = torch.argmax(output, dim=1)