Update GAN/diffusion.py

GAN/diffusion.py

@@ -5,6 +5,7 @@ from tqdm.auto import tqdm
 import tensorflow as tf
 from tensorflow import keras
 from tensorflow.keras import layers
+from tensorflow.keras.layers import Layer
 from GAN.utils import linear_beta_schedule, cosine_beta_schedule
 import matplotlib.pyplot as plt
 import os
@@ -152,6 +153,26 @@ def TimeMLP(units, activation_fn=keras.activations.swish):
 
     return apply
 
+# 创建一个新的层来包装时间嵌入操作
+class TimeEmbeddingLayer(Layer):
+    def __init__(self, d_model, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        
+    def call(self, timesteps):
+        # 扩展维度
+        timesteps = tf.expand_dims(timesteps, -1)
+        
+        # 计算不同频率的正弦信号
+        dim = self.d_model // 2
+        inv_freq = 1.0 / (10000 ** (tf.range(0, dim, dtype=tf.float32) / dim))
+        
+        # 计算正弦波
+        sinusoid_inp = tf.einsum("i,j->ij", tf.squeeze(timesteps), inv_freq)
+        pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], axis=-1)
+        
+        return pos_emb
+
 # Kernel initializer to use
 def kernel_init(scale):
     scale = max(scale, 1e-10)
@@ -256,10 +277,13 @@ def build_model(time_len, fea_num, d_model=16, n_heads=2, encoder_type='dual'):
     
     # Input layers
     x_input = layers.Input(shape=(time_len, fea_num))
-    time_input = layers.Input(shape=())
     
-    # Time step embeddings
-    time_emb = get_time_embedding(time_input, d_model)
+    #time_input = layers.Input(shape=())
+    #time_emb = get_time_embedding(time_input, d_model)
+
+    time_input = Input(shape=(1,), name='time_input')
+    feature_input = Input(shape=(time_len, fea_num), name='feature_input')
+    time_emb = TimeEmbeddingLayer(d_model)(time_input)
     
     encoded_features = []