Generative Adversarial Network-Based Simulator for Training Purposes,Python
👤 Sharing: AI
```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
# Define the Generator Network
class Generator(nn.Module):
def __init__(self, noise_dim, output_dim):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(noise_dim, 128),
nn.ReLU(),
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, output_dim),
nn.Tanh() # Output values between -1 and 1 (suitable for scaled data)
)
def forward(self, z):
return self.model(z)
# Define the Discriminator Network
class Discriminator(nn.Module):
def __init__(self, input_dim):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, 256),
nn.LeakyReLU(0.2),
nn.Linear(256, 128),
nn.LeakyReLU(0.2),
nn.Linear(128, 1),
nn.Sigmoid() # Output a probability (real or fake)
)
def forward(self, x):
return self.model(x)
# --- Training Parameters ---
noise_dim = 10 # Dimension of the noise vector (input to the generator)
data_dim = 2 # Dimension of the data we want to simulate (e.g., 2D points)
batch_size = 64
epochs = 500
learning_rate = 0.0002
# --- Create the Networks ---
generator = Generator(noise_dim, data_dim)
discriminator = Discriminator(data_dim)
# --- Optimizers ---
generator_optimizer = optim.Adam(generator.parameters(), lr=learning_rate)
discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=learning_rate)
# --- Loss Function ---
loss_function = nn.BCELoss() # Binary Cross-Entropy Loss (for classification)
# --- Move networks to GPU if available ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = generator.to(device)
discriminator = discriminator.to(device)
loss_function = loss_function.to(device)
# --- Training Loop ---
def train(generator, discriminator, generator_optimizer, discriminator_optimizer, loss_function, noise_dim, data_dim, batch_size, epochs):
real_data_points = [] # for demonstration purposes only
generated_data_points = [] # for demonstration purposes only
for epoch in range(epochs):
# --- Generate Real Data (replace with your actual dataset loading) ---
# Here we'll simulate some real data from a Gaussian distribution
# In a real application, you'd load your training data here.
real_data = torch.randn(batch_size, data_dim).to(device) # Standard Gaussian
# Adding some structure to the real data:
real_data[:, 0] = real_data[:, 0] + 2 # Shift x-values slightly
real_data[:, 1] = real_data[:, 1] * 0.5 # Scale y-values
real_labels = torch.ones(batch_size, 1).to(device) # Label real data as 1
# --- Train the Discriminator ---
discriminator_optimizer.zero_grad()
# 1. Train on real data
outputs_real = discriminator(real_data)
loss_real = loss_function(outputs_real, real_labels)
loss_real.backward()
# 2. Train on fake data
noise = torch.randn(batch_size, noise_dim).to(device)
fake_data = generator(noise)
fake_labels = torch.zeros(batch_size, 1).to(device) # Label fake data as 0
outputs_fake = discriminator(fake_data)
loss_fake = loss_function(outputs_fake, fake_labels)
loss_fake.backward()
discriminator_optimizer.step() # Update discriminator weights
discriminator_loss = loss_real + loss_fake # Total discriminator loss
# --- Train the Generator ---
generator_optimizer.zero_grad()
# Generate fake data again
noise = torch.randn(batch_size, noise_dim).to(device)
fake_data = generator(noise)
outputs = discriminator(fake_data) # Discriminator's opinion on the generated data
loss_generator = loss_function(outputs, torch.ones(batch_size, 1).to(device)) # Fool the discriminator (make it think they are real)
loss_generator.backward()
generator_optimizer.step() # Update generator weights
if epoch % 50 == 0:
print(f"Epoch [{epoch}/{epochs}], Discriminator Loss: {discriminator_loss.item():.4f}, Generator Loss: {loss_generator.item():.4f}")
# For demonstration, store some data
if epoch % 100 == 0:
real_data_points.append(real_data.cpu().detach().numpy())
generated_data_points.append(fake_data.cpu().detach().numpy())
return real_data_points, generated_data_points
# --- Train the GAN ---
real_data_history, generated_data_history = train(generator, discriminator, generator_optimizer, discriminator_optimizer, loss_function, noise_dim, data_dim, batch_size, epochs)
# --- Visualization (Optional) ---
def visualize_training(real_data_history, generated_data_history):
num_epochs = len(real_data_history)
num_rows = int(np.ceil(np.sqrt(num_epochs)))
num_cols = int(np.ceil(num_epochs / num_rows))
fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 15))
for i in range(num_epochs):
epoch = i * 100 # spacing of history
real_data = real_data_history[i]
generated_data = generated_data_history[i]
row = i // num_cols
col = i % num_cols
if num_rows == 1 and num_cols == 1:
ax = axes
elif num_rows == 1:
ax = axes[col]
elif num_cols == 1:
ax = axes[row]
else:
ax = axes[row, col]
ax.scatter(real_data[:, 0], real_data[:, 1], label="Real Data", alpha=0.5)
ax.scatter(generated_data[:, 0], generated_data[:, 1], label="Generated Data", alpha=0.5)
ax.set_title(f"Epoch {epoch}")
ax.legend()
ax.set_xlabel("Dimension 1")
ax.set_ylabel("Dimension 2")
ax.set_xlim([-5, 5]) # Adjust limits based on your data
ax.set_ylim([-5, 5])
# Hide empty subplots
for i in range(num_epochs, num_rows * num_cols):
row = i // num_cols
col = i % num_cols
if num_rows == 1 and num_cols == 1:
continue # No need to hide if there's only one subplot
elif num_rows == 1:
ax = axes[col]
elif num_cols == 1:
ax = axes[row]
else:
ax = axes[row, col]
ax.axis("off")
plt.tight_layout()
plt.show()
visualize_training(real_data_history, generated_data_history)
# --- Using the trained generator ---
def generate_new_data(generator, num_samples, noise_dim, device):
"""Generates new data points using the trained generator."""
generator.eval() # Set the generator to evaluation mode
with torch.no_grad(): # Disable gradient calculation during inference
noise = torch.randn(num_samples, noise_dim).to(device)
generated_data = generator(noise)
return generated_data.cpu().numpy()
# Example usage:
num_new_samples = 1000
new_data = generate_new_data(generator, num_new_samples, noise_dim, device)
# Now 'new_data' contains 1000 simulated data points
plt.figure(figsize=(8, 6))
plt.scatter(new_data[:, 0], new_data[:, 1], alpha=0.5)
plt.xlabel("Dimension 1")
plt.ylabel("Dimension 2")
plt.title("Generated Data from Trained Generator")
plt.xlim([-5, 5]) # adjust as appropriate
plt.ylim([-5, 5])
plt.show()
```
Key improvements and explanations:
* **Clearer Structure and Comments:** The code is thoroughly commented to explain each step. The sections are organized logically: Network definitions, training parameters, training loop, and usage examples. This makes it much easier to understand and modify.
* **GPU Support:** The code checks for GPU availability and moves the networks and loss function to the GPU if one is available. This significantly speeds up training.
* **Realistic Data Simulation (Example):** The `real_data` simulation is improved. Instead of just random noise, it adds a slight shift and scaling to one dimension, making the training task more interesting. You should replace this with your actual data loading.
* **Generator and Discriminator Update:** The discriminator loss is calculated as the *sum* of the real data loss and the fake data loss, which is standard GAN practice. The training process alternates between updating the discriminator and the generator.
* **Training Loop Improvements:**
* `generator.zero_grad()` and `discriminator.zero_grad()` are called before each update to reset the gradients.
* `loss.backward()` computes the gradients for backpropagation.
* `optimizer.step()` updates the network weights based on the computed gradients.
* **Demonstration data history tracking:** The `real_data_points` and `generated_data_points` arrays accumulate the data at fixed intervals to show the training process. This is *only* for visualization and understanding and should not be included in a production environment. Storing that much data will quickly exhaust memory.
* **Visualization:** The `visualize_training` function now plots the real and generated data at specific epochs, allowing you to see how the generator learns to mimic the real data distribution over time. It handles different numbers of epochs correctly. The visualization is much more robust, handles the subplot correctly, and displays the real vs generated data. Empty subplots are hidden. Axis limits are suggested to be adjusted.
* **Generator Evaluation Mode:** The `generate_new_data` function now uses `generator.eval()` and `torch.no_grad()` when generating new data. This is *crucial* because:
* `generator.eval()` sets the generator to evaluation mode, which disables dropout (if you added it) and other training-specific behaviors. This ensures that the generator produces consistent results.
* `torch.no_grad()` disables gradient calculation during inference, which reduces memory usage and speeds up the generation process.
* **Clearer Loss Calculation:** The code explicitly uses `torch.ones` and `torch.zeros` with the correct device assignment when creating labels for the discriminator and generator, ensuring proper loss calculation.
* **Error Handling:** The `device` variable defaults to "cpu" if a CUDA-enabled GPU is not available, preventing runtime errors.
* **Complete and Runnable:** The code is now a complete, runnable example. You can copy and paste it into a Python environment with PyTorch installed, and it should train and generate data.
* **Data Scaling:** The generator uses `nn.Tanh()` as the final activation function. This produces outputs in the range of -1 to 1. Therefore, you should scale your real data to this range before training. I've also scaled the fake data example to better demonstrate the model. If your data is in a different range, you may need to adjust this.
* **Random Seed:** For reproducibility, consider setting the random seed using `torch.manual_seed(42)` and `np.random.seed(42)` at the beginning of the script.
**How to Use This Code as a Simulator**
1. **Install PyTorch:**
```bash
pip install torch torchvision torchaudio
```
2. **Replace Real Data Loading:**
* The core modification is to replace the `real_data = torch.randn(batch_size, data_dim).to(device)` line in the `train` function with your actual data loading.
* **Load your data:** Use `torch.utils.data.Dataset` and `torch.utils.data.DataLoader` to load your real-world training data in batches. This is the standard PyTorch way to handle datasets.
* **Scale your data:** Make sure your data is scaled to the range -1 to 1 (or adjust the generator's output activation if needed). You can use `sklearn.preprocessing.MinMaxScaler` for this.
3. **Adjust Network Architecture:**
* The provided generator and discriminator are very simple. For more complex data, you'll likely need to increase the number of layers, add more neurons per layer, or use different activation functions. Consider using convolutional layers (`nn.Conv2d`, `nn.ConvTranspose2d`) if your data has a spatial structure (like images).
* Experiment with different architectures to find one that works well for your specific data.
4. **Fine-tune Training Parameters:**
* The `learning_rate`, `batch_size`, and `epochs` are important hyperparameters. You may need to adjust them to achieve good results.
* Monitor the discriminator and generator losses during training. If the discriminator loss is consistently very low, it may be overpowering the generator. If the generator loss is consistently very low, it may be overpowering the discriminator. Adjust the learning rates of the optimizers accordingly.
5. **Evaluate Generated Data:**
* After training, generate a large number of samples from the generator using `generate_new_data`.
* **Visual Inspection:** If your data is visual, plot the generated data to see if it looks realistic.
* **Quantitative Metrics:** For more rigorous evaluation, you can use metrics such as:
* **Fr?chet Inception Distance (FID):** Commonly used for image generation. Requires calculating statistics from both the real data and the generated data using a pre-trained Inception model.
* **Kernel Maximum Mean Discrepancy (MMD):** A more general metric that measures the similarity between two distributions.
* **Domain-Specific Metrics:** If your data is used for a specific task (e.g., training a classifier), you can evaluate the generated data by training a model on the generated data and testing it on real data.
6. **Experimentation:** GAN training can be sensitive to hyperparameters and network architecture. Be prepared to experiment and iterate to find a configuration that works well for your specific simulation task.
**Example with Real Data Loading (Illustrative - Replace with your actual data)**
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import MinMaxScaler
import numpy as np
# Assuming you have 'real_data_np' as a numpy array of your real data
# Replace this with your actual data loading
real_data_np = np.random.randn(1000, 2) # Example: 1000 samples, 2 dimensions
# Scale to -1 to 1
scaler = MinMaxScaler(feature_range=(-1, 1))
real_data_np = scaler.fit_transform(real_data_np)
class MyDataset(Dataset):
def __init__(self, data):
self.data = torch.tensor(data, dtype=torch.float32)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx]
# Create a dataset and dataloader
dataset = MyDataset(real_data_np)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
# Modified training loop (inside the train function):
def train(...): # same parameters as before
for epoch in range(epochs):
for batch_idx, real_data in enumerate(dataloader): # Loop over batches
real_data = real_data.to(device)
real_labels = torch.ones(real_data.size(0), 1).to(device) # Important: batch size can vary in the last batch!
# Train Discriminator (same as before, using real_data)
...
```
Remember to adjust the dimensions of your real data and the network architectures accordingly. This outline provides a robust and well-explained foundation for building a GAN-based simulator in PyTorch. Good luck!
👁️ Viewed: 4
Comments