This is my GAN model code:
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
# Define hyperparameters
batch_size = 1 # I am putting it as 1 bcoz I want to save the images 1 by 1
image_size = 256
num_epochs = 200
noise_dim = 50
lr = 0.0002
beta1 = 0.5
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import glob
import os
import matplotlib.pyplot as plt
import numpy as np
import torchvision.utils as vutils
# Define a custom dataset for your SMT resistor x-ray images
class MyDataset(Dataset):
def __init__(self, path):
self.transforms = transforms.Compose([
transforms.Resize((image_size,image_size)), # resize to 64x64
transforms.Grayscale(), # convert to grayscale
transforms.ToTensor(), # convert to PyTorch tensor
transforms.Normalize((0.5,), (0.5,)) # normalize to range [-1, 1]
])
self.img_paths = sorted(glob.glob(os.path.join(path, '*.jpg'))) # get list of image paths
def __getitem__(self, index):
img_path = self.img_paths[index]
img = Image.open(img_path)
img = self.transforms(img)
return img
def __len__(self):
return len(self.img_paths)
# Set the path to your dataset
data_path = path_images
# Initialize the custom dataset
dataset = MyDataset(data_path)
# Initialize the dataloader
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
# Define the generator model
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.fc1 = nn.Linear(noise_dim, 128)
self.fc2 = nn.Linear(128, 256)
self.fc3 = nn.Linear(256, 512)
self.fc4 = nn.Linear(512, image_size*image_size)
def forward(self, x):
x = x.view(x.size(0), -1)
x = nn.functional.leaky_relu(self.fc1(x), 0.2)
x = nn.functional.leaky_relu(self.fc2(x), 0.2)
x = nn.functional.leaky_relu(self.fc3(x), 0.2)
x = nn.functional.tanh(self.fc4(x))
x = x.view(x.size(0), 1, image_size, image_size)
return x
# Define the discriminator model
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.conv1 = nn.Conv2d(1, 64, 4, 2, 1)
self.conv2 = nn.Conv2d(64, 128, 4, 2, 1)
self.conv3 = nn.Conv2d(128, 256, 4, 2, 1)
self.conv4 = nn.Conv2d(256, 512, 4, 2, 1)
self.fc = nn.Linear(image_size//16*image_size//16*512, 1)
def forward(self, x):
x = nn.functional.leaky_relu(self.conv1(x), 0.2)
x = nn.functional.leaky_relu(self.conv2(x), 0.2)
x = nn.functional.leaky_relu(self.conv3(x), 0.2)
x = nn.functional.leaky_relu(self.conv4(x), 0.2)
x = x.view(-1, image_size//16*image_size//16*512)
x = nn.functional.sigmoid(self.fc(x))
return x
# Initialize networks
G = Generator().cuda()
D = Discriminator().cuda()
# Set the loss function and optimizer
criterion = nn.BCELoss()
optimizer_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))
print(G)
print(D)
# Create a directory to save the generated images
os.makedirs('generated_images', exist_ok=True)
# Start the training loop
for epoch in range(num_epochs):
for i, images in enumerate(dataloader):
# Initialize the noise vector
noise = torch.randn(images.size(0), noise_dim, 1, 1).cuda()
# Train the discriminator with real images
D.zero_grad()
real_images = images.cuda()
real_labels = torch.ones((images.size(0), 1)).cuda()
real_outputs = D(real_images)
d_loss_real = criterion(real_outputs, real_labels.view(-1, 1))
d_loss_real.backward()
# Train the discriminator with fake images
fake_images = G(noise).detach()
fake_labels = torch.zeros((images.size(0), 1)).cuda()
fake_outputs = D(fake_images)
d_loss_fake = criterion(fake_outputs, fake_labels)
d_loss_fake.backward()
# Update the discriminator parameters
d_loss = (d_loss_real + d_loss_fake) / 2
optimizer_D.step()
# Train the generator
G.zero_grad()
noise = torch.randn(images.size(0), noise_dim, 1, 1).cuda()
fake_images = G(noise)
fake_labels = torch.ones((images.size(0), 1)).cuda()
fake_outputs = D(fake_images)
g_loss = criterion(fake_outputs, fake_labels.view(-1, 1))
g_loss.backward()
optimizer_G.step()
print(
"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
% (epoch, num_epochs, i, len(dataloader), d_loss.item(), g_loss.item())
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
fixed_noise = torch.randn(batch_size, noise_dim, 1, 1).to(device)
# Generate and save the images
generated_images = G(fixed_noise)
for j, img in enumerate(generated_images): # j = num of fake images
save_image(img, f'generated_images/image_{epoch}_{i}_{j}.png')
The generated_images is only able to save the first image from my dataset. I don't know which part is playing the role to save all the generated images from the original images. For example, I have 5 images in my dataset. When comes to the training loop, the generated_images should be showing all the 5 fake images, but it only shows the fake image for the first original image. Please help me to look into the code. Thanks in advance!
