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] * 4 + [outputs.shape[1]] dropout_rate = 0.00 model = NeuralNetwork(layer_sizes, dropout_rate=dropout_rate, 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': dropout_rate } 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').to(torch.device('cpu')) # data = Dataset() # data = Dataset() # print(np.unique(data.df['Fiber_Volume_Fractions'].to_numpy())[:10]) # test_input = torch.tensor([[2, 0.6, 450.0, 100.0, 500.0]], dtype=torch.float32) # test_output = model.predict((test_input - torch.tensor(data.input_mean)) / torch.tensor(data.input_std)) # test_output = test_output * torch.tensor(data.output_std) + torch.tensor(data.output_mean) # print(f"Test Prediction for fixed input {test_input.numpy()}: {test_output.numpy()}")