67 lines
3.4 KiB
Python
67 lines
3.4 KiB
Python
import os
|
|
import sys
|
|
|
|
import scipy.io as sio
|
|
import tensorflow as tf
|
|
from spektral.layers import GCNConv, GeneralConv, GATConv
|
|
from spektral.utils import sp_matrix_to_sp_tensor
|
|
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
|
|
GlobalAveragePooling2D, Layer,
|
|
LayerNormalization, Flatten, Reshape, UpSampling2D)
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
|
|
class Smoother(tf.keras.Model):
|
|
def __init__(self, ):
|
|
super().__init__()
|
|
self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
|
|
self.conv1 = GeneralConv(channels=32, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True)
|
|
self.conv2 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True)
|
|
self.conv3 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
|
|
self.conv4 = GeneralConv(channels=16, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True)
|
|
|
|
self.conv5 = GCNConv(3, activation='tanh', use_bias=False)
|
|
self.conv6 = GCNConv(3, activation='tanh', use_bias=False)
|
|
|
|
def call(self, inputs, training=None):
|
|
x = inputs
|
|
x = self.conv1([x, self.a], training=training)
|
|
x = self.conv2([x, self.a], training=training)
|
|
x = self.conv3([x, self.a], training=training)
|
|
x = self.conv4([x, self.a], training=training)
|
|
x = self.conv5([x, self.a], training=training)
|
|
x = self.conv6([x, self.a], training=training)
|
|
return x
|
|
|
|
|
|
class GNNDecoder(tf.keras.Model):
|
|
def __init__(self, ):
|
|
super().__init__()
|
|
self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
|
|
self.a_dense = sio.loadmat("custom_utils/ssmpl0.mat")['L'].todense()
|
|
|
|
self.flatten = Flatten()
|
|
self.conv1 = GATConv(channels=1024, attn_heads=4, concat_heads=True, dropout_rate=0.0, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
|
|
self.conv2 = GATConv(channels=512, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
|
|
self.conv3 = GATConv(channels=256, attn_heads=4, concat_heads=True, dropout_rate=0.1, return_attn_coef=False, add_self_loops=True, activation=None, use_bias=True)
|
|
self.conv4 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='sum', activation='prelu', use_bias=True)
|
|
self.conv5 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='sum', activation='prelu', use_bias=True)
|
|
self.conv6 = GeneralConv(channels=3, batch_norm=False, dropout=0.0, aggregate='sum', activation='tanh', use_bias=True)
|
|
|
|
def call(self, inputs, training=None):
|
|
x = inputs
|
|
x = self.flatten(x)
|
|
x = tf.tile(x[:,None,:], [1,14475,1])
|
|
|
|
x = self.conv1([x, self.a_dense], training=training)
|
|
x = self.conv2([x, self.a_dense], training=training)
|
|
x = self.conv3([x, self.a_dense], training=training)
|
|
x = self.conv4([x, self.a], training=training)
|
|
x = self.conv5([x, self.a], training=training)
|
|
x = self.conv6([x, self.a], training=training)
|
|
return x
|
|
|
|
|
|
|