models.py

# -*- coding: utf-8 -*-
"""Untitled10.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1jP0hfA5O4j1AounpfAvq69V5YAC7RT8Z
"""

# models.py
import torch.nn as nn
from channel import power_normalize, awgn

class Encoder(nn.Module):
    """
    Fig.2 Encoder:
      5x5x16/2 -> 5x5x32/2 -> 5x5x32/1 -> 5x5x32/1 -> 5x5xc/1
      each: conv + PReLU
    """
    def __init__(self, latent_ch=8):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, stride=2, padding=2),
            nn.PReLU(16),

            nn.Conv2d(16, 32, kernel_size=5, stride=2, padding=2),
            nn.PReLU(32),

            nn.Conv2d(32, 32, kernel_size=5, stride=1, padding=2),
            nn.PReLU(32),

            nn.Conv2d(32, 32, kernel_size=5, stride=1, padding=2),
            nn.PReLU(32),

            nn.Conv2d(32, latent_ch, kernel_size=5, stride=1, padding=2),
            nn.PReLU(latent_ch),
        )

    def forward(self, x):
        return self.net(x)

class Decoder(nn.Module):
    """
    Fig.2 Decoder:
      5x5x32/1 -> 5x5x32/1 -> 5x5x32/1 -> 5x5x16/2 -> 5x5x3/2(sigmoid)
      trans conv + PReLU, last: trans conv + sigmoid
    """
    def __init__(self, latent_ch=8):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose2d(latent_ch, 32, kernel_size=5, stride=1, padding=2, output_padding=0),
            nn.PReLU(32),

            nn.ConvTranspose2d(32, 32, kernel_size=5, stride=1, padding=2, output_padding=0),
            nn.PReLU(32),

            nn.ConvTranspose2d(32, 32, kernel_size=5, stride=1, padding=2, output_padding=0),
            nn.PReLU(32),

            # 8x8 -> 16x16 (stride=2). output_padding=1 to match size exactly
            nn.ConvTranspose2d(32, 16, kernel_size=5, stride=2, padding=2, output_padding=1),
            nn.PReLU(16),

            # 16x16 -> 32x32 (stride=2). output_padding=1 to match size exactly
            nn.ConvTranspose2d(16, 3, kernel_size=5, stride=2, padding=2, output_padding=1),
            nn.Sigmoid(),
        )

    def forward(self, z):
        return self.net(z)

class DeepJSCC(nn.Module):
    def __init__(self, latent_ch=8):
        super().__init__()
        self.enc = Encoder(latent_ch)
        self.dec = Decoder(latent_ch)

    def forward(self, x, snr_db):
        z = self.enc(x)
        z = power_normalize(z)
        y = awgn(z, snr_db)
        xhat = self.dec(y)
        return xhat