import tensorflow as tf
from ..model_unit import BaseModel
from ..model_unit.GraphModelLayers import GCL
[docs]class ST_MGCN(BaseModel):
"""
Args:
T(int): Input sequence length
input_dim(int): Input feature dimension
num_graph(int): Number of graphs used in the model.
gcl_k(int): The highest order of Chebyshev Polynomial approximation in GCN.
gcl_l(int): Number of GCN layers.
lstm_units(int): Number of hidden units of RNN.
lstm_layers(int): Number of LSTM layers.
lr(float): Learning rate
external_dim(int): Dimension of the external feature, e.g. temperature and wind are two dimension.
code_version(str): Current version of this model code, which will be used as filename for saving the model
model_dir(str): The directory to store model files. Default:'model_dir'.
gpu_device(str): To specify the GPU to use. Default: '0'.
"""
def __init__(self,
T,
input_dim,
num_graph,
gcl_k,
gcl_l,
lstm_units,
lstm_layers,
lr,
external_dim,
code_version,
model_dir,
gpu_device):
super(ST_MGCN, self).__init__(code_version=code_version,
model_dir=model_dir,
gpu_device=gpu_device)
self._T = T
self._input_dim = input_dim
self._num_graph = num_graph
self._gcl_k = gcl_k
self._gcl_l = gcl_l
self._lstm_units = lstm_units
self._lstm_layers = lstm_layers
self._lr = lr
self._external_dim = external_dim
[docs] def build(self, init_vars=True, max_to_keep=5):
with self._graph.as_default():
# [batch, T, num_stations, input_dim]
traffic_flow = tf.placeholder(tf.float32, [None, self._T, None, self._input_dim])
laplacian_matrix = tf.placeholder(tf.float32, [self._num_graph, None, None])
target = tf.placeholder(tf.float32, [None, None, 1])
self._input['traffic_flow'] = traffic_flow.name
self._input['laplace_matrix'] = laplacian_matrix.name
self._input['target'] = target.name
station_number = tf.shape(traffic_flow)[-2]
outputs = []
for graph_index in range(self._num_graph):
with tf.variable_scope('CGRNN_Graph%s' % graph_index, reuse=False):
f_k_g = GCL.add_multi_gc_layers(tf.reshape(traffic_flow, [-1, station_number, self._input_dim]),
gcn_k=1, gcn_l=1, output_size=self._input_dim,
laplacian_matrix=laplacian_matrix[graph_index],
activation=tf.nn.tanh)
f_k_g = tf.reshape(f_k_g, [-1, self._T, station_number, self._input_dim])
x_hat = tf.concat([f_k_g, traffic_flow], axis=-1)
z = tf.reduce_mean(x_hat, axis=-2, keepdims=True)
s = tf.layers.dense(tf.layers.dense(z, units=4, use_bias=False, activation=tf.nn.relu),
units=1, use_bias=False, activation=tf.nn.sigmoid)
x_rnn = tf.multiply(traffic_flow, tf.tile(s, [1, 1, station_number, self._input_dim]))
x_rnn = tf.reshape(tf.transpose(x_rnn, perm=[0, 2, 1, 3]), [-1, self._T, self._input_dim])
for lstm_layer_index in range(self._lstm_layers):
x_rnn = tf.keras.layers.LSTM(units=self._lstm_units,
activation='tanh',
dropout=0.1,
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
return_sequences=True if lstm_layer_index<self._lstm_layers-1
else False)\
(x_rnn)
H = tf.reshape(x_rnn, [-1, station_number, self._lstm_units])
outputs.append(H)
outputs = tf.reduce_sum(outputs, axis=0)
# external dims
if self._external_dim is not None and self._external_dim > 0:
external_input = tf.placeholder(tf.float32, [None, self._external_dim])
self._input['external_feature'] = external_input.name
external_dense = tf.layers.dense(inputs=external_input, units=10)
external_dense = tf.tile(tf.reshape(external_dense, [-1, 1, 10]),
[1, tf.shape(outputs)[-2], 1])
outputs = tf.concat([outputs, external_dense], axis=-1)
prediction = tf.layers.dense(outputs, units=1)
loss = tf.sqrt(tf.reduce_mean(tf.square(prediction - target)))
train_operation = tf.train.AdamOptimizer(self._lr).minimize(loss, name='train_op')
# record output
self._output['prediction'] = prediction.name
self._output['loss'] = loss.name
# record train operation
self._op['train_op'] = train_operation.name
super(ST_MGCN, self).build(init_vars=init_vars, max_to_keep=max_to_keep)
# Step 1 : Define your '_get_feed_dict function‘, map your input to the tf-model
def _get_feed_dict(self, traffic_flow, laplace_matrix, target=None, external_feature=None):
feed_dict = {
'traffic_flow': traffic_flow,
'laplace_matrix': laplace_matrix,
}
if target is not None:
feed_dict['target'] = target
if external_feature is not None:
feed_dict['external_feature'] = external_feature
return feed_dict