Semantic Segmentation with Attention based CycleGAN

Idrees11: <p>Hi, Please read my scenario : I am having Original Images as Gray scale Images, and Target Images as gray scales images. However, Target Images are segmentation maps [background=1, segmentation object=2, another segmentation object=3]. I mean to say Target images are segmentation maps whos pixels range from 1 to 3 but Target or labelled images are also Gray scale images. I am designing custom Pixelwise Generator and Pixelwise Discriminators for my Model. Note I am using Cross entropy loss for Segmentation loss (for multi-class semantic segmentation). Therefore, I have to send that loss function gradients as Raw logits : Please see my code</p>
<p>class CustomSegmentationDataset(Dataset):<br>
def <strong>init</strong>(self, data_dir, labels_dir, image_transform=None, label_transform=None):<br>
self.data_dir = data_dir<br>
self.labels_dir = labels_dir<br>
self.image_transform = image_transform<br>
self.label_transform = label_transform</p>
<pre><code>    # List all image files in the data directory
    self.image_paths = [os.path.join(data_dir, fname)   for fname in os.listdir(data_dir) if fname.endswith('.jpg') or fname.endswith('.png')]
    self.label_paths = [os.path.join(labels_dir, fname) for fname in os.listdir(labels_dir) if fname.endswith('.jpg') or fname.endswith('.png')]

    # Ensure that the number of data and labels match
    assert len(self.image_paths) == len(self.label_paths), "Mismatch between number of images and labels"

def __len__(self):
    return len(self.image_paths)

def __getitem__(self, idx):
    image = Image.open(self.image_paths[idx]).convert('L')
    label = Image.open(self.label_paths[idx]).convert('L')  # Assuming labels are images
    
    # Apply the transformations if any
    if self.image_transform:
        image = self.image_transform(image)
    
    if self.label_transform:
        label = self.label_transform(label)
    return image, label   
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467026-dual-attention-channel-attention-spatial-attention-1" class="anchor" href="#p-467026-dual-attention-channel-attention-spatial-attention-1"></a>Dual ATTENTION  (Channel Attention + Spatial Attention)</h1>
<p><span class="hashtag-raw">#--------------------------------------------------------------------------------</span>#<br>
class DualAttention(nn.Module):<br>
def <strong>init</strong>(self, in_channels):<br>
super(DualAttention, self).<strong>init</strong>()<br>
self.channel_attention = ChannelAttention(in_channels)<br>
self.spatial_attention = SpatialAttention()</p>
<pre><code>def forward(self, x):
    x = self.channel_attention(x)
    x = self.spatial_attention(x)
    return x
