CNN.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. from torch import device, cuda
  2. import torch
  3. from torch import add
  4. import torch.nn as nn
  5. import utils.CNN_Layers as CustomLayers
  6. import torch.nn.functional as F
  7. import torch.optim as optim
  8. import pandas as pd
  9. import matplotlib.pyplot as plt
  10. import time
  11. import numpy as np
  12. from sklearn.metrics import roc_curve, auc, confusion_matrix, classification_report
  13. import seaborn as sns
  14. class CNN_Net(nn.Module):
  15. def __init__(self, prps, final_layer_size=5):
  16. super(CNN_Net, self).__init__()
  17. self.final_layer_size = final_layer_size
  18. self.device = device('cuda:0' if cuda.is_available() else 'cpu')
  19. print("CNN Initialized. Using: " + str(self.device))
  20. # LAYERS
  21. print(f"CNN Model Initialization")
  22. self.conv1 = CustomLayers.Conv_elu_maxpool_drop(1, 192, (11, 13, 11), stride=(4,4,4), pool=True, prps=prps)
  23. self.conv2 = CustomLayers.Conv_elu_maxpool_drop(192, 384, (5, 6, 5), stride=(1,1,1), pool=True, prps=prps)
  24. self.conv3_mid_flow = CustomLayers.Mid_flow(384, 384, prps=prps)
  25. self.conv4_sepConv = CustomLayers.Conv_elu_maxpool_drop(384, 96,(3, 4, 3), stride=(1,1,1), pool=True, prps=prps,
  26. sep_conv=True)
  27. self.conv5_sepConv = CustomLayers.Conv_elu_maxpool_drop(96, 48, (3, 4, 3), stride=(1, 1, 1), pool=True,
  28. prps=prps, sep_conv=True)
  29. self.fc1 = CustomLayers.Fc_elu_drop(113568, 20, prps=prps, softmax=False) # TODO, concatenate clinical data after this
  30. self.fc2 = CustomLayers.Fc_elu_drop(20, final_layer_size, prps=prps, softmax=True) # For now this works as output layer, though may be incorrect
  31. # FORWARDS
  32. def forward(self, x):
  33. x = self.conv1(x)
  34. x = self.conv2(x)
  35. x = self.conv3_mid_flow(x)
  36. x = self.conv4_sepConv(x)
  37. x = self.conv5_sepConv(x)
  38. # FLATTEN x
  39. flatten_size = x.size(1) * x.size(2) * x.size(3) * x.size(4)
  40. x = x.view(-1, flatten_size)
  41. x = self.fc1(x)
  42. x = self.fc2(x)
  43. return x