Employing UNet architecture to train my model

Bfuraide: <p>I have succeded in creating heatmaps to annotate key points on my medical cardiograms. I want to use UNet architecture to train my model to return heatmaps for each keyponts on the image. How can i go about creating this model. I have coded like below but it has failed to run withoit errors. Please advise.</p>
<pre><code>    for layer in self.decoder:
        if isinstance(layer, nn.ConvTranspose2d):
          x = layer(x)
          skip_connection = encoder_outputs.pop()
        # Adjust the size of the skip connection tensor
          skip_connection = F.interpolate(skip_connection, size=x.size()[2:], mode='bilinear', align_corners=False)
        # Ensure the number of channels matches for concatenation
          x = torch.cat([x, skip_connection], dim=1)
        else:
          x = layer(x)
    x = self.final_conv(x)
    return x
</code></pre>
<h1><a name="define-your-model-with-the-required-arguments-1" class="anchor" href="#define-your-model-with-the-required-arguments-1"></a>Define your model with the required arguments</h1>
<p>model = UNet(in_channels=3, out_channels=3, channels=[64, 128, 256, 512], strides=[2, 2, 2])</p>
<h1><a name="print-the-model-architecture-2" class="anchor" href="#print-the-model-architecture-2"></a>Print the model architecture</h1>
<p>print(model)</p>
<h1><a name="set-the-model-to-training-mode-3" class="anchor" href="#set-the-model-to-training-mode-3"></a>Set the model to training mode</h1>
<p>model.train()</p>
<h1><a name="define-optimizer-and-loss-function-4" class="anchor" href="#define-optimizer-and-loss-function-4"></a>Define optimizer and loss function</h1>
<p>optimizer = optim.Adam(model.parameters(), lr=1e-4)<br>
criterion = nn.MSELoss()</p>
<h1><a name="cast-images-to-torchfloat-before-passing-to-the-model-5" class="anchor" href="#cast-images-to-torchfloat-before-passing-to-the-model-5"></a>Cast images to torch.float before passing to the model</h1>
<h1><a name="training-loop-6" class="anchor" href="#training-loop-6"></a>Training loop</h1>
<p>num_epochs = 10<br>
for epoch in range(num_epochs):<br>
for images, true_heatmaps in data_loader:<br>
optimizer.zero_grad()<br>
# Cast images to torch.float before passing to the model<br>
images = images.float()</p>
<pre><code>    # Transpose the input tensor to match the expected shape (batch_size, channels, height, width)
    images = images.permute(0, 3, 1, 2)  # Assuming images has shape (batch_size, height, width, channels)


    # Predict
    pred_heatmaps = model(images)

    # Calculate loss
    loss = criterion(pred_heatmaps, true_heatmaps)

    # Backpropagation
    loss.backward()
    optimizer.step()

print(f"Epoch {epoch+1}, Loss: {loss.item()}")
</code></pre>

ptrblck: <p>What is the error you are seeing?</p>

Bfuraide: <p>Traceback (most recent call last)<br>
 in &lt;cell line: 66&gt;()<br>
75<br>
76         # Predict<br>
—&gt; 77         pred_heatmaps = model(images)<br>
78<br>
79         # Calculate loss</p>
<p>6 frames<br>
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)<br>
1509             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]<br>
1510         else:<br>
 → 1511             return self._call_impl(*args, **kwargs)<br>
1512<br>
1513     def _call_impl(self, *args, **kwargs):</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)<br>
1518                 or _global_backward_pre_hooks or _global_backward_hooks<br>
1519                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1520             return forward_call(*args, **kwargs)<br>
1521<br>
1522         try:</p>
<p> in forward(self, x)<br>
40               x = torch.cat([x, skip_connection], dim=1)<br>
41             else:<br>
—&gt; 42               x = layer(x)<br>
43         x = self.final_conv(x)<br>
44         return x</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)<br>
1509             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]<br>
1510         else:<br>
 → 1511             return self._call_impl(*args, **kwargs)<br>
1512<br>
1513     def _call_impl(self, *args, **kwargs):</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)<br>
1518                 or _global_backward_pre_hooks or _global_backward_hooks<br>
1519                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1520             return forward_call(*args, **kwargs)<br>
1521<br>
1522         try:</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in forward(self, input)<br>
458<br>
459     def forward(self, input: Tensor) → Tensor:<br>
 → 460         return self._conv_forward(input, self.weight, self.bias)<br>
461<br>
462 class Conv3d(_ConvNd):</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)<br>
454                             weight, bias, self.stride,<br>
455                             _pair(0), self.dilation, self.groups)<br>
 → 456         return F.conv2d(input, weight, bias, self.stride,<br>
457                         self.padding, self.dilation, self.groups)<br>
458</p>
<p>RuntimeError: Given groups=1, weight of size [256, 512, 3, 3], expected input[4, 768, 192, 256] to have 512 channels, but got 768 channels instead</p>

Bfuraide: <pre><code class="lang-auto"> Traceback (most recent call last)

</code></pre>
<pre><code class="lang-auto">[&lt;ipython-input-14-58faa5b773cb&gt;](https://localhost:8080/#) in &lt;cell line: 66&gt;()
     75 
     76         # Predict
---&gt; 77         pred_heatmaps = model(images)
     78 
     79         # Calculate loss

</code></pre>