lara roth

Upload ChebAutoencoder model.

238676a 2 months ago

6.49 kB

import torch
import torch.nn as nn
# import torch.nn.functional as F
# import torch.optim as optim
import torch
# import torchvision.transforms as transforms
# import numpy as np
# import os
# from os import listdir
from os.path import join
# import pandas as pd
# import cv2 as cv
# from torch.utils.data import DataLoader, Dataset, random_split
# from PIL import Image, ImageOps
#from torchmetrics import Accuracy
from pytorch_lightning import LightningModule
# from pytorch_lightning.callbacks.early_stopping import EarlyStopping
# from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
#from pytorch_lightning.callbacks.progress import TQDMProgressBar
# from pytorch_lightning.loggers import TensorBoardLogger
# from torchvision import models
# import general_functions
# from data_loader import DatasetCostum
# import time
# import optuna
# from pytorch_lightning.callbacks.progress import TQDMProgressBar
# from pytorch_lightning.callbacks.early_stopping import EarlyStopping
# import statistics
# import panorama.datasets
# from optuna.integration import PyTorchLightningPruningCallback
# from graph_attention_layer_or import GATLayerImp3

# from pytorch_forecasting import TimeSeriesDataSet
# from pytorch_forecasting.data.encoders import TorchNormalizer
# from torch_geometric import EdgeIndex
# from graph.astgcn import ASTGCN
# from graph.gconv_gru import GConvGRU
#import ray
#from ray import tune
#from ray.tune.tuner import Tuner, TuneConfig
#from ray.tune.search.optuna import OptunaSearch
# import pickle
# import sqlite3
# import json
# from torch.utils.tensorboard import SummaryWriter
# from tgcn import TGCN2
from torch_geometric.nn import ChebConv, Sequential
# from torch.optim.lr_scheduler import ReduceLROnPlateau
# from GAT_layer import GATLayer, GraphAttentionLayer


class AutoEncoderModel(LightningModule):
    def __init__(self, cuda_true, batch_size):
        super().__init__()
    

        self.epochs, self.conditions = list(), list()
        self.recon_loss_test_step_list = list()
        self.num_step = 0
        
        if cuda_true:
            self.dev = "cuda"
        else:
            self.dev = "cpu"
        self.num_nodes = 15
        # self.optimizer_name = hyper_params["optimizer_name"]
      

        # self.tb_writer = summary_writer

        self.edge_index_att = None
        
        self.batch_size = batch_size  
        self.criterion = nn.MSELoss(reduction='mean')

        self.steps = 0
        self.recon_loss_train_step = 0
        self.recon_loss_train_step_list = list()
        self.recon_loss_tain_epoch_list = list()
        self.recon_loss_val_step = 0
        self.recon_loss_val_step_list = list()
        self.recon_loss_val_epoch_list = list()
        self.recon_loss_test_step_list = list()
        self.epoch = 0
        
        self.window = 64
        self.automatic_optimization = True
        self.test_target_data, self.test_predict_data = list(), list()
        # self.output_decoder = torch.rand(self.batch_size, self.num_nodes, self.window) # check size
        #  self.output_decoder.requires_grad_()
        self.output_first_layer_decoder = torch.rand(self.batch_size, self.num_nodes, self.window*4) # check size
        self.output_first_layer_decoder.requires_grad_()
        self.output_first_layer_decoder.to(self.dev)

        self.node_num_featues = 5
        self.total_feat = self.node_num_featues * self.window
        
        
        
        ### Original Code
        self.k = 4
        latent_dim = 104
        self.beta = 0.009256865323169841
        self.encoder = Sequential('x, edge_index', [
            (ChebConv(in_channels=self.window*self.node_num_featues, out_channels=self.window*2, K=self.k), 'x, edge_index -> x'),
            nn.ReLU(inplace=True),
            (ChebConv(in_channels=self.window*2, out_channels=self.window*4, K=self.k), 'x, edge_index -> x'),
            nn.ReLU(inplace=True),
        ])

        self.encoder_2 = Sequential('x, edge_index', [
            (ChebConv(in_channels=self.window, out_channels=self.window*4, K=self.k), 'x, edge_index -> x'),
            nn.ReLU(inplace=True)
        ])
        self.latent = nn.Sequential(
            nn.Flatten(),
            nn.Linear(self.window*4*self.num_nodes, latent_dim),

            nn.Linear(latent_dim, self.window*4*self.num_nodes),
            nn.Unflatten(-1, (int(self.num_nodes), int(self.window*4)))
            )
            
        
        self.latent.to(self.dev)
        self.decoder_2 = Sequential('x, edge_index' ,[
        (ChebConv(in_channels=self.window*4, out_channels=self.window, K=self.k), 'x, edge_index -> x'),
        nn.ReLU(inplace=True)
        ])
        self.decoder = Sequential('x, edge_index' ,[
        (ChebConv(in_channels=self.window*4, out_channels=self.window*2, K=self.k), 'x, edge_index -> x'),
        nn.ReLU(inplace=True),
        (ChebConv(in_channels=self.window*2, out_channels=self.window*self.node_num_featues, K=self.k), 'x, edge_index -> x')
        ])

    
        
        self.softmax = nn.Softmax()
        
 
    def forward(self, input_data, edge_indices, adj_matrix):
        self.edge_indices = edge_indices
        self.adj_matrix = adj_matrix
        # print("input data", input_data)

        

       
        input_data_reshaped = torch.reshape(input=input_data, shape=(input_data.shape[0], input_data.shape[2], input_data.shape[3] * input_data.shape[1]))
        input_data_reshaped = input_data_reshaped.to(self.dev)
        output_encoder = self.encoder(input_data_reshaped, self.edge_indices)
        scaled_encoder = torch.mul(output_encoder, self.beta)
        output_latent = self.latent(scaled_encoder)
        output_decoder = self.decoder(output_latent, self.edge_indices)
        self.output_decoder = torch.reshape(input=output_decoder, shape=(output_decoder.shape[0], self.window, self.num_nodes, self.node_num_featues))

        recon_loss_list = list()
        for i in range(input_data.shape[1]):
            recon_loss_list.append(self.criterion(self.output_decoder[:,i,:,:], input_data[:,i,:,:]).to(self.dev))
        recon_loss = sum(recon_loss_list)/len(recon_loss_list)
    
        return recon_loss, self.output_decoder
 


