Merge pull request ML4ITS#3 from ML4ITS/trainer_class

Trainer class

.gitignore

@@ -6,4 +6,5 @@ datasets/ecg_data.txt
 eda/.ipynb_checkpoints/
 ServerMachineDataset/processed/
 plots/
-output/
+*.pkl
+models/

modules.py

@@ -2,17 +2,21 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from datetime import datetime
 
 
 class ConvLayer(nn.Module):
-	"""
-	1-D Convolution layer to extract high-level features of each time-series input
+	""" 1-D Convolution layer to extract high-level features of each time-series input
+		:param num_nodes: Number of input features/nodes
+		:param window_size: length of the input sequence
+		:param: kernel_size: size of kernel to use in the convolution operation
+		:param device: which device to use (cpu or cuda)
 	"""
 
-	def __init__(self, num_nodes, window_size, kernel_size=7, device='cpu'):
+	def __init__(self, n_features, window_size, kernel_size=7, device='cpu'):
 		super(ConvLayer, self).__init__()
 		self.padding = nn.ConstantPad1d((kernel_size - 1) // 2, 0.0)
-		self.conv = nn.Conv1d(in_channels=num_nodes, out_channels=num_nodes, kernel_size=kernel_size)
+		self.conv = nn.Conv1d(in_channels=n_features, out_channels=n_features, kernel_size=kernel_size)
 		self.relu = nn.ReLU()
 
 	def forward(self, x):
@@ -23,10 +27,14 @@ def forward(self, x):
 
 
 class FeatureAttentionLayer(nn.Module):
-	"""
-	Single Graph Feature/Spatial Attention Layer
-	"""
+	""" Single Graph Feature/Spatial Attention Layer
+		:param num_nodes: Number of input features/nodes
+		:param window_size: length of the input sequence
+	 	:param dropout: percentage of nodes to dropout
+	 	:param alpha: negative slope used in the leaky rely activation function
+		:param device: which device to use (cpu or cuda)
 
+	"""
 	def __init__(self, num_nodes, window_size, dropout, alpha, device='cpu'):
 		super(FeatureAttentionLayer, self).__init__()
 		self.num_nodes = num_nodes
@@ -37,7 +45,7 @@ def __init__(self, num_nodes, window_size, dropout, alpha, device='cpu'):
 		self.w = nn.Parameter(torch.empty((2 * window_size, 1)))
 		nn.init.xavier_uniform_(self.w.data, gain=1.414)
 
-	# self.a = nn.Linear(2 * out_dim, 1, bias=False)
+		self.sigmoid = nn.Sigmoid()
 
 	def forward(self, x):
 		# x has shape (b, n, k) where n is window size and k is number of nodes
@@ -59,7 +67,7 @@ def forward(self, x):
 		attention = torch.dropout(attention, self.dropout, train=self.training)
 
 		# Computing new node features using the attention
-		h = torch.matmul(attention, v)
+		h = self.sigmoid(torch.matmul(attention, v))
 
 		return h
 
@@ -90,8 +98,13 @@ def _make_attention_input(self, v):
 
 
 class TemporalAttentionLayer(nn.Module):
-	"""
-	Single Graph Temporal Attention Layer
+	""" Single Graph Temporal Attention Layer
+		:param num_nodes: Number of input features/nodes
+		:param window_size: length of the input sequence
+	 	:param dropout: percentage of nodes to dropout
+	 	:param alpha: negative slope used in the leaky rely activation function
+		:param device: which device to use (cpu or cuda)
+
 	"""
 
 	def __init__(self, num_nodes, window_size, dropout, alpha, device='cpu'):
@@ -104,6 +117,8 @@ def __init__(self, num_nodes, window_size, dropout, alpha, device='cpu'):
 		self.w = nn.Parameter(torch.empty((2 * num_nodes, 1)))
 		nn.init.xavier_uniform_(self.w.data, gain=1.414)
 
+		self.sigmoid = nn.Sigmoid()
+
 	# self.a = nn.Linear(2 * out_dim, 1, bias=False)
 
 	def forward(self, x):
@@ -121,12 +136,10 @@ def forward(self, x):
 		attention = torch.dropout(attention, self.dropout, train=self.training)
 
 		# Computing new node features using the attention
-		h = torch.matmul(attention, x)
+		h = self.sigmoid(torch.matmul(attention, x))
 
 		return h
 
-
-
 	def _make_attention_input(self, v):
 		""" Preparing the temporal attention mechanism.
 			Creating matrix with all possible combinations of concatenations of node values:
@@ -152,8 +165,12 @@ def _make_attention_input(self, v):
 
 
 class GRU(nn.Module):
-	"""
-	Gated Recurrent Unit (GRU)
+	""" Gated Recurrent Unit (GRU)
+		:param in_dim: number of input features
+		:param hid_dim: hidden size of the GRU
+		:param n_layers: number of layers in GRU
+		:param dropout: dropout rate
+		:param device: which device to use (cpu or cuda)
 	"""
 
 	def __init__(self, in_dim, hid_dim, n_layers, dropout, device='cpu'):