</code></pre>
<p>class ChannelAttention(nn.Module):<br>
def <strong>init</strong>(self, in_channels):<br>
super(ChannelAttention, self).<strong>init</strong>()<br>
self.avg_pool = nn.AdaptiveAvgPool2d(1)<br>
self.max_pool = nn.AdaptiveMaxPool2d(1)<br>
self.fc = nn.Sequential(<br>
nn.Linear(in_channels, in_channels // 16, bias=False),<br>
nn.ReLU(inplace=True),<br>
nn.Linear(in_channels // 16, in_channels, bias=False),<br>
nn.Sigmoid()<br>
)</p>
<pre><code>def forward(self, x):
    avg_out = self.avg_pool(x)
    max_out = self.max_pool(x)
    out = avg_out + max_out
    out = out.view(out.size(0), -1)  # Flatten to (batch_size, channels)
    out = self.fc(out).view(out.size(0), out.size(1), 1, 1)
    return x * out  # Apply attention to the original feature map
</code></pre>
<p>class SpatialAttention(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(SpatialAttention, self).<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(2, 1, kernel_size=7, padding=3)<br>
self.sigmoid = nn.Sigmoid()</p>
<pre><code>def forward(self, x):
    avg_out = torch.mean(x, dim=1, keepdim=True)
    max_out, _ = torch.max(x, dim=1, keepdim=True)
    x_cat = torch.cat([avg_out, max_out], dim=1)
    out = self.conv1(x_cat)
    return x * self.sigmoid(out)  # Apply attention to the original feature map
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467026-custom-generator-encoder-decoder-based-2" class="anchor" href="#p-467026-custom-generator-encoder-decoder-based-2"></a>Custom Generator (Encoder Decoder Based)</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
class conv_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(in_c, out_c, kernel_size=3, padding=1, padding_mode=‘reflect’)<br>
self.bn1 = nn.BatchNorm2d(out_c, affine=True)<br>
self.relu = nn.ReLU(inplace=True)</p>
<pre><code>def forward(self, inputs):
    x = self.conv1(inputs)
    x = self.bn1(x)
    x = self.relu(x)
    return x
</code></pre>
<p>class conv1_cross_1(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super().<strong>init</strong>()<br>
self.conv1x1 = nn.Sequential(<br>
nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)<br>
)</p>
<pre><code>def forward(self, x):
    return self.conv1x1(x)
</code></pre>
<p>class encoder_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.conv = conv_block(in_c, out_c)<br>
self.attention = DualAttention(out_c)  # Add attention mechanism<br>
self.pool = nn.MaxPool2d((2, 2))</p>
<pre><code>def forward(self, inputs):
    x = self.conv(inputs)
    p = self.pool(x)
    return x, p
</code></pre>
<p>class decoder_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.up = nn.Sequential(<br>
nn.ConvTranspose2d(in_c, out_c, kernel_size=2, stride=2, padding=0),<br>
nn.BatchNorm2d(out_c, affine=True),<br>
nn.ReLU(inplace=True)<br>
)<br>
self.attention = DualAttention(out_c)  # Add attention mechanism</p>
<pre><code>def forward(self, inputs, skip):
    x = self.up(inputs)
    x = x + skip  # Summation-based skip connection
    x = self.attention(x)  # Apply attention after the skip connection
    return x
</code></pre>
<p>class classifier_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.outputs = nn.Sequential(<br>
nn.Conv2d(in_c, out_c, kernel_size=1, padding=0)<br>
)</p>
<pre><code>def forward(self, inputs):
    x = self.outputs(inputs)
    return x
</code></pre>
<p>class Generator(nn.Module):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()<br>
self.e1 = encoder_block(1, 96)<br>
self.e2 = encoder_block(96, 96)<br>
self.e3 = encoder_block(96, 96)<br>
self.e4 = encoder_block(96, 96)</p>
<pre><code>    # Bottleneck
    self.bottleneck = nn.Sequential(
        nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1, bias=False),
        nn.BatchNorm2d(96),
        nn.ReLU(inplace=True)
    )
    
    self.d1 = decoder_block(96, 96)
    self.d2 = decoder_block(96, 96)
    self.d3 = decoder_block(96, 96)
    self.d4 = decoder_block(96, 96)
    self.outputs = classifier_block(96, 3)  # Output channels = 3 for 3 classes

def forward(self, inputs):  
    
    # Encoder
    s1, p1 = self.e1(inputs)
    s2, p2 = self.e2(p1)
    s3, p3 = self.e3(p2)
    s4, p4 = self.e4(p3)
    
    # Bottleneck
    b = self.bottleneck(p4)
    
    # Decoder with skip connections
    d1 = self.d1(b, s4)
    d2 = self.d2(d1, s3)
    d3 = self.d3(d2, s2)
    d4 = self.d4(d3, s1)
    
    # Output
    x = self.outputs(d4)  # Output shape: (batch_size, 3, height, width)
    return x
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467026-custom-discriminator-encoder-decoder-based-3" class="anchor" href="#p-467026-custom-discriminator-encoder-decoder-based-3"></a>Custom Discriminator (encoder decoder based)</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467026-custom-pixelwise-encoder-decoder-based-discriminator-4" class="anchor" href="#p-467026-custom-pixelwise-encoder-decoder-based-discriminator-4"></a>custom pixelwise encoder-decoder based discriminator</h1>
<p><span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#<br>
class EncoderBlock(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super(EncoderBlock, self).<strong>init</strong>()<br>
self.conv = nn.Sequential(<br>
nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),<br>
nn.BatchNorm2d(out_channels),<br>
nn.LeakyReLU(0.2, inplace=True)<br>
)</p>
<pre><code>def forward(self, x):
    return self.conv(x)
</code></pre>
<p>class DecoderBlock(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super(DecoderBlock, self).<strong>init</strong>()<br>
self.conv = nn.Sequential(<br>
nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),<br>
nn.BatchNorm2d(out_channels),<br>
nn.ReLU(inplace=True)<br>
)</p>
<pre><code>def forward(self, x, skip):
    x = self.conv(x)
    x = x + skip  # addition based skip connection
    # x = torch.cat((x, skip), dim=1)  # Skip connection
    return x
</code></pre>
<p>class Discriminator(nn.Module):<br>
def <strong>init</strong>(self, in_channels=3):<br>
super(Discriminator, self).<strong>init</strong>()</p>
<pre><code>    # Encoder layers
    self.e1 = EncoderBlock(in_channels, 96)
    self.e2 = EncoderBlock(96, 96)
    self.e3 = EncoderBlock(96, 96)
    self.e4 = EncoderBlock(96, 96)
    
    # Bottleneck
    self.bottleneck = nn.Sequential(
        nn.Conv2d(96, 96, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(96),
        nn.LeakyReLU(0.2, inplace=True)
    )
    
    # Decoder layers (for skip connections)
    self.d1 = DecoderBlock(96, 96)
    self.d2 = DecoderBlock(96, 96)
    self.d3 = DecoderBlock(96, 96)
    self.d4 = DecoderBlock(96, 96)
    
    # Final output layer (pixel-wise)
    self.outputs = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0, bias=False)  # Output 1 channel per pixel

def forward(self, x):
    # Encoder forward pass
    s1 = self.e1(x)
    s2 = self.e2(s1)
    s3 = self.e3(s2)
    s4 = self.e4(s3)
    
    # Bottleneck (intermediate layer)
    b = self.bottleneck(s4)
    
    # Decoder with skip connections
    d1 = self.d1(b, s4)
    d2 = self.d2(d1, s3)
    d3 = self.d3(d2, s2)
    d4 = self.d4(d3, s1)
    
    # Final pixel-wise output (raw logits)
    x = self.outputs(d4)
    # print (x.shape)
    return x
</code></pre>
<p>But After running I am getting following error: Since Generator expects input shape 1 channel, and output shape as 3 (as I am having three classes in Labelled Gray image [1,2,3])… DIscriminator Input as 3 and output 1 But it shows following error</p>
<p>RuntimeError: Given groups=1, weight of size [96, 1, 3, 3], expected input[1, 3, 258, 258] to have 1 channels, but got 3 channels instead</p>

KFrank: <p>Hi Idrees!</p>
<aside class="quote no-group quote-modified" data-username="Idrees11" data-post="1" data-topic="217647">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/idrees11/48/74448_2.png" class="avatar"> Idrees11:</div>
<blockquote>
<p>self.conv1 = nn.Conv2d(in_c, out_c, kernel_size=3, padding=1, padding_mode=‘reflect’)<br>
…<br>
RuntimeError: Given groups=1, weight of size [96, 1, 3, 3], expected input[1, 3, 258, 258] to have 1 channels, but got 3 channels instead</p>
</blockquote>
</aside>
<p>Based on the code you’ve posted, the channels mismatch appears to be in the first layer<br>
of your <code>Generator</code>.  However, you don’t show where you instantiate a <code>Generator</code> nor to<br>
what image you might apply that <code>Generator</code>.</p>
<p>If, as might be typical for a CycleGAN, you are feeding the output of your “transformation”<br>
<code>Generator</code> to a second “reconstruction” <code>Generator</code> of the same architecture, you would<br>
be feeding a three-channel image to a generator that expects a single channel.</p>
<p>Print out the shapes of any images to which you are applying one of your <code>Generator</code>s,<br>
see if any of them do have three channels, and then ask how many channels they logically<br>
should have.</p>
<p>For your use case would it make sense to instantiate a <code>Generator</code> with three channels?<br>
Or could it make sense to sum or average the three channels of one of those “input”<br>
images to get a one-channel image compatible with the <code>Generator</code>?</p>
<p>Best.</p>
<p>K. Frank</p>

Idrees11: <p>Please see my full code:</p>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-dependencies-1" class="anchor" href="#p-467097-dependencies-1"></a>DEPENDENCIES</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
import torch<br>
import torch.nn as nn<br>
import torch.optim as optim<br>
from torchvision import models<br>
import torch.nn.functional as F<br>
from torch.utils.data import DataLoader, Dataset<br>
from torchvision import transforms<br>
import os<br>
from PIL import Image<br>
import torchvision.models as models<br>
from torchvision.models.segmentation import deeplabv3_resnet50<br>
from matplotlib import pyplot as plt</p>
<h1><a name="p-467097-from-sklearnmodel_selection-import-train_test_split-2" class="anchor" href="#p-467097-from-sklearnmodel_selection-import-train_test_split-2"></a>from sklearn.model_selection import train_test_split</h1>
<p>import random<br>
from torch.utils.data import random_split<br>
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, ReduceLROnPlateau<br>
import torch.nn.init as init<br>
import random<br>
from collections import deque<br>
import numpy as np<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-dataset-class-3" class="anchor" href="#p-467097-custom-dataset-class-3"></a>Custom Dataset Class</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
class CustomSegmentationDataset(Dataset):<br>
def <strong>init</strong>(self, data_dir, labels_dir, image_transform=None, label_transform=None):<br>
self.data_dir = data_dir<br>
self.labels_dir = labels_dir<br>
self.image_transform = image_transform<br>
self.label_transform = label_transform</p>
<pre><code>    # List all image files in the data directory
    self.image_paths = [os.path.join(data_dir, fname)   for fname in os.listdir(data_dir) if fname.endswith('.jpg') or fname.endswith('.png')]
    self.label_paths = [os.path.join(labels_dir, fname) for fname in os.listdir(labels_dir) if fname.endswith('.jpg') or fname.endswith('.png')]

    # Ensure that the number of data and labels match
    assert len(self.image_paths) == len(self.label_paths), "Mismatch between number of images and labels"

def __len__(self):
    return len(self.image_paths)

def __getitem__(self, idx):
    image = Image.open(self.image_paths[idx]).convert('L')
    label = Image.open(self.label_paths[idx]).convert('L')  # Assuming labels are images
    
    # Apply the transformations if any
    if self.image_transform:
        image = self.image_transform(image)
    
    if self.label_transform:
        label = self.label_transform(label)
    return image, label   
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-dual-attention-channel-attention-spatial-attention-4" class="anchor" href="#p-467097-dual-attention-channel-attention-spatial-attention-4"></a>Dual ATTENTION  (Channel Attention + Spatial Attention)</h1>
<p><span class="hashtag-raw">#--------------------------------------------------------------------------------</span>#<br>
class DualAttention(nn.Module):<br>
def <strong>init</strong>(self, in_channels):<br>
super(DualAttention, self).<strong>init</strong>()<br>
self.channel_attention = ChannelAttention(in_channels)<br>
self.spatial_attention = SpatialAttention()</p>
<pre><code>def forward(self, x):
    x = self.channel_attention(x)
    x = self.spatial_attention(x)
    return x
</code></pre>
<p>class ChannelAttention(nn.Module):<br>
def <strong>init</strong>(self, in_channels):<br>
super(ChannelAttention, self).<strong>init</strong>()<br>
self.avg_pool = nn.AdaptiveAvgPool2d(1)<br>
self.max_pool = nn.AdaptiveMaxPool2d(1)<br>
self.fc = nn.Sequential(<br>
nn.Linear(in_channels, in_channels // 16, bias=False),<br>
nn.ReLU(inplace=True),<br>
nn.Linear(in_channels // 16, in_channels, bias=False),<br>
nn.Sigmoid()<br>
)</p>
<pre><code>def forward(self, x):
    avg_out = self.avg_pool(x)
    max_out = self.max_pool(x)
    out = avg_out + max_out
    out = out.view(out.size(0), -1)  # Flatten to (batch_size, channels)
    out = self.fc(out).view(out.size(0), out.size(1), 1, 1)
    return x * out  # Apply attention to the original feature map
</code></pre>
<p>class SpatialAttention(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(SpatialAttention, self).<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(2, 1, kernel_size=7, padding=3)<br>
self.sigmoid = nn.Sigmoid()</p>
<pre><code>def forward(self, x):
    avg_out = torch.mean(x, dim=1, keepdim=True)
    max_out, _ = torch.max(x, dim=1, keepdim=True)
    x_cat = torch.cat([avg_out, max_out], dim=1)
    out = self.conv1(x_cat)
    return x * self.sigmoid(out)  # Apply attention to the original feature map
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-generator-encoder-decoder-based-5" class="anchor" href="#p-467097-custom-generator-encoder-decoder-based-5"></a>Custom Generator (Encoder Decoder Based)</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
class conv_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(in_c, out_c, kernel_size=3, padding=1, padding_mode=‘reflect’)<br>
self.bn1 = nn.BatchNorm2d(out_c, affine=True)<br>
self.relu = nn.ReLU(inplace=True)</p>
<pre><code>def forward(self, inputs):
    x = self.conv1(inputs)
    x = self.bn1(x)
    x = self.relu(x)
    return x
</code></pre>
<p>class conv1_cross_1(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super().<strong>init</strong>()<br>
self.conv1x1 = nn.Sequential(<br>
nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)<br>
)</p>
<pre><code>def forward(self, x):
    return self.conv1x1(x)
</code></pre>
<p>class encoder_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.conv = conv_block(in_c, out_c)<br>
self.attention = DualAttention(out_c)  # Add attention mechanism<br>
self.pool = nn.MaxPool2d((2, 2))</p>
<pre><code>def forward(self, inputs):
    x = self.conv(inputs)
    p = self.pool(x)
    return x, p
</code></pre>
<p>class decoder_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.up = nn.Sequential(<br>
nn.ConvTranspose2d(in_c, out_c, kernel_size=2, stride=2, padding=0),<br>
nn.BatchNorm2d(out_c, affine=True),<br>
nn.ReLU(inplace=True)<br>
)<br>
self.attention = DualAttention(out_c)  # Add attention mechanism</p>
<pre><code>def forward(self, inputs, skip):
    x = self.up(inputs)
    x = x + skip  # Summation-based skip connection
    x = self.attention(x)  # Apply attention after the skip connection
    return x
</code></pre>
<p>class classifier_block(nn.Module):<br>
def <strong>init</strong>(self, in_c, out_c):<br>
super().<strong>init</strong>()<br>
self.outputs = nn.Sequential(<br>
nn.Conv2d(in_c, out_c, kernel_size=1, padding=0)<br>
)</p>
<pre><code>def forward(self, inputs):
    x = self.outputs(inputs)
    return x
</code></pre>
<p>class Generator(nn.Module):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()<br>
self.e1 = encoder_block(1, 96)<br>
self.e2 = encoder_block(96, 96)<br>
self.e3 = encoder_block(96, 96)<br>
self.e4 = encoder_block(96, 96)</p>
<pre><code>    # Bottleneck
    self.bottleneck = nn.Sequential(
        nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1, bias=False),
        nn.BatchNorm2d(96),
        nn.ReLU(inplace=True)
    )
    
    self.d1 = decoder_block(96, 96)
    self.d2 = decoder_block(96, 96)
    self.d3 = decoder_block(96, 96)
    self.d4 = decoder_block(96, 96)
    self.outputs = classifier_block(96, 3)  # Output channels = 3 for 3 classes

def forward(self, inputs):  
    
    # Encoder
    s1, p1 = self.e1(inputs)
    s2, p2 = self.e2(p1)
    s3, p3 = self.e3(p2)
    s4, p4 = self.e4(p3)
    
    # Bottleneck
    b = self.bottleneck(p4)
    
    # Decoder with skip connections
    d1 = self.d1(b, s4)
    d2 = self.d2(d1, s3)
    d3 = self.d3(d2, s2)
    d4 = self.d4(d3, s1)
    
    # Output
    x = self.outputs(d4)  # Output shape: (batch_size, 3, height, width)
    # print ('at generation level')
    # print (x.shape)
    return x
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-discriminator-encoder-decoder-based-6" class="anchor" href="#p-467097-custom-discriminator-encoder-decoder-based-6"></a>Custom Discriminator (encoder decoder based)</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-pixelwise-encoder-decoder-based-discriminator-7" class="anchor" href="#p-467097-custom-pixelwise-encoder-decoder-based-discriminator-7"></a>custom pixelwise encoder-decoder based discriminator</h1>
<p><span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#<br>
class EncoderBlock(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super(EncoderBlock, self).<strong>init</strong>()<br>
self.conv = nn.Sequential(<br>
nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),<br>
nn.BatchNorm2d(out_channels),<br>
nn.LeakyReLU(0.2, inplace=True)<br>
)</p>
<pre><code>def forward(self, x):
    return self.conv(x)
</code></pre>
<p>class DecoderBlock(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super(DecoderBlock, self).<strong>init</strong>()<br>
self.conv = nn.Sequential(<br>
nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),<br>
nn.BatchNorm2d(out_channels),<br>
nn.ReLU(inplace=True)<br>
)</p>
<pre><code>def forward(self, x, skip):
    x = self.conv(x)
    x = x + skip  # addition based skip connection
    # x = torch.cat((x, skip), dim=1)  # Skip connection
    return x
</code></pre>
<p>class Discriminator(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Discriminator, self).<strong>init</strong>()</p>
<pre><code>    # Encoder layers
    self.e1 = EncoderBlock(3, 96)
    self.e2 = EncoderBlock(96, 96)
    self.e3 = EncoderBlock(96, 96)
    self.e4 = EncoderBlock(96, 96)
    
    # Bottleneck
    self.bottleneck = nn.Sequential(
        nn.Conv2d(96, 96, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(96),
        nn.LeakyReLU(0.2, inplace=True)
    )
    
    # Decoder layers (for skip connections)
    self.d1 = DecoderBlock(96, 96)
    self.d2 = DecoderBlock(96, 96)
    self.d3 = DecoderBlock(96, 96)
    self.d4 = DecoderBlock(96, 96)
    
    # Final output layer (pixel-wise)
    self.outputs = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0, bias=False)  # Output 1 channel per pixel

def forward(self, x):
    # Encoder forward pass
    # print ("before discrinaation",x.shape)
    s1 = self.e1(x)
    s2 = self.e2(s1)
    s3 = self.e3(s2)
    s4 = self.e4(s3)
    # print ('before bootleneck at disc level',s4.shape)
    # Bottleneck (intermediate layer)
    b = self.bottleneck(s4)
    # print ('After bootleneck at disc level',b.shape)
    
    # Decoder with skip connections
    d1 = self.d1(b, s4)
    d2 = self.d2(d1, s3)
    d3 = self.d3(d2, s2)
    d4 = self.d4(d3, s1)
    
    # Final pixel-wise output (raw logits)
    x = self.outputs(d4)
    # print ('at discriminator level')
    # print (x.shape)
    return x
</code></pre>
<p><span class="hashtag-raw">#-----------------------------------------------------------------------------</span>#</p>
<p><span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-segmentation-unet-8" class="anchor" href="#p-467097-custom-segmentation-unet-8"></a>Custom Segmentation Unet</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
class SegmentationModel(nn.Module):<br>
def <strong>init</strong>(self, in_channels=1, out_channels=3):<br>
super(SegmentationModel, self).<strong>init</strong>()<br>
self.encoder1 = self.conv_block(in_channels, 96)<br>
self.encoder2 = self.conv_block(96, 96)<br>
self.encoder3 = self.conv_block(96, 96)<br>
self.encoder4 = self.conv_block(96, 96)<br>
self.bottleneck = self.conv_block(96, 96)<br>
self.decoder4 = self.conv_block(96, 96)<br>
self.decoder3 = self.conv_block(96, 96)<br>
self.decoder2 = self.conv_block(96, 96)<br>
self.decoder1 = self.conv_block(96, 96)<br>
self.final_conv = nn.Conv2d(96, out_channels, kernel_size=1)</p>
<pre><code>def conv_block(self, in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True)
    )

def forward(self, x):
    enc1 = self.encoder1(x)
    enc2 = self.encoder2(enc1)
    enc3 = self.encoder3(enc2)
    enc4 = self.encoder4(enc3)
    bottleneck = self.bottleneck(enc4)
    dec4 = self.decoder4(bottleneck)
    dec3 = self.decoder3(dec4 + enc4)
    dec2 = self.decoder2(dec3 + enc3)
    dec1 = self.decoder1(dec2 + enc2)
    x = self.final_conv(dec1 + enc1)
    # print ('at segmentation')
    # print (x.shape)
    return x
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-custom-cyclegan-9" class="anchor" href="#p-467097-custom-cyclegan-9"></a>Custom CycleGAN</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
class CycleGAN(nn.Module):<br>
def <strong>init</strong>(self, num_classes):<br>
super(CycleGAN, self).<strong>init</strong>()<br>
self.G_A2B = Generator()<br>
self.G_B2A = Generator()<br>
self.D_A = Discriminator()<br>
self.D_B = Discriminator()<br>
self.Seg = SegmentationModel()<br>
self.criterion_adv = nn.MSELoss()<br>
self.criterion_cycle = nn.L1Loss()<br>
self.criterion_identity = nn.L1Loss()<br>
self.criterion_seg = nn.CrossEntropyLoss()</p>
<pre><code>def forward(self, real_A, real_B, seg_A, seg_B):
    fake_B = self.G_A2B(real_A)
    fake_A = self.G_B2A(real_B)
    D_A_fake = self.D_A(fake_A.detach())
    D_B_fake = self.D_B(fake_B.detach())
    cycle_A = self.G_B2A(fake_B)
    cycle_B = self.G_A2B(fake_A)
    identity_A = self.G_B2A(real_A)
    identity_B = self.G_A2B(real_B)
    D_A_real = self.D_A(real_A)
    D_B_real = self.D_B(real_B)
    seg_fake_B = self.Seg(fake_B)
    seg_fake_A = self.Seg(fake_A)
    return {
        'fake_A': fake_A,
        'fake_B': fake_B,
        'cycle_A': cycle_A,
        'cycle_B': cycle_B,
        'identity_A': identity_A,
        'identity_B': identity_B,
        'D_A_real': D_A_real,
        'D_A_fake': D_A_fake,
        'D_B_real': D_B_real,
        'D_B_fake': D_B_fake,
        'seg_fake_A': seg_fake_A,
        'seg_fake_B': seg_fake_B
    }
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-iou-calculation-10" class="anchor" href="#p-467097-iou-calculation-10"></a>IoU Calculation</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
def calculate_iou(pred, target, num_classes):<br>
pred = pred.view(-1)             # Flatten the prediction tensor<br>
target = target.view(-1)         # Flatten the target tensor<br>
ious = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
for c in range(num_classes):<br>
pred_c = (pred == c)         # Boolean mask for predicted class<br>
target_c = (target == c)     # Boolean mask for target class</p>
<pre><code>    intersection = (pred_c &amp; target_c).sum().float()  # Intersection
    union = (pred_c | target_c).sum().float()         # Union

    # Handle the case where union is zero (no pixels for this class)
    if union == 0:
        ious.append(0.0)  # IoU is 0 if there are no pixels for this class
    else:
        iou = intersection / union
        ious.append(iou)
# Compute the mean IoU across all classes
mean_iou = torch.mean(torch.tensor(ious))
return mean_iou  # Mean IoU is guaranteed to be in [0, 1]
</code></pre>
<p><span class="hashtag-raw">#-------------------------------------------------------------------------------</span>#</p>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-training-loop-11" class="anchor" href="#p-467097-training-loop-11"></a>Training Loop</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
def train_cyclegan(real_A, real_B, seg_A, seg_B, model, optimizer_G, optimizer_D, epoch, batch_idx):<br>
model.train()<br>
device = next(model.parameters()).device</p>
<pre><code># Move data to the correct device
real_A = real_A.to(device)  # No need for requires_grad=True for inputs
real_B = real_B.to(device)  # No need for requires_grad=True for inputs
seg_A = seg_A.long().to(device) if seg_A is not None else None
seg_B = seg_B.long().to(device) if seg_B is not None else None

# Forward pass
outputs = model(real_A, real_B, seg_A, seg_B)

# Generator losses
loss_G_A2B = model.criterion_adv(outputs['D_B_fake'], torch.ones_like(outputs['D_B_fake']).to(device))
loss_G_B2A = model.criterion_adv(outputs['D_A_fake'], torch.ones_like(outputs['D_A_fake']).to(device))
loss_G_adv = loss_G_A2B + loss_G_B2A

loss_cycle_A = model.criterion_cycle(outputs['cycle_A'], real_A)
loss_cycle_B = model.criterion_cycle(outputs['cycle_B'], real_B)
loss_cycle = loss_cycle_A + loss_cycle_B

loss_identity_A = model.criterion_identity(outputs['identity_A'], real_A)
loss_identity_B = model.criterion_identity(outputs['identity_B'], real_B)
loss_identity = loss_identity_A + loss_identity_B

# Segmentation loss and IoU calculation
loss_seg = 0.0
mean_iou = 0.0
if seg_A is not None and seg_B is not None:
    seg_fake_A = outputs['seg_fake_A'].float()
    seg_fake_B = outputs['seg_fake_B'].float()
    seg_A = seg_A.squeeze(1).long()  # Ensure it is in shape (batch_size, height, width)
    seg_B = seg_B.squeeze(1).long()  # Ensure it is in shape (batch_size, height, width)

    loss_seg_A = model.criterion_seg(seg_fake_A, seg_A)
    loss_seg_B = model.criterion_seg(seg_fake_B, seg_B)
    loss_seg = loss_seg_A + loss_seg_B

    pred_A = torch.argmax(outputs['seg_fake_A'], dim=1)  # Get the predicted classes
    pred_B = torch.argmax(outputs['seg_fake_B'], dim=1)  # Same for B

    mean_iou_A = calculate_iou(pred_A, seg_A, num_classes=3)
    mean_iou_B = calculate_iou(pred_B, seg_B, num_classes=3)
    mean_iou = (mean_iou_A + mean_iou_B) / 2

# Total generator loss
loss_G = loss_G_adv + 10.0 * loss_cycle + 5.0 * loss_identity + 1.0 * loss_seg

# Discriminator losses
loss_D_A = model.criterion_adv(outputs['D_A_real'], torch.ones_like(outputs['D_A_real']).to(device)) + \
           model.criterion_adv(outputs['D_A_fake'], torch.zeros_like(outputs['D_A_fake']).to(device))
loss_D_B = model.criterion_adv(outputs['D_B_real'], torch.ones_like(outputs['D_B_real']).to(device)) + \
           model.criterion_adv(outputs['D_B_fake'], torch.zeros_like(outputs['D_B_fake']).to(device))
loss_D = loss_D_A + loss_D_B

# Backpropagation
optimizer_G.zero_grad()
loss_G.backward(retain_graph=True)  # Retain the graph for the next backward pass
optimizer_G.step()

optimizer_D.zero_grad()
loss_D.backward()  # No need to retain the graph here
optimizer_D.step()

return loss_G.item(), loss_D.item(), mean_iou
</code></pre>
<p><span class="hashtag-raw">#----------------------------------------------------------------------------------</span>#<br>
def evaluate_cyclegan(real_A, real_B, seg_A, seg_B, model):<br>
# Forward pass<br>
outputs = model(real_A, real_B, seg_A, seg_B)</p>
<pre><code># Generator losses
loss_G_A2B = model.criterion_adv(outputs['D_B_fake'], torch.ones_like(outputs['D_B_fake']).to(device))
loss_G_B2A = model.criterion_adv(outputs['D_A_fake'], torch.ones_like(outputs['D_A_fake']).to(device))
loss_G_adv = loss_G_A2B + loss_G_B2A

loss_cycle_A = model.criterion_cycle(outputs['cycle_A'], real_A)
loss_cycle_B = model.criterion_cycle(outputs['cycle_B'], real_B)
loss_cycle = loss_cycle_A + loss_cycle_B

loss_identity_A = model.criterion_identity(outputs['identity_A'], real_A)
loss_identity_B = model.criterion_identity(outputs['identity_B'], real_B)
loss_identity = loss_identity_A + loss_identity_B

# Segmentation loss and IoU calculation
loss_seg = 0.0
mean_iou = 0.0
if seg_A is not None and seg_B is not None:
    seg_fake_A = outputs['seg_fake_A'].float()
    seg_fake_B = outputs['seg_fake_B'].float()
    seg_A = seg_A.squeeze(1).long()
    seg_B = seg_B.squeeze(1).long()

    loss_seg_A = model.criterion_seg(seg_fake_A, seg_A)
    loss_seg_B = model.criterion_seg(seg_fake_B, seg_B)
    loss_seg = loss_seg_A + loss_seg_B

    pred_A = torch.argmax(outputs['seg_fake_A'], dim=1)
    pred_B = torch.argmax(outputs['seg_fake_B'], dim=1)

    mean_iou_A = calculate_iou(pred_A, seg_A, num_classes=3)
    mean_iou_B = calculate_iou(pred_B, seg_B, num_classes=3)
    mean_iou = (mean_iou_A + mean_iou_B) / 2

# Total generator loss
loss_G = loss_G_adv + 10.0 * loss_cycle + 5.0 * loss_identity + 1.0 * loss_seg

# Discriminator losses (but we don’t need to calculate gradients here)
loss_D_A = model.criterion_adv(outputs['D_A_real'], torch.ones_like(outputs['D_A_real']).to(device)) + \
           model.criterion_adv(outputs['D_A_fake'], torch.zeros_like(outputs['D_A_fake']).to(device))
loss_D_B = model.criterion_adv(outputs['D_B_real'], torch.ones_like(outputs['D_B_real']).to(device)) + \
           model.criterion_adv(outputs['D_B_fake'], torch.zeros_like(outputs['D_B_fake']).to(device))
loss_D = loss_D_A + loss_D_B

return loss_G.item(), loss_D.item(), mean_iou
</code></pre>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-initialisation-of-weights-zero-centered-initialization-12" class="anchor" href="#p-467097-initialisation-of-weights-zero-centered-initialization-12"></a>INITIALISATION OF WEIGHTS  Zero-Centered Initialization</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
def zero_centered_init(m):<br>
“”"<br>
Apply zero-centered initialization to the weights of a layer.<br>
“”"<br>
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)):<br>
nn.init.normal_(m.weight, mean=0.0, std=0.02)  # Zero-centered initialization<br>
if m.bias is not None:<br>
nn.init.zeros_(m.bias)                     # Initialize biases to zero<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-main-execution-13" class="anchor" href="#p-467097-main-execution-13"></a>Main Execution</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
if <strong>name</strong> == “<strong>main</strong>”:<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-hyperparameters-14" class="anchor" href="#p-467097-hyperparameters-14"></a>Hyperparameters</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
num_classes = 3<br>
num_epochs = 1<br>
batch_size = 1<br>
learning_rate_G = 0.0001<br>
learning_rate_D = 0.0001<br>
replay_buffer_size = 50  # Size of the replay buffer<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-initialize-model-and-optimizers-15" class="anchor" href="#p-467097-initialize-model-and-optimizers-15"></a>Initialize model and optimizers</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
model = CycleGAN(num_classes).to(device)<br>
model.apply(zero_centered_init)<br>
# model.apply(init_weights)<br>
optimizer_G = optim.Adam(list(model.G_A2B.parameters()) + list(model.G_B2A.parameters()), lr=learning_rate_G, betas=(0.5, 0.999))<br>
optimizer_D = optim.Adam(list(model.D_A.parameters()) + list(model.D_B.parameters()), lr=learning_rate_D, betas=(0.5, 0.999))<br>
<span class="hashtag-raw">#---------------------------------------------------------</span>#<br>
#   SCHEDULERS<br>
<span class="hashtag-raw">#----------------</span>#<br>
scheduler_G = StepLR(optimizer_G, step_size=30, gamma=0.1)<br>
scheduler_D = StepLR(optimizer_D, step_size=30, gamma=0.1)<br>
<span class="hashtag-raw">#------------------------------------------------------------------------------</span>#</p>
<h1><a name="p-467097-dataset-and-dataloader-16" class="anchor" href="#p-467097-dataset-and-dataloader-16"></a>Dataset and DataLoader</h1>
<p><span class="hashtag-raw">#------------------------------------------------------------------------------</span>#<br>
data = r’C:\Users\Idrees Bhat\Desktop\reduced_dataset\source’<br>
labels = r’C:\Users\Idrees Bhat\Desktop\reduced_dataset\target’<br>
image_transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5], std=[0.5]),])<br>
label_transform = transforms.Compose([transforms.ToTensor()])<br>
# Load the full dataset<br>
dataset = CustomSegmentationDataset(data,labels,image_transform=image_transform, label_transform=label_transform )</p>
<pre><code># Set the split ratios
train_size = int(0.8 * len(dataset))   # 80% for training
val_size   = int(0.1 * len(dataset))   # 10% for validation
test_size  = len(dataset) - train_size - val_size  # Remaining 10% for testing
# test_size = len(dataset) - train_size
# Split the dataset
train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])
                                                      
