Chapter 12 fix: Added support for running scripts from different paths and referencing model files correctly.

Author: Rishal Hurbans

ch12-generative_image_models/numpy_u_net/toy_numpy_u_net_generate.py

@@ -1,7 +1,9 @@
-import numpy as np
-import matplotlib.pyplot as plt
 import math
 import pickle
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
 
 # Step 0: Redefine the NumPy Layers and UNet Class
 # We must have the exact same class definitions as the training script
@@ -160,13 +162,14 @@ def set_params(self, params):
 # Step 1: Load Model and Define Helpers
 print("\nStep 1: Loading Model and Defining Helpers")
 model = NumPyUNet()
+MODEL_PATH = Path(__file__).resolve().parent / "trained_numpy_unet_model.pkl"
 try:
-    with open('trained_numpy_unet_model.pkl', 'rb') as f:
+    with MODEL_PATH.open('rb') as f:
         params = pickle.load(f)
     model.set_params(params)
     print("Successfully loaded trained NumPy U-Net model weights.")
 except FileNotFoundError:
-    print("Error: 'trained_numpy_unet_model.pkl' not found.")
+    print(f"Error: '{MODEL_PATH}' not found.")
     print("Please run the NumPy training script first.")
     exit()
 

ch12-generative_image_models/numpy_u_net/toy_numpy_u_net_train.py

@@ -1,7 +1,9 @@
-import numpy as np
-import matplotlib.pyplot as plt
 import math
 import pickle
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
 
 # Seed randomness for reproducibility
 np.random.seed(42)
@@ -288,6 +290,9 @@ def set_params(self, params):
             if f'{name}_gamma' in params: layer.gamma = params[f'{name}_gamma']
             if f'{name}_beta' in params: layer.beta = params[f'{name}_beta']
 
+# Where to persist the trained model so the script works regardless of CWD
+MODEL_PATH = Path(__file__).resolve().parent / "trained_numpy_unet_model.pkl"
+
 # Step 1: Define The Training Data
 print("\nStep 1: Defining The Training Data")
 training_data = [
@@ -577,6 +582,6 @@ def create_timestep_embedding(t, embedding_dim):
 plt.show()
 
 print("\n--- Saving Model Weights ---")
-with open('trained_numpy_unet_model.pkl', 'wb') as f:
+with MODEL_PATH.open('wb') as f:
     pickle.dump(model.get_params(), f)
-print("Model weights saved to 'trained_numpy_unet_model.pkl'")
+print(f"Model weights saved to '{MODEL_PATH}'")

ch12-generative_image_models/pytorch-u_net/toy_pytorch_u_net_generate.py

@@ -1,8 +1,10 @@
+import math
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
-import matplotlib.pyplot as plt
-import math
 
 # Step 1: Setup and Model Definition
 print("Step 1: Setting up environment and model definition")
@@ -54,11 +56,12 @@ def forward(self, x, t_emb, c_emb):
 print("\nStep 2: Loading Model and Defining Helpers")
 # Initialize the model and load the saved weights
 model = UNet().to(device)
+MODEL_PATH = Path(__file__).resolve().parent / "trained_pytorch_unet_model.pth"
 try:
-    model.load_state_dict(torch.load('/trained_pytorch_unet_model.pth', map_location=device))
+    model.load_state_dict(torch.load(str(MODEL_PATH), map_location=device))
     print("Successfully loaded trained U-Net model weights.")
 except FileNotFoundError:
-    print("Error: 'trained_unet_model.pth' not found.")
+    print(f"Error: '{MODEL_PATH}' not found.")
     print("Please run the training script first to create the model file.")
     exit()
 model.eval() # Set the model to evaluation mode

ch12-generative_image_models/pytorch-u_net/toy_pytorch_u_net_train.py

@@ -1,13 +1,18 @@
+import math
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
-import matplotlib.pyplot as plt
-import math
 
 # Use a GPU if available, otherwise use the CPU
 device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
 print(f"Using device: {device}")
 
+# Where to persist the trained model so the script works regardless of CWD
+MODEL_PATH = Path(__file__).resolve().parent / "trained_pytorch_unet_model.pth"
+
 # Step 1: Define the Training Data
 print("\nStep 1: Defining The Training Data")
 training_data = [
@@ -225,5 +230,5 @@ def forward(self, x, t_emb, c_emb):
 
 # Save the Trained Model Weights
 print("\nSaving Model Weights")
-torch.save(model.state_dict(), '/trained_pytorch_unet_model.pth')
-print("Model weights saved to 'trained_unet_model.pth'")
+torch.save(model.state_dict(), str(MODEL_PATH))
+print(f"Model weights saved to '{MODEL_PATH}'")