@@ -162,22 +179,22 @@ def __init__(self, in_dim, hid_dim, n_layers, dropout, device='cpu'):
 		self.n_layers = n_layers
 		self.device = device
 
-		self.gru = nn.GRU(in_dim, hid_dim, n_layers, batch_first=True, dropout=dropout)
+		self.gru = nn.GRU(in_dim, hid_dim, num_layers=n_layers, batch_first=True, dropout=dropout)
 		# self.hidden = None
 
 	def forward(self, x):
 		h0 = torch.zeros(self.n_layers, x.shape[0], self.hid_dim).to(self.device)
 		out, h = self.gru(x, h0)
+		out, h =  out[-1, :, :], h[-1, :, :]  # Extracting from last layer
 		return out, h
 
 
 class RNNEncoder(nn.Module):
-	"""
-	    Encoder network containing enrolled GRU
+	""" Encoder network containing enrolled GRU
 	    :param in_dim: number of input features
 	    :param hid_dim: hidden size of the RNN
 	    :param n_layers: number of layers in RNN
-	    :param dropout: percentage of nodes to dropout
+	    :param dropout: dropout rate
 	    :param device: which device to use (cpu or cuda)
 	"""
 
@@ -201,17 +218,14 @@ def forward(self, x):
 		return h_end
 
 class RNNDecoder(nn.Module):
+	""" GRU-based Decoder network that converts latent vector into output
+		:param in_dim: number of input features
+		:param n_layers: number of layers in RNN
+		:param hid_dim: hidden size of the RNN
+		:param dropout: dropout rate
+		:param device: which device to use (cpu or cuda)
 	"""
-		    Decoder network converts latent vector into output
-		    :param window_size: length of the input sequence
-		    :param batch_size: batch size of the input sequence
-		    :param in_dim: number of input features
-		    :param n_layers: number of layers in RNN
-		    :param hid_dim: hidden size of the RNN
-		    :param batch_size: batch size
-		    :param dropout: percentage of nodes to dropout
-		    :param device: which device to use (cpu or cuda)
-		"""
+
 	def __init__(self, in_dim, n_layers, hid_dim, out_dim, dropout=0.0, device='cpu'):
 		super(RNNDecoder, self).__init__()
 
@@ -231,14 +245,19 @@ def forward(self, h):
 		return out
 
 class RNNAutoencoder(nn.Module):
-	"""
-	Reconstruction Model using a GRU-based Encoder-Decoder.
-	TODO
+	""" Reconstruction Model using a GRU-based Encoder-Decoder.
+		:param window_size: length of the input sequence
+		:param in_dim: number of input features
+		:param n_layers: number of layers in RNN
+		:param hid_dim: hidden size of the RNN
+		:param in_dim: number of output features
+		:param dropout: dropout rate
+		:param device: which device to use (cpu or cuda)
 	"""
 
 	def __init__(self,
 				 window_size,
-				 n_features,
+				 in_dim,
 				 n_layers,
 				 hid_dim,
 				 out_dim,
@@ -250,25 +269,29 @@ def __init__(self,
 		self.dropout = dropout
 		self.device = device
 
-		self.encoder = RNNEncoder(3*n_features, n_layers, hid_dim, dropout=dropout, device=device)
-		self.decoder = RNNDecoder(hid_dim, n_layers, hid_dim, out_dim, dropout=dropout, device=device)
+		# self.encoder = RNNEncoder(in_dim, n_layers, hid_dim, dropout=dropout, device=device)
+		self.decoder = RNNDecoder(in_dim, n_layers, hid_dim, out_dim, dropout=dropout, device=device)
 
 	def forward(self, x):
 		# x shape: (b, n, k)
 
-		h_end = self.encoder(x).squeeze(0).unsqueeze(2)
-
-		# print(f'h_end: {h_end.shape}')
+		# h_end = self.encoder(x).squeeze(0).unsqueeze(2)
+		# Use last hidden state from GRU module as the encoding
+		h_end = x
 		h_end_rep = h_end.repeat_interleave(self.window_size, dim=1).view(x.size(0), self.window_size, -1)
-		# print(h_end_rep.shape)
 
 		out = self.decoder(h_end_rep)
 		# print(f'out: {out.shape}')
 		return out
 
 class Forecasting_Model(nn.Module):
-	"""
-	Forecasting model (fully-connected network)
+	""" Forecasting model (fully-connected network)
+		:param in_dim: number of input features
+		:param hid_dim: hidden size of the FC network
+		:param out_dim: number of output features
+		:param n_layers: number of FC layers
+		:param dropout: dropout rate
+		:param device: which device to use (cpu or cuda)
 	"""
 	def __init__(self, in_dim, hid_dim, out_dim, n_layers, dropout, device='cpu'):
 		super(Forecasting_Model, self).__init__()

mtad_gat.py

@@ -5,24 +5,34 @@
 
 
 class MTAD_GAT(nn.Module):
-	def __init__(self,  num_nodes, window_size, horizon, out_dim, batch_size,
+	def __init__(self,  n_features, window_size, horizon, out_dim, batch_size,
 				 kernel_size=7,
 				 gru_n_layers=3,
 				 gru_hid_dim=64,
-				 autoencoder_n_layers=1,
-				 autoencoder_hid_dim=128,
-				 forecasting_n_layers=3,
-				 forecasting_hid_dim=32,
+				 autoenc_n_layers=1,
+				 autoenc_hid_dim=128,
+				 forecast_n_layers=3,
+				 forecast_hid_dim=32,
 				 dropout=0.2,
 				 alpha=0.2,
-				 device='cpu'):
+				 use_cuda=True):
 		super(MTAD_GAT, self).__init__()
-		self.conv = ConvLayer(num_nodes, window_size, kernel_size, device)
-		self.feature_gat = FeatureAttentionLayer(num_nodes, window_size, dropout, alpha, device)
-		self.temporal_gat = TemporalAttentionLayer(num_nodes, window_size, dropout, alpha, device)
-		self.gru = GRU(3*num_nodes, gru_hid_dim, gru_n_layers, dropout, device)
-		self.forecasting_model = Forecasting_Model(gru_hid_dim, forecasting_hid_dim, out_dim*horizon, forecasting_n_layers, dropout, device)
-		self.recon_model = RNNAutoencoder(window_size, num_nodes, autoencoder_n_layers, autoencoder_hid_dim, num_nodes, dropout, device)
+
+		device = 'cuda' if use_cuda and torch.cuda.is_available() else 'cpu'
+
+		self.conv = ConvLayer(n_features, window_size, kernel_size, device)
+		self.feature_gat = FeatureAttentionLayer(n_features, window_size, dropout, alpha, device)
+		self.temporal_gat = TemporalAttentionLayer(n_features, window_size, dropout, alpha, device)
+		self.gru = GRU(3*n_features, gru_hid_dim, gru_n_layers, dropout, device)
+		self.forecasting_model = Forecasting_Model(gru_hid_dim, forecast_hid_dim, out_dim*horizon, forecast_n_layers, dropout, device)
+		self.recon_model = RNNAutoencoder(window_size, gru_hid_dim, autoenc_n_layers, autoenc_hid_dim, n_features, dropout, device)
+
+		self.model_args = f'MTAD-GAT(k_size={kernel_size}, gru_n_layers={gru_n_layers}, gru_hid_dim={gru_hid_dim}, ' \
+						   f'autoencoder_n_layers={autoenc_n_layers}, autoencoder_hid_dim={autoenc_hid_dim}, ' \
+						   f'forecast_n_layers={forecast_n_layers}, forecast_hid_dim={forecast_hid_dim}, ' \
+						   f'dropout={dropout}, alpha={alpha})'
+	def __repr__(self):
+		return self.model_args
 
 	def forward(self, x):
 		# x shape (b, n, k): b - batch size, n - window size, k - number of nodes/features
@@ -31,20 +41,14 @@ def forward(self, x):
 		h_feat = self.feature_gat(x)
 		h_temp = self.temporal_gat(x)
 
-		#h_cat = x
-		#h_cat = torch.cat([x, h_temp], dim=2)
 		h_cat = torch.cat([x, h_feat.permute(0, 2, 1), h_temp], dim=2) # (b, n, 3k)
 
-		gru_out, last_hidden = self.gru(h_cat)
-
-		forecasting_in = last_hidden.view(x.shape[0], -1) #gru_out[:, -1, :]  # Hidden state for to last timestamp
+		gru_out, h_end = self.gru(h_cat)  # (
 
-		# Flatten
-		predictions = self.forecasting_model(forecasting_in)
+		h_end = h_end.view(x.shape[0], -1) #gru_out[:, -1, :]  # Hidden state for last timestamp
 
-		# TODO: Reconstruction model
-		#recons = None
-		recons = self.recon_model(h_cat)
+		predictions = self.forecasting_model(h_end)
+		recons = self.recon_model(h_end)
 
 		return predictions, recons
 

requirements.txt

@@ -1,3 +1,4 @@
 matplotlib
 sklearn
 pandas
+tensorboard