Upload dcgan.py with huggingface_hub

dcgan.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+
+# The networks are taken from 
+# https://arxiv.org/abs/1511.06434
+class TBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size,stride,pad,norm_type):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.ConvTranspose2d(in_ch,out_ch,kernel_size,stride, pad,bias=False),
+            norm_layer(out_ch,norm_type),
+            nn.ReLU()
+        )
+    def forward(self,x):
+        return self.net(x)
+class CBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size,stride,pad,norm_type: str = "batch"):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_ch,out_ch,kernel_size,stride, pad, bias=False),
+            norm_layer(out_ch,norm_type),
+            nn.LeakyReLU(0.2)
+        )
+    def forward(self,x):
+        return self.net(x)   
+    
+class Generator(nn.Module):
+    #Outputs 64x64 pixel images
+
+    def __init__(
+        self,img_size=64,
+        out_ch=3,zdim=100,norm_type="BatchNorm2d",final_activation=None
+    ):
+        super().__init__()
+        # self.nf_g = nf_g
+        # self.z_dim = z_dim
+        # self.out_ch = out_ch
+        nf_g=2*img_size
+        self.final_activation=None if final_activation is None else getattr(torch,final_activation)
+
+        self.net = nn.Sequential(
+            # * Layer 1: 1x1
+            TBlock(zdim,8*nf_g, 4,1, 0,norm_type),
+            # * Layer 2: 4x4
+            TBlock(8*nf_g,4*nf_g,4,2,1,norm_type),
+            # * Layer 3: 8x8
+            TBlock(4*nf_g,2*nf_g,4,2,1,norm_type),
+            # * Layer 4: 16x16
+            TBlock(2*nf_g,nf_g,4,2,1,norm_type),
+            # * Layer 5: 32x32
+            nn.ConvTranspose2d(nf_g, out_ch, 4, 2, 1, bias=False),
+            # * Output: 64x64
+        )
+
+    def forward(self, x):
+        x = self.net(x)
+        return x if self.final_activation is None else self.final_activation(x)
+        
+        #return torch.tanh(x)
+      
+
+class Discriminator(nn.Module):
+    def __init__(self, img_size=64,in_ch=3,norm_type="BatchNorm2d",final_activation=None):
+        super().__init__()
+        nf_d=img_size
+        self.final_activation=None if final_activation is None else getattr(torch,final_activation)
+        self.net = nn.Sequential(
+            # * 64x64
+            nn.Conv2d(in_ch, nf_d, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2),
+            # * 32x32
+            CBlock(nf_d,2*nf_d,4,2,1,norm_type),
+            # * 16x16
+            CBlock(2*nf_d,4*nf_d,4,2,1,norm_type),
+            # * 8x8
+            CBlock(4*nf_d,8*nf_d,4,2,1,norm_type),
+            # * 4x4
+            nn.Conv2d(8*nf_d, 1, 4, 1, 0, bias=False),
+        )
+        
+
+    def forward(self, x):
+          x = self.net(x)
+          return x if self.final_activation is None else self.final_activation(x)
+       
+class norm_layer(nn.Module):
+    def __init__(self, num_channels,norm_type: str = None):
+        super().__init__()
+        if norm_type == "BatchNorm2d":
+            self.norm = nn.BatchNorm2d(num_channels)
+        elif norm_type == "GroupNorm":
+            self.norm = nn.GroupNorm(num_channels, num_channels)
+        elif norm_type is None or norm_type == "None":
+            self.norm=None
+        else:
+            raise ValueError(f"Unknown normalization type: {norm_type}")
+        
+    def forward(self, x):
+        return x if self.norm is None else self.norm(x)