model update

main_thermo.py

@@ -0,0 +1,166 @@
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from Dataset import Dataset
+from model import NeuralNetwork
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Set global plotting parameters
+plt.rcParams.update({'font.size': 14,
+                     'figure.figsize': (10, 8),
+                     'lines.linewidth':  2,
+                     'lines.markersize': 6,
+                     'axes.grid': True,
+                     'axes.labelsize': 16,
+                     'legend.fontsize': 14,
+                     'xtick.labelsize': 14,
+                     'ytick.labelsize': 14,
+                     'figure.autolayout': True
+                     })
+
+def set_seed(seed=42):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)    
+
+def train_neural_network(model, inputs, outputs, optimizer, epochs=1000, lr_scheduler=None):
+    model.train()
+    for epoch in range(epochs):
+        optimizer.zero_grad()
+        predictions = model(inputs)
+        loss = torch.mean(torch.square(predictions - outputs))
+        loss.backward()
+        optimizer.step()
+
+        if lr_scheduler:
+            lr_scheduler.step()
+
+        if epoch % 100 == 0:
+            print(f'Epoch {epoch}, Loss: {loss.item()}, Learning Rate: {optimizer.param_groups[0]["lr"]}')
+
+def main():
+    set_seed(5324)
+    dataset = Dataset()
+    inputs = dataset.get_input(normalize=True)
+    outputs = dataset.get_output(normalize=True)
+
+    idx_train = np.random.choice(len(inputs), size=int(0.98 * len(inputs)), replace=False)
+    idx_test = np.setdiff1d(np.arange(len(inputs)), idx_train)
+
+    inputs_train = torch.tensor(inputs[idx_train], dtype=torch.float32).to(DEVICE)
+    outputs_train = torch.tensor(outputs[idx_train], dtype=torch.float32).to(DEVICE)
+
+    inputs_test = torch.tensor(inputs[idx_test], dtype=torch.float32).to(DEVICE)
+    outputs_test = torch.tensor(outputs[idx_test], dtype=torch.float32).to(DEVICE)
+    
+    layer_sizes = [inputs.shape[1]] + [64] * 3 + [outputs.shape[1]]
+    model = NeuralNetwork(layer_sizes, dropout_rate=0.0, activation=torch.nn.ReLU).to(DEVICE)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.9)
+
+    # Create a proper dataset that keeps input-output pairs together
+    train_dataset = torch.utils.data.TensorDataset(inputs_train, outputs_train)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True)
+
+    # Train the model
+    epochs = 20000
+    for epoch in range(epochs):
+        model.train()
+        for inputs_batch, outputs_batch in train_loader:
+            inputs_batch = inputs_batch.to(DEVICE)
+            outputs_batch = outputs_batch.to(DEVICE)
+            optimizer.zero_grad()
+            predictions = model(inputs_batch)
+            loss = torch.mean(torch.square(predictions - outputs_batch))
+            loss.backward()
+            optimizer.step()
+
+        if lr_scheduler:
+            lr_scheduler.step()
+
+        if epoch % 500 == 0:
+            train_pred = model(inputs_train)
+            train_loss = torch.mean(torch.square(train_pred - outputs_train))
+            test_pred = model(inputs_test)
+            test_loss = torch.mean(torch.square(test_pred - outputs_test))
+            print(f'Epoch {epoch}, Train Loss: {train_loss.item():.6f}, Test Loss: {test_loss.item():.6f}')
+            # print(f'Learning Rate: {optimizer.param_groups[0]["lr"]}')
+
+
+    predictions = model.predict(inputs_test)
+    test_loss = torch.mean(torch.square(predictions - outputs_test))
+    print(f'Test Loss: {test_loss.item()}. Samples: {idx_test}')
+
+    x = np.arange(0, len(idx_test))
+
+    outputs_test = dataset.denormalize_output(outputs_test.cpu().numpy())
+    predictions = dataset.denormalize_output(predictions.cpu().numpy())
+    # for sample in outputs_test:
+    #     print(f'Test samples: {sample}')
+    plt.figure(figsize=(10, 6))
+    plt.plot(x, outputs_test[:, 0], color='b', linestyle='--', label='True A1')
+    plt.plot(x, predictions[:, 0], color='b', linestyle='-', label='Predicted A1')
+    plt.plot(x, outputs_test[:, 1], color='r', linestyle='--', label='True B1')
+    plt.plot(x, predictions[:, 1], color='r', linestyle='-', label='Predicted B1')
+    plt.plot(x, outputs_test[:, 2], color='g', linestyle='--', label='True C1')
+    plt.plot(x, predictions[:, 2], color='g', linestyle='-', label='Predicted C1')
+    plt.gca().xaxis.set_major_locator(plt.MaxNLocator(integer=True))    
+    plt.xlabel('Sample Index')
+    plt.xticks(ticks=range(len(idx_test)),labels=idx_test + 1)
+    plt.ylabel('Springback Angle (Degrees)')
+    plt.title('Springback Angle Prediction')
+    plt.legend(loc='upper right')
+    plt.savefig('springback_angle_prediction.png')
+
+
+    plt.figure(figsize=(10, 6))
+    plt.plot(x, outputs_test[:, 3], color='m', linestyle='--', label='True Stress(Max)')
+    plt.plot(x, predictions[:, 3], color='m', linestyle='-', label='Predicted Stress(Max)')
+    plt.xlabel('Sample Index')
+    plt.xticks(ticks=range(len(idx_test)),labels=idx_test + 1)
+    plt.ylabel('Stress (MPa)')
+    plt.legend(loc='upper left')
+    plt.savefig('stress_max_prediction.png')
+
+
+
+    # MSE
+    mse = np.mean((predictions - outputs_test) ** 2, axis=0)
+    print(f'Mean Squared Error for A1: {mse[0]:.6f}, B1: {mse[1]:.6f}, C1: {mse[2]:.6f}, Stress(Max): {mse[3]:.6f}')
+
+    # R 2 score
+    ss_ress = np.sum((outputs_test - predictions) ** 2, axis=0)
+    ss_tots = np.sum((outputs_test - np.mean(outputs_test, axis=0)) ** 2, axis=0)
+    r2_scores = 1 - ss_ress / ss_tots
+    print(f'R² Score for A1: {r2_scores[0]:.6f}, B1: {r2_scores[1]:.6f}, C1: {r2_scores[2]:.6f}, Stress(Max): {r2_scores[3]:.6f}')
+
+    # Error
+
+    # Save the model
+    model_save_path = './model_checkpoint.pth'
+    model_config = {'layer_sizes': layer_sizes,
+                    'dropout_rate': 0.05
+                    }
+    checkpoint = {
+        'model_state_dict': model.state_dict(),
+        'model_config': model_config
+    }
+    torch.save(checkpoint, model_save_path)
+    # Load the model
+    # model = NeuralNetwork(layer_sizes)
+    # model.load_state_dict(torch.load(model_save_path))
+
+def load_model(model_path):
+    checkpoint = torch.load(model_path)
+    model_config = checkpoint['model_config']
+    model = NeuralNetwork(model_config['layer_sizes'], dropout_rate=model_config['dropout_rate'], activation=torch.nn.ReLU).to(DEVICE)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    print(f"Model loaded from {model_path}")
+    return model
+
+
+if __name__ == "__main__":
+    main()
+    # model = load_model('./model_checkpoint.pth')
+

