File size: 6,492 Bytes
238676a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | 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 |