60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import torch
|
|
from torch import nn, optim
|
|
from torch.autograd.variable import Variable
|
|
from torchvision import transforms, datasets
|
|
from utils import real_data_target, fake_data_target
|
|
|
|
class DiscriminatorNet(torch.nn.Module):
|
|
"""
|
|
A three hidden-layer discriminative neural network
|
|
"""
|
|
def __init__(self, in_features, leakyRelu=0.2, dropout=0.3, hidden_layers=[1024, 512, 256]):
|
|
super(DiscriminatorNet, self).__init__()
|
|
|
|
out_features = 1
|
|
self.layers = hidden_layers.copy()
|
|
self.layers.insert(0, in_features)
|
|
|
|
for count in range(0, len(self.layers)-1):
|
|
self.add_module("hidden_" + str(count),
|
|
nn.Sequential(
|
|
nn.Linear(self.layers[count], self.layers[count+1]),
|
|
nn.LeakyReLU(leakyRelu),
|
|
nn.Dropout(dropout)
|
|
)
|
|
)
|
|
|
|
self.add_module("out",
|
|
nn.Sequential(
|
|
nn.Linear(self.layers[-1], out_features),
|
|
torch.nn.Sigmoid()
|
|
)
|
|
)
|
|
|
|
def forward(self, x):
|
|
for name, module in self.named_children():
|
|
x = module(x)
|
|
return x
|
|
|
|
# train_discriminator(d_optimizer, discriminator, loss, real_data, fake_data)
|
|
def train_discriminator(optimizer, discriminator, loss, real_data, fake_data):
|
|
# Reset gradients
|
|
optimizer.zero_grad()
|
|
|
|
# 1.1 Train on Real Data
|
|
prediction_real = discriminator(real_data)
|
|
# Calculate error and backpropagate
|
|
error_real = loss(prediction_real, real_data_target(real_data.size(0)))
|
|
error_real.backward()
|
|
|
|
# 1.2 Train on Fake Data
|
|
prediction_fake = discriminator(fake_data)
|
|
# Calculate error and backpropagate
|
|
error_fake = loss(prediction_fake, fake_data_target(real_data.size(0)))
|
|
error_fake.backward()
|
|
|
|
# 1.3 Update weights with gradients
|
|
optimizer.step()
|
|
|
|
# Return error
|
|
return error_real + error_fake, prediction_real, prediction_fake |