model_checkpoint.pth

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

model_inverse.py

@@ -90,7 +90,7 @@ def load_model(model_path):
     return model
 
 def inverse_design(ply_number, y_target, n_restarts=20, epochs=1000, use_lbfgs=False):
-    model = load_model('./model_inverse_ckpt.pth')
+    model = load_model('./model_checkpoint.pth')
 
     data = Dataset()
     y_target_norm = data.normalize_output(y_target) # (A1, B1, C1, Stress)
@@ -101,7 +101,7 @@ def inverse_design(ply_number, y_target, n_restarts=20, epochs=1000, use_lbfgs=F
     output_std = torch.tensor(data.output_std)
     
     weights = torch.tensor([1.0, 1.0, 1.0, 0.0], dtype=torch.float32)
-    bounds = torch.tensor([[350., 450.], [100., 500.], [450., 550.]], dtype=torch.float32) # Initial_Temp, Punch_Velocity, Cooling_Time
+    bounds = torch.tensor([[100., 600.], [100., 600.], [100., 600.]], dtype=torch.float32) # Initial_Temp, Punch_Velocity, Cooling_Time
     best = {"loss": float('inf'), "input": None, "output": None}
 
     for restart in range(n_restarts):
@@ -235,15 +235,15 @@ def inverse_model():
     # Error
 
     # Save the model
-    # model_save_path = './model_inverse_ckpt.pth'
-    # model_config = {'layer_sizes': layer_sizes,
-    #                 'dropout_rate': 0.05
-    #                 }
-    # checkpoint = {
-    #     'model_state_dict': model.state_dict(),
-    #     'model_config': model_config
-    # }
-    # torch.save(checkpoint, model_save_path)
+    model_save_path = './model_inverse_ckpt.pth'
+    model_config = {'layer_sizes': layer_sizes,
+                    'dropout_rate': 0.05
+                    }
+    checkpoint = {
+        'model_state_dict': model.state_dict(),
+        'model_config': model_config
+    }
+    torch.save(checkpoint, model_save_path)
 
 
 if __name__ == "__main__":