<filename>pyramid_attention.py
import torch
from torch import nn
import numpy as np
from axial_attention import SelfAttention
from positional_encodings.positional_encodings import PositionalEncoding1D, PositionalEncoding2D, PositionalEncoding3D
from common_attention.attention import attention
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.layers import Reshape, InputSpec

import tfpyth

import argparse
import os

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)

class Pyramid_attention_layer(layers.Layer):
    def __init__(self, **kwargs):
        super(Pyramid_attention_layer, self).__init__(**kwargs)
        self.input_spec = InputSpec(min_ndim=4)

    def torch_2_tensorflow(self, data):
        '''
        Transfer PositionalEncoding3D to tensorflow version
        '''
        model = PositionalEncoding3D(11)
        model = tfpyth.eager_tensorflow_from_torch(model)

        model = tf.py_function(model, [data], tf.float32)
        model.set_shape(data.get_shape())
        return model

    def attend(self, data1, data2):
        '''
        calculate directional structural attention mechanism
        '''
        data_attend = tfpyth.tensorflow_from_torch(attention.attend, [data1, data2], tf.float32)
        return data_attend

    def stack_tensor(self, data):
        tmp_data = []
        for a in data:
            tmp_a = []
            for b in a:
                tmp_a.append(tf.stack(b, 1))
            tmp_data.append(tf.stack(tmp_a, 1))
        data = tf.stack(tmp_data, 1)
        return data


    def data_reconstruction(self, data, k_s = 3):

        (_, n_g1, n_g2, n_t, _) = data.shape
        

        # data = data + self.torch_2_tensorflow(data) # data + positional information

        data = tf.pad(data,[[0,0],[k_s-1,k_s-1],[k_s-1,k_s-1],[0,k_s-1],[0,0]],constant_values = 0)

        target = []
        for i in range(n_t):
            target.append([])
            for j in range(k_s-1, k_s-1+n_g1):
                target[i].append([])
                for k in range(k_s-1, k_s-1+n_g2):
                    target[i][j-(k_s-1)].append([])
                    t = []
                    for m in range(k_s):
                        horiz_start = j-m
                        horiz_end = j+m+1
                        vert_start = k-m
                        vert_end = k+m+1
                        dep = i+m
                        tmp = data[:, horiz_start:horiz_end, vert_start:vert_end, dep, :] # the slice for each layer in pyramid structure 
                        tmp = tf.reshape(tmp, (-1, tmp.shape[1]*tmp.shape[2], tmp.shape[3])) # reshape it into an array
                        t.append(tmp) # concatenate all the slices
                    t = tf.concat(t, 1)
                    target[i][j-(k_s-1)][k-(k_s-1)] = t
        target = self.stack_tensor(target)
        return target
    
    def call(self, inputs):
        '''
        The shape of data: [1, grid_1, grid_2, n_t, n_features]
        The shape of outputs: [n_t, grid_1, grid_2, n_features]
        '''

        data = self.data_reconstruction(inputs) # none * 20 * 10 * 10 * 35 * 1

        batch_size, n_g1, n_g2, n_t, n_g, dim = data.shape # n_g: the grouped attention blocks

        attenlist = []
        for i in range(n_t):
            atten_i = []
            for j in range(n_g1):
                atten_j = []
                for k in range(n_g2):
                    data1 = tf.expand_dims(data[:,j,k,i,0], axis=-1) # target
                    data2 = data[:,j,k,i,1:] # the sources of traffic

                    attened_data = self.attend(data1, data2)
                    attened_data.set_shape(data1.get_shape())
                    atten_j.append(attened_data)
                atten_i.append(tf.stack(atten_j, 1))
            attenlist.append(tf.stack(atten_i, 1))

        data_attn = tf.stack(attenlist, 1)
        return data_attn

# if __name__ == '__main__':
# # add one more axis to restore pyramid-shaped data info
#     g = tf.random.get_global_generator()
#     data = g.normal(shape=[1, 16, 16, 8, 4])
    
#     pyramid_attention(data)

<filename>pyramid_attention.py
import torch
from torch import nn
import numpy as np
from axial_attention import SelfAttention
from positional_encodings.positional_encodings import PositionalEncoding1D, PositionalEncoding2D, PositionalEncoding3D
from common_attention.attention import attention
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.layers import Reshape, InputSpec

import tfpyth

import argparse
import os

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)