# Create DataLoader for each split
train_loader = DataLoader(train_dataset, batch_size= batch_size, shuffle=True)
val_loader   = DataLoader(val_dataset,batch_size= batch_size,    shuffle=False)
test_loader  = DataLoader(test_dataset,batch_size= batch_size,   shuffle=False)
# print ('')
# print (type(train_data))
#-------------------------------------------------------------------------------------------------------#
#  tracking losses
#--------------------------------------------------------------------------------------------------------#
# Initialize lists to store values
gen_losses = []
disc_losses = []
iou_values = []
lr_G_values = []
lr_D_values = []
</code></pre>
<p><span class="hashtag-raw">#-----------------------------------------------------------------------------------------------------</span>------#</p>
<h1><a name="p-467097-training-loop-17" class="anchor" href="#p-467097-training-loop-17"></a>TRAINING LOOP</h1>
<p><span class="hashtag-raw">#-----------------------------------------------------------------------------------------------------</span>-------#<br>
# Training Loop<br>
for epoch in range(num_epochs):<br>
total_gen_loss = 0.0<br>
total_disc_loss = 0.0<br>
total_iou = 0.0<br>
for i, (real_A, real_B) in enumerate(train_loader):<br>
real_A = real_A.to(device)<br>
real_B = real_B.to(device)<br>
# Forward pass to generate fake images<br>
fake_B = model.G_A2B(real_A)<br>
fake_A = model.G_B2A(real_B)</p>
<pre><code>        # Pass real_A and real_B as segmentation maps
        loss_G, loss_D, mean_iou = train_cyclegan(real_A, real_B, real_B, real_B, model, optimizer_G, optimizer_D,epoch, i)
        
         # Get intermediate outputs from the model
        outputs = model(real_A, real_B, real_B, real_B)

        if i % 5 == 0:
             print(f"Epoch [{epoch}/{num_epochs}] Batch [{i}/{len(train_loader)}] Loss_G: {loss_G:.4f} Loss_D: {loss_D:.4f} Mean IoU: {mean_iou:.4f}")
             
        scheduler_G.step()
        scheduler_D.step()
