udithhaputhanthri

class Generator(nn.Module):
  def __init__(self,noise_channels,img_channels,hidden_G):
    super(Generator,self).__init__()
    self.G=nn.Sequential(
        conv_trans_block(noise_channels,hidden_G*16,kernal_size=4,stride=1,padding=0),
        conv_trans_block(hidden_G*16,hidden_G*8),
        conv_trans_block(hidden_G*8,hidden_G*4),
        conv_trans_block(hidden_G*4,hidden_G*2),
        nn.ConvTranspose2d(hidden_G*2,img_channels,kernel_size=4,stride=2,padding=1),
        nn.Tanh()

udithhaputhanthri

D=Discriminator(img_channels,hidden_D).to(device)
G=Generator(noise_channels,img_channels,hidden_G).to(device)

criterion=nn.BCELoss()
opt_D=torch.optim.Adam(D.parameters(),lr=lr, betas=(0.5,0.999))
opt_G=torch.optim.Adam(G.parameters(),lr=lr, betas=(0.5,0.999))

udithhaputhanthri

class conv_trans_block(nn.Module):
  def __init__(self,in_channels,out_channels,kernal_size=4,stride=2,padding=1):
    super(conv_trans_block,self).__init__()
    self.block=nn.Sequential(
        nn.ConvTranspose2d(in_channels,out_channels,kernal_size,stride,padding),
        nn.BatchNorm2d(out_channels),
        nn.ReLU())
  def forward(self,x):
    return self.block(x)

udithhaputhanthri

import zipfile

!mkdir data_faces && wget https://s3-us-west-1.amazonaws.com/udacity-dlnfd/datasets/celeba.zip 
with zipfile.ZipFile("celeba.zip","r") as zip_ref:
  zip_ref.extractall("data_faces/")
  
from torch.utils.data import DataLoader
transform = transforms.Compose([
                               transforms.Resize((img_size,img_size)),
                               transforms.ToTensor()])

udithhaputhanthri

D.train()
G.train()
noise_for_generate=torch.randn(batch_size,noise_channels,1,1).to(device)
for epoch in range(epochs):
  for idx,(x,_) in enumerate(data_loader):
    x=x.to(device)
    x_len=x.shape[0]

    ### Train D
    D.zero_grad()

udithhaputhanthri

# for 128,128 size
class Generator(nn.Module):
  def __init__(self):
    super(Generator,self).__init__()

    self.enc_layer1=conv_block(4,8)
    self.enc_layer2=conv_block(8,16)
    self.enc_layer3=conv_block(16,32)
    self.enc_layer4=conv_block(32,64)
    self.enc_layer5=conv_block(64,128)

udithhaputhanthri

bce=nn.BCELoss().to(device)
L1=nn.L1Loss().to(device)
k=50
def criterion(y_hat,y):
  loss=bce(y_hat,y)+k*L1(y_hat,y)
  return loss

udithhaputhanthri

#G(x,z), D(x,y)
factor=1
G.train()
D.train()
for epoch in range(50):
  for i,(x,y) in enumerate(dataloader):
    opt_D.zero_grad()
    opt_G.zero_grad()

    x=x.to(device)

udithhaputhanthri

class get_dataset(torch.utils.data.Dataset):
  def __init__(self,name='edges2shoes',type_='train',transform=None):
    self.dir_=name+'/'+name+'/'+type_
    self.img_list=sorted(os.listdir(self.dir_))
    self.transform=transform
  def __len__(self):
    return len(self.img_list)
  def __getitem__(self,idx):
    both=plt.imread(self.dir_+'/'+self.img_list[idx]).astype('uint8')
    x=both[:,:both.shape[1]//2,:]

udithhaputhanthri

class conv_block(nn.Module):
  def __init__(self,in_channels,out_channels,kernel_size=4,stride=2,padding=1):
    super(conv_block,self).__init__()
    self.conv_block=nn.Sequential(
        nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding),
        nn.LeakyReLU(0.2),
        nn.BatchNorm2d(out_channels)
    )  
  def forward(self,x):
    return self.conv_block(x)