class Pyramid_attention_layer(layers.Layer):
    def __init__(self, **kwargs):
        super(Pyramid_attention_layer, self).__init__(**kwargs)
        self.input_spec = InputSpec(min_ndim=4)

    def torch_2_tensorflow(self, data):
        '''
        Transfer PositionalEncoding3D to tensorflow version
        '''
        model = PositionalEncoding3D(11)
        model = tfpyth.eager_tensorflow_from_torch(model)

        model = tf.py_function(model, [data], tf.float32)
        model.set_shape(data.get_shape())
        return model

    def attend(self, data1, data2):
        '''
        calculate directional structural attention mechanism
        '''
        data_attend = tfpyth.tensorflow_from_torch(attention.attend, [data1, data2], tf.float32)
        return data_attend

    def stack_tensor(self, data):
        tmp_data = []
        for a in data:
            tmp_a = []
            for b in a:
                tmp_a.append(tf.stack(b, 1))
            tmp_data.append(tf.stack(tmp_a, 1))
        data = tf.stack(tmp_data, 1)
        return data


    def data_reconstruction(self, data, k_s = 3):

        (_, n_g1, n_g2, n_t, _) = data.shape
        

        # data = data + self.torch_2_tensorflow(data) # data + positional information

        data = tf.pad(data,[[0,0],[k_s-1,k_s-1],[k_s-1,k_s-1],[0,k_s-1],[0,0]],constant_values = 0)

        target = []
        for i in range(n_t):
            target.append([])
            for j in range(k_s-1, k_s-1+n_g1):
                target[i].append([])
                for k in range(k_s-1, k_s-1+n_g2):
                    target[i][j-(k_s-1)].append([])
                    t = []
                    for m in range(k_s):
                        horiz_start = j-m
                        horiz_end = j+m+1
                        vert_start = k-m
                        vert_end = k+m+1
                        dep = i+m
                        tmp = data[:, horiz_start:horiz_end, vert_start:vert_end, dep, :] # the slice for each layer in pyramid structure 
                        tmp = tf.reshape(tmp, (-1, tmp.shape[1]*tmp.shape[2], tmp.shape[3])) # reshape it into an array
                        t.append(tmp) # concatenate all the slices
                    t = tf.concat(t, 1)
                    target[i][j-(k_s-1)][k-(k_s-1)] = t
        target = self.stack_tensor(target)
        return target
    
    def call(self, inputs):
        '''
        The shape of data: [1, grid_1, grid_2, n_t, n_features]
        The shape of outputs: [n_t, grid_1, grid_2, n_features]
        '''

        data = self.data_reconstruction(inputs) # none * 20 * 10 * 10 * 35 * 1

        batch_size, n_g1, n_g2, n_t, n_g, dim = data.shape # n_g: the grouped attention blocks

        attenlist = []
        for i in range(n_t):
            atten_i = []
            for j in range(n_g1):
                atten_j = []
                for k in range(n_g2):
                    data1 = tf.expand_dims(data[:,j,k,i,0], axis=-1) # target
                    data2 = data[:,j,k,i,1:] # the sources of traffic

                    attened_data = self.attend(data1, data2)
                    attened_data.set_shape(data1.get_shape())
                    atten_j.append(attened_data)
                atten_i.append(tf.stack(atten_j, 1))
            attenlist.append(tf.stack(atten_i, 1))

        data_attn = tf.stack(attenlist, 1)
        return data_attn

# if __name__ == '__main__':
# # add one more axis to restore pyramid-shaped data info
#     g = tf.random.get_global_generator()
#     data = g.normal(shape=[1, 16, 16, 8, 4])
    
#     pyramid_attention(data)

Transfer PositionalEncoding3D to tensorflow version calculate directional structural attention mechanism # data = data + self.torch_2_tensorflow(data) # data + positional information # the slice for each layer in pyramid structure # reshape it into an array # concatenate all the slices The shape of data: [1, grid_1, grid_2, n_t, n_features]         The shape of outputs: [n_t, grid_1, grid_2, n_features] # none * 20 * 10 * 10 * 35 * 1 # n_g: the grouped attention blocks # target # the sources of traffic # if __name__ == '__main__': # # add one more axis to restore pyramid-shaped data info #     g = tf.random.get_global_generator() #     data = g.normal(shape=[1, 16, 16, 8, 4]) #     pyramid_attention(data)