Fix loss functions

src/model/GANModel.py

@@ -45,12 +45,12 @@ def eval_mode(self):
  def _generator_criterion(d_g_z: FloatTensor) -> FloatTensor:
  """
  Loss function for Generator G.
- Calculates 1/m · ∑ log(1 - D(G(z))
+ Calculates 1/m · ∑ log(D(G(z))
  where *z* is the uniform random vector (noise) ∈ [0, 1]^T
  :param d_g_z: Tensor corresponding to the discriminator prediction D(G(z))
  :return: Loss of G
  """
- return torch.mean(torch.log(1 - d_g_z))
+ return torch.mean(torch.log(d_g_z))
 
  @staticmethod
  def _discriminator_criterion(d_x: FloatTensor, d_g_z: FloatTensor) -> FloatTensor:

src/model/train.py

@@ -61,12 +61,12 @@ def train_generator(model: GANModel, data):
  prediction = model.discriminator(fake_data)
 
  # Calculate gradients w.r.t parameters and backpropagate
- loss = model.g_criterion(d_g_z=prediction)
- (-loss).backward()
+ loss = -model.g_criterion(d_g_z=prediction)
+ loss.backward()
 
  # Update parameters
  model.g_optimizer.step()
- return loss
+ return -loss
 
 
 def train_discriminator(model: GANModel, data):
@@ -88,13 +88,13 @@ def train_discriminator(model: GANModel, data):
  fake_predictions = model.discriminator(fake_data)
 
  # Calculate loss and optimize
- loss = model.d_criterion(d_x=real_predictions, d_g_z=fake_predictions)
- (-loss).backward()
+ loss = -model.d_criterion(d_x=real_predictions, d_g_z=fake_predictions)
+ loss.backward()
 
  # Update parameters
  model.d_optimizer.step()
 
- return loss
+ return -loss
 
 
 def train_epoch(model: GANModel, loader: DataLoader) -> Tuple[float, float]: