New training

.gitattributes

@@ -1 +1,2 @@
 checkpoints/latest_checkpoint.pth.tar filter=lfs diff=lfs merge=lfs -text
+checkpoints/latest_checkpoint1.pth.tar filter=lfs diff=lfs merge=lfs -text

checkpoints/latest_checkpoint.pth.tar

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c8de65df3e4879e931cdf7f3de2fdc3d05298c0e955b39b2281627f36e27fcff
-size 686252474
+oid sha256:9b32b1f4363aad01e662d468989e7e0b8f41afec20ffcbf1e87b6a6147454cbd
+size 686253114

checkpoints/latest_checkpoint1.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8de65df3e4879e931cdf7f3de2fdc3d05298c0e955b39b2281627f36e27fcff
+size 686252474

colorizer_pipeline.py

@@ -230,13 +230,22 @@ def visualize_results(epoch, generator, train_loader, device):
     generator.train()
 
 def save_checkpoint(state, filename="checkpoint.pth.tar"):
-    torch.save(state, filename)
+    # Only save the necessary state
+    save_state = {
+        'epoch': state['epoch'],
+        'generator_state_dict': state['generator_state_dict'],
+        'discriminator_state_dict': state['discriminator_state_dict'],
+        'optimizerG_state_dict': state['optimizerG_state_dict'],
+        'optimizerD_state_dict': state['optimizerD_state_dict'],
+    }
+    torch.save(save_state, filename)
     mlflow.log_artifact(filename)
 
-def load_checkpoint(filename, generator, discriminator, optimizerG, optimizerD):
+def load_checkpoint(filename, generator, discriminator, optimizerG, optimizerD, device):
     if os.path.isfile(filename):
         print(f"Loading checkpoint '{filename}'")
-        checkpoint = torch.load(filename)
+        # Use weights_only=True for safer loading
+        checkpoint = torch.load(filename, map_location=device, weights_only=True)
         start_epoch = checkpoint['epoch'] + 1
         generator.load_state_dict(checkpoint['generator_state_dict'])
         discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
@@ -248,6 +257,11 @@ def load_checkpoint(filename, generator, discriminator, optimizerG, optimizerD):
         print(f"No checkpoint found at '{filename}'")
         return 0
 
+# Global variables
+checkpoint_dir = "checkpoints"
+os.makedirs(checkpoint_dir, exist_ok=True)
+os.makedirs("results", exist_ok=True)
+
 # Training function
 def train(generator, discriminator, train_loader, num_epochs, device, lr=0.0002, beta1=0.5):
     criterion = nn.BCEWithLogitsLoss()
@@ -256,12 +270,8 @@ def train(generator, discriminator, train_loader, num_epochs, device, lr=0.0002,
     optimizerG = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
     optimizerD = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))
 
-    checkpoint_dir = "checkpoints"
-    os.makedirs(checkpoint_dir, exist_ok=True)
-    os.makedirs("results", exist_ok=True)
-
     checkpoint_path = os.path.join(checkpoint_dir, "latest_checkpoint.pth.tar")
-    start_epoch = load_checkpoint(checkpoint_path, generator, discriminator, optimizerG, optimizerD)
+    start_epoch = load_checkpoint(checkpoint_path, generator, discriminator, optimizerG, optimizerD, device)
 
     experiment_id = setup_mlflow()
     with mlflow.start_run(experiment_id=experiment_id, run_name="training_run") as run:
@@ -270,7 +280,7 @@ def train(generator, discriminator, train_loader, num_epochs, device, lr=0.0002,
                 generator.train()
                 discriminator.train()
 
-                num_iterations = 2
+                num_iterations = 2000
                 pbar = tqdm(enumerate(islice(train_loader, num_iterations)), total=num_iterations, desc=f"Epoch {epoch+1}/{num_epochs}")
                 
                 for i, (real_L, real_AB) in pbar:

convert_checkpoint.py

@@ -0,0 +1,31 @@
+import torch
+import os
+
+def load_and_save_checkpoint(input_filename, output_filename, device):
+    if os.path.isfile(input_filename):
+        print(f"Loading checkpoint '{input_filename}'")
+        checkpoint = torch.load(input_filename, map_location=device)
+        
+        # Extract only the necessary state
+        save_state = {
+            'epoch': checkpoint['epoch'],
+            'generator_state_dict': checkpoint['generator_state_dict'],
+            'discriminator_state_dict': checkpoint['discriminator_state_dict'],
+            'optimizerG_state_dict': checkpoint['optimizerG_state_dict'],
+            'optimizerD_state_dict': checkpoint['optimizerD_state_dict'],
+        }
+        
+        # Save the checkpoint
+        torch.save(save_state, output_filename)
+        print(f"Saved checkpoint to '{output_filename}'")
+    else:
+        print(f"No checkpoint found at '{input_filename}'")
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    input_checkpoint = "checkpoints/latest_checkpoint.pth.tar"
+    output_checkpoint = "checkpoints/converted_checkpoint.pth.tar"
+
+    load_and_save_checkpoint(input_checkpoint, output_checkpoint, device)