</code></pre>

Idrees11: <p>I am doing sementic segmentation with Attention based CycleGAN. I am expecting Gray scale image as input to Generator since my target images are also Gray scale images with semantic segmentation map [1,2,3]. As a multiclass seg problem My Generators output shape should be equal to 3 and Discriminator also expect 3 shape as input and produce at last 1 shape as out put… Here it is showing error. But When I am running with Generator output as 1 shape Discriminators input shape as 1 it works fine… Note, I am using Cross entropy loss fucntion as additional seg loss .</p>

KFrank: <p>Hi Idrees!</p>
<aside class="quote no-group" data-username="Idrees11" data-post="3" data-topic="217647">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/idrees11/48/74448_2.png" class="avatar"> Idrees11:</div>
<blockquote>
<p>Please see my full code:</p>
</blockquote>
</aside>
<p>More code isn’t helpful.</p>
<p>My recommendation:</p>
<p>Find the spot in your code where the error actually occurs.  It is probably something like:</p>
<pre data-code-wrap="python"><code class="lang-python">output = myGenerator (myThreeChannelImage)
</code></pre>
<p>where <code>myThreeChannelImage</code> is a tensor of shape <code>[1, 3, 258, 258]</code>, as reported in<br>
the error message you posted.  (That is, <code>myThreeChannelImage</code> is a batch containing<br>
a single sample that is a 258x258 three-channel image).  <code>myGenerator</code> would be one<br>
of the instances of your <code>Generator</code> class in your CycleGAN.</p>
<p>Then verify that this is causing the issue with something like:</p>
<pre data-code-wrap="python"><code class="lang-python">singleChannelTestImage = torch.randn (1, 1, 258, 258)
print (singleChannelTestImage.shape)
print ('calling myGenerator (singleChannelTestImage) ...')
_ = myGenerator (singleChannelTestImage)
print ('myGenerator (singleChannelTestImage) returned.')

print (myThreeChannelImage.shape)
print ('calling myGenerator (myThreeChannelImage) ...')
output = myGenerator (myThreeChannelImage)
print ('expected error before here')
</code></pre>
<p>If this confirms the analysis of the cause of your error, you then need to understand why<br>
you are trying to pass a three-channel image to a model that expects a single-channel<br>
input.  Based on the purpose of this CycleGAN, how many channels should this image<br>
have?</p>
<p>Best.</p>
<p>K. Frank</p>