What I need to change the CNN model for a regression problem

akib62: <p>I have a CNN model to classify the MNIST digits problem. The model has 10 output nodes, used <code>CrossEntropyLoss</code> as a loss function and <code>Adam</code> optimizer. The structure of the model is given below</p>
<pre><code class="lang-auto">class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(         
            nn.Conv2d(
                in_channels=1,              
                out_channels=16,            
                kernel_size=5,              
                stride=1,                   
                padding=2,                  
            ),                              
            nn.ReLU(),                      
            nn.MaxPool2d(kernel_size=2),    
        )
        self.conv2 = nn.Sequential(         
            nn.Conv2d(16, 32, 5, 1, 2),     
            nn.ReLU(),                      
            nn.MaxPool2d(2),                
        )
        # fully connected layer, output 10 classes
        self.out = nn.Linear(32 * 7 * 7, 10)
        # self.softmax = torch.nn.Softmax(dim=1)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # flatten the output of conv2 to (batch_size, 32 * 7 * 7)
        x = x.view(x.size(0), -1)       
        output = self.out(x)
        # output = self.softmax(output)
        return output, x    # return x for visualization
</code></pre>
<p>As, we can use CNN to predict the <code>house price prediction</code> from images. Now, I want to consider this classification problem as a regression problem. I want to predict the digit as a number (whatever the model predicts, i.e, 1.0, 0.5, 8.5, etc). To do this I need only <code>1 output neuron</code> instead of <code>10</code>.</p>
<p>If I change this line for the expected output (1 regression output), I am getting an error (as the target labels are from 0 to 9).</p>
<pre><code class="lang-auto"># fully connected layer, 1 output
self.out = nn.Linear(32 * 7 * 7, 1)
</code></pre>
<p>Error</p>
<pre><code class="lang-auto">IndexError: Target 2 is out of bounds.
</code></pre>
<p>Could you tell me what I have to do, if I want to predict only one number as digits? Moreover, do I need to change the loss function to <code>MSELoss</code> and optimizer to <code>SGD</code>?</p>
<p>The data-loader</p>
<pre><code class="lang-auto">def data_loaders():
    train_data = datasets.MNIST(
        root = 'data',
        train = True,                         
        transform = transforms.ToTensor(), 
        download = True,            
    )
    test_data = datasets.MNIST(
        root = 'data', 
        train = False, 
        transform = transforms.ToTensor()
    )

    train_test_loaders = {
        'train' : torch.utils.data.DataLoader(train_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
        
        'test'  : torch.utils.data.DataLoader(test_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
    }
    return train_test_loaders
</code></pre>
<p>Later, at the time of training</p>
<pre><code class="lang-auto">def train(NB_EPOCS, model, loaders):
    model.train()
        
    # Train the model
    total_step = len(loaders['train'])
    for epoch in range(NB_EPOCS):
        for i, (images, labels) in enumerate(loaders['train']):
            b_x = Variable(images)   # batch x
            b_y = Variable(labels)   # batch y
            print(np.shape(b_x), np.shape(b_y))
            output = model(b_x)[0]               
            loss = criterion(output, b_y)
</code></pre>
<p>The shape of the <code>b_x</code> and <code>b_y</code> is</p>
<pre><code class="lang-auto">torch.Size([100, 1, 28, 28]) torch.Size([100])
</code></pre>

KFrank: <p>Hi Akib!</p>
<aside class="quote no-group" data-username="akib62" data-post="1" data-topic="132084" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3bc359/48.png" class="avatar"> akib62:</div>
<blockquote>
<p>I have a CNN model to classify the MNIST digits problem.</p>
</blockquote>
</aside>
<p>Perhaps some of the posts in the following thread could be helpful:</p>
<aside class="quote quote-modified" data-post="1" data-topic="131977">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3bc359/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-get-continuous-value-from-output-layer/131977">How to get continuous value from output layer</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    I have a code (from <a href="https://medium.com/@nutanbhogendrasharma/pytorch-convolutional-neural-network-with-mnist-dataset-4e8a4265e118" rel="noopener nofollow ugc">here</a>) to classify the MINST digits. The code is working fine. Here they used CrossEntropyLoss and Adam optimizer. 
The model code is given below 
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(         
            nn.Conv2d(
                in_channels=1,              
                out_channels=16,            
                kernel_size=5,              
                stride=1,                   
    …
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

akib62: <p><a class="mention" href="/u/kfrank">@KFrank</a> the link you gave me (this is my post and I haven’t received an answer for the given question)</p>