def calc_edge_weight(edge_index, adj_matrix):
    edge_weight = torch.rand(edge_index.shape[1])
    for i, element in enumerate(edge_index.T): 
        edge_weight[i] = (adj_matrix[element[0]][element[1]] + adj_matrix[element[1]][element[0]])/2.0
    return edge_weight

	import torch
	import torch.nn as nn
	# import torch.nn.functional as F
	# import torch.optim as optim
	import torch
	# import torchvision.transforms as transforms
	# import numpy as np
	# import os
	# from os import listdir
	from os.path import join
	# import pandas as pd
	# import cv2 as cv
	# from torch.utils.data import DataLoader, Dataset, random_split
	# from PIL import Image, ImageOps
	#from torchmetrics import Accuracy
	from pytorch_lightning import LightningModule
	# from pytorch_lightning.callbacks.early_stopping import EarlyStopping
	# from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
	#from pytorch_lightning.callbacks.progress import TQDMProgressBar
	# from pytorch_lightning.loggers import TensorBoardLogger
	# from torchvision import models
	# import general_functions
	# from data_loader import DatasetCostum
	# import time
	# import optuna
	# from pytorch_lightning.callbacks.progress import TQDMProgressBar
	# from pytorch_lightning.callbacks.early_stopping import EarlyStopping
	# import statistics
	# import panorama.datasets
	# from optuna.integration import PyTorchLightningPruningCallback
	# from graph_attention_layer_or import GATLayerImp3

	# from pytorch_forecasting import TimeSeriesDataSet
	# from pytorch_forecasting.data.encoders import TorchNormalizer
	# from torch_geometric import EdgeIndex
	# from graph.astgcn import ASTGCN
	# from graph.gconv_gru import GConvGRU
	#import ray
	#from ray import tune
	#from ray.tune.tuner import Tuner, TuneConfig
	#from ray.tune.search.optuna import OptunaSearch
	# import pickle
	# import sqlite3
	# import json
	# from torch.utils.tensorboard import SummaryWriter
	# from tgcn import TGCN2
	from torch_geometric.nn import ChebConv, Sequential
	# from torch.optim.lr_scheduler import ReduceLROnPlateau
	# from GAT_layer import GATLayer, GraphAttentionLayer


	class AutoEncoderModel(LightningModule):
	def __init__(self, cuda_true, batch_size):
	super().__init__()


	self.epochs, self.conditions = list(), list()
	self.recon_loss_test_step_list = list()
	self.num_step = 0

	if cuda_true:
	self.dev = "cuda"
	else:
	self.dev = "cpu"
	self.num_nodes = 15
	# self.optimizer_name = hyper_params["optimizer_name"]


	# self.tb_writer = summary_writer

	self.edge_index_att = None

	self.batch_size = batch_size
	self.criterion = nn.MSELoss(reduction='mean')

	self.steps = 0
	self.recon_loss_train_step = 0
	self.recon_loss_train_step_list = list()
	self.recon_loss_tain_epoch_list = list()
	self.recon_loss_val_step = 0
	self.recon_loss_val_step_list = list()
	self.recon_loss_val_epoch_list = list()
	self.recon_loss_test_step_list = list()
	self.epoch = 0

	self.window = 64
	self.automatic_optimization = True
	self.test_target_data, self.test_predict_data = list(), list()
	# self.output_decoder = torch.rand(self.batch_size, self.num_nodes, self.window) # check size
	# self.output_decoder.requires_grad_()
	self.output_first_layer_decoder = torch.rand(self.batch_size, self.num_nodes, self.window*4) # check size
	self.output_first_layer_decoder.requires_grad_()
	self.output_first_layer_decoder.to(self.dev)

	self.node_num_featues = 5
	self.total_feat = self.node_num_featues * self.window



	### Original Code
	self.k = 4
	latent_dim = 104
	self.beta = 0.009256865323169841
	self.encoder = Sequential('x, edge_index', [
	(ChebConv(in_channels=self.windowself.node_num_featues, out_channels=self.window2, K=self.k), 'x, edge_index -> x'),
	nn.ReLU(inplace=True),
	(ChebConv(in_channels=self.window2, out_channels=self.window4, K=self.k), 'x, edge_index -> x'),
	nn.ReLU(inplace=True),
	])

	self.encoder_2 = Sequential('x, edge_index', [
	(ChebConv(in_channels=self.window, out_channels=self.window*4, K=self.k), 'x, edge_index -> x'),
	nn.ReLU(inplace=True)
	])
	self.latent = nn.Sequential(
	nn.Flatten(),
	nn.Linear(self.window4self.num_nodes, latent_dim),

	nn.Linear(latent_dim, self.window4self.num_nodes),
	nn.Unflatten(-1, (int(self.num_nodes), int(self.window*4)))
	)


	self.latent.to(self.dev)
	self.decoder_2 = Sequential('x, edge_index' ,[
	(ChebConv(in_channels=self.window*4, out_channels=self.window, K=self.k), 'x, edge_index -> x'),
	nn.ReLU(inplace=True)
	])
	self.decoder = Sequential('x, edge_index' ,[
	(ChebConv(in_channels=self.window4, out_channels=self.window2, K=self.k), 'x, edge_index -> x'),
	nn.ReLU(inplace=True),
	(ChebConv(in_channels=self.window2, out_channels=self.windowself.node_num_featues, K=self.k), 'x, edge_index -> x')
	])



	self.softmax = nn.Softmax()


	def forward(self, input_data, edge_indices, adj_matrix):
	self.edge_indices = edge_indices
	self.adj_matrix = adj_matrix
	# print("input data", input_data)




	input_data_reshaped = torch.reshape(input=input_data, shape=(input_data.shape[0], input_data.shape[2], input_data.shape[3] * input_data.shape[1]))
	input_data_reshaped = input_data_reshaped.to(self.dev)
	output_encoder = self.encoder(input_data_reshaped, self.edge_indices)
	scaled_encoder = torch.mul(output_encoder, self.beta)
	output_latent = self.latent(scaled_encoder)
	output_decoder = self.decoder(output_latent, self.edge_indices)
	self.output_decoder = torch.reshape(input=output_decoder, shape=(output_decoder.shape[0], self.window, self.num_nodes, self.node_num_featues))

	recon_loss_list = list()
	for i in range(input_data.shape[1]):
	recon_loss_list.append(self.criterion(self.output_decoder[:,i,:,:], input_data[:,i,:,:]).to(self.dev))
	recon_loss = sum(recon_loss_list)/len(recon_loss_list)

	return recon_loss, self.output_decoder



	def calc_edge_weight(edge_index, adj_matrix):
	edge_weight = torch.rand(edge_index.shape[1])
	for i, element in enumerate(edge_index.T):
	edge_weight[i] = (adj_matrix[element[0]][element[1]] + adj_matrix[element[1]][element[0]])/2.0
	return edge_weight