5
$\begingroup$

I know that fully convolutional networks can be used for image segmentation and similar but I wondered if you could also apply them to simple image classification tasks. And if so, what is the proper way to do it (in pytorch)?

Imagine that we have RGB images of cats and dogs of size 30x30 and the following (demo) networks.

class NormalNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.first_layer = nn.Conv2d(3, 1, (3, 3))
        self.fc = nn.Linear(784, 2)

    def forward(self, x):
        x = self.first_layer(x)
        x = x.reshape(-1, 784) 
        return self.fc(x)


class FullyConvolutionalNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Conv2d(3, 2, (30, 30))

    def forward(self, x):
        x = self.layer(x)
        return x.reshape(-1, 2)

Then you could use:

sample_image = torch.rand((10, 3, 30, 30))
normal_network = NormalNetwork()
fully_convolutional_network = FullyConvolutionalNetwork()
ret1 = normal_network(sample_image)
ret2 = fully_convolutional_network(sample_image)

and both of these networks would return a tensor of dimension torch.Size([10, 2]) which can then be fed into a BCELoss or similar. Is this a proper way to do it or is this bad for some reason?

$\endgroup$
4
$\begingroup$

First of all, here is the answer to your question: YES

You have a beginning of CNN (Convolutional Neural Network) in your code. Just don't forget to add activation functions (usually ReLU) after convolutional layers.

Another unconventional thing is you use 2 outputs for your binary classification, which is not the way to go, we would rather use 1 output that goes through sigmoid function and gives 1 if it is a cat and 0 if it is a dog (or reverse). And we apply BCE loss to train the network.

I would use transfer learning (use of pretrained architecture) for your task, VGG or ResNet are architectures that work well on classification tasks. We usually fine tune these pre-trained model to get the best results on classification tasks. Here is a Kaggle that may help you understand how to do it with Pytorch.

Hope this helps.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.