Spaces:
Sleeping
Sleeping
| 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()}") | |