main.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. print("--- INITIALIZING LIBRARIES ---")
  2. import torch
  3. from torch import nn
  4. # GENERAL PURPOSE
  5. import numpy as np
  6. from datetime import datetime
  7. import pandas as pd
  8. import os
  9. import tomli as tl
  10. # FOR TRAINING
  11. import torch.optim as optim
  12. import utils.models as models
  13. from utils.training import train_model, test_model, initalize_dataloaders, plot_results
  14. print("--- LIBRARIES INITIALIZED ---")
  15. #GET CONFIG SETTINGS
  16. if os.getenv('ADL_CONFIG_PATH') is None:
  17. with open ('config.toml', 'rb') as f:
  18. config = tl.load(f)
  19. else:
  20. with open(os.getenv('ADL_CONFIG_PATH'), 'rb') as f:
  21. config = tl.load(f)
  22. cuda_device = torch.device(config['cuda']['device'])
  23. torch.set_default_device(cuda_device)
  24. model_name = config['model']['name']
  25. print("--- RUNNING ---")
  26. print("Pytorch Version: " + torch. __version__)
  27. # data & training properties:
  28. val_split = config['dataset']['validation_split'] # % of val and test, rest will be train
  29. runs = config['training']['runs']
  30. epochs = config['training']['epochs']
  31. seeds = [np.random.randint(0, 1000) for _ in range(runs)]
  32. #paths
  33. mri_path = config['paths']['mri_data']
  34. xls_path = config['paths']['xls_data']
  35. saved_model_path = config['paths']['model_output']
  36. plot_path = config['paths']['plot_output']
  37. training_record_path = config['paths']['training_record_output']
  38. #model
  39. model_CNN = models.CNN_Net(1, 2, config['model']['droprate']).to(cuda_device)
  40. criterion = nn.BCELoss()
  41. optimizer = optim.Adam(model_CNN.parameters(), lr=config['training']['learning_rate'])
  42. for seed in seeds:
  43. #get time stamp for model
  44. time_stamp = datetime.now().strftime('%Y%m%d+%H%M%S')
  45. #initialize dataloaders, train model, and test model
  46. train_loader, val_loader, test_loader, _ = initalize_dataloaders(mri_path, xls_path, val_split, seed, cuda_device=cuda_device, batch_size=config['training']['batch_size'])
  47. print("--- TRAINING MODEL ---")
  48. print("Seed: ", seed)
  49. train_results = train_model(model_CNN, seed, time_stamp, epochs, train_loader, val_loader, saved_model_path, model_name, optimizer, criterion, cuda_device=cuda_device)
  50. print("--- TESTING MODEL ---")
  51. predicted, actual, correct, incorrect, _, _ = test_model(model_CNN, test_loader, cuda_device=cuda_device)
  52. print("Accuracy: " + str(correct / (correct + incorrect)))
  53. #Plot results and confusion matrix
  54. plot_results(train_results["train_acc"], train_results["train_loss"], train_results["val_acc"], train_results["val_loss"], model_name, time_stamp, plot_path)
  55. #Save training results
  56. if not os.path.exists(training_record_path):
  57. os.makedirs(training_record_path)
  58. train_results.to_csv(training_record_path + model_name + "_t-" + time_stamp + "_s-" + str(seed) + "_e-" + str(epochs) + ".csv")
  59. print("--- END ---")