File size: 2,309 Bytes
de86909 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch
from torch import nn
from torch.utils.data import DataLoader
# Hyperparameters
image_size = (224, 224, 3) # Adjust based on your data
# Define the Generator Network
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
# Define convolutional layers with appropriate filters and activations
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
# ... Add more convolutional layers as needed
self.conv_final = nn.Conv2d(128, 3, kernel_size=3, stride=1, padding=1, activation=nn.Tanh) # Tanh for shadow intensity
def forward(self, x):
# Define the forward pass through the convolutional layers
x = self.conv1(x)
# ... Forward pass through remaining convolutional layers
return self.conv_final(x)
# Define the Discriminator Network
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
# Define convolutional layers with appropriate filters and activations
self.conv1 = nn.Conv2d(6, 32, kernel_size=3, stride=1, padding=1)
# ... Add more convolutional layers as needed
self.linear = nn.Linear(128, 1) # Final layer with sigmoid activation
def forward(self, car, shadow):
# Concatenate car and shadow features
x = torch.cat([car, shadow], dim=1)
# Define the forward pass through the convolutional layers
x = self.conv1(x)
# ... Forward pass through remaining convolutional layers
return torch.sigmoid(self.linear(x))
# Create data loaders for training and validation data
# ... (Implement data loading logic using PyTorch's DataLoader)
# Create the models
generator = Generator()
discriminator = Discriminator()
# Define loss function and optimizer
criterion = nn.BCELoss()
g_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0002)
d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0002)
# Training loop
for epoch in range(epochs):
# Train the Discriminator
# ... (Implement discriminator training logic with loss calculation and updates)
# Train the Generator
# ... (Implement generator training logic with loss calculation and updates)
# Print training progress
# ... (Print loss values or other metrics)
# Save the trained generator
torch.save(generator.state_dict(), 'generator.pt')
|