Files
cvpr/models/gnn.py
T
2023-02-25 12:33:18 +00:00

34 lines
1.6 KiB
Python

import os
import sys
import scipy.io as sio
import tensorflow as tf
from spektral.layers import GCNConv, GeneralConv
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, features=64, activation_middle='tanh', activation_final='tanh', learn_weights=False):
super().__init__()
self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
self.conv1 = GCNConv(features, activation=activation_middle, use_bias=False)
self.conv2 = GCNConv(features, activation=activation_middle, use_bias=False)
# self.conv3 = GCNConv(features, activation=activation_middle, use_bias=False)
# self.conv4 = GCNConv(64, activation='tanh', use_bias=False)
# self.conv5 = GCNConv(64, activation='tanh', use_bias=False)
self.conv6 = GCNConv(3 if not learn_weights else 24, activation=activation_final, 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