Merge pull request #57 from renan-siqueira/develop

Merge develop into main

Author: renan-siqueira

src/app/training.py

@@ -85,7 +85,8 @@ def train_model(
     lambda_gp = 10  # Default value to WGAN-GP
 
     fixed_noise = Variable(torch.randn(sample_size, z_dim, 1, 1)).to(device)
-    
+    fid_score = 'Unavailable'
+
     for epoch in range(last_epoch, num_epochs + 1):
         start_time = time.time()
 
@@ -118,8 +119,9 @@ def train_model(
             fake_images = generator(z)
             real_images = images
 
-            # FID Test
-            fid_score = calculate_fid(real_images, fake_images, inception_model)
+            if inception_model:
+                # FID Test
+                fid_score = calculate_fid(real_images, fake_images, inception_model)
 
             outputs = discriminator(fake_images).squeeze()
             g_loss = -torch.mean(outputs)

src/modules/run_training.py

@@ -23,9 +23,11 @@ def main(params, path_data, path_dataset, path_train_params):
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-    # Frechet Inception Distance (FID)
-    inception_model = models.inception_v3(weights='Inception_V3_Weights.DEFAULT', transform_input=False, init_weights=False).to(device)
-    inception_model = inception_model.eval()
+    inception_model = None
+    if params['image_size'] == 128:
+        # Frechet Inception Distance (FID)
+        inception_model = models.inception_v3(weights='Inception_V3_Weights.DEFAULT', transform_input=False, init_weights=False).to(device)
+        inception_model = inception_model.eval()
 
     generator = Generator(params["z_dim"], params["channels_img"], params["features_g"], img_size=params['image_size']).to(device)
     generator.apply(weights_init)