A lot of change and a new implemantation from scratch of MPS and MPO …

…for further work on tensor network

layers/TensorAlexNet.py

@@ -0,0 +1,10 @@
+from keras.layers import Layer, Conv2D, BatchNormalization, MaxPooling2D, Flatten, Dropout
+from keras import Model
+
+class TensorAlexNet(Model):
+
+    def __init__(self):
+        super(TensorAlexNet, self).__init__()
+
+    def call(self):
+        pass

layers/TensorAutoEncoder.py

@@ -0,0 +1,182 @@
+import tensorflow as tf
+import numpy as np
+
+from keras.datasets import mnist
+from keras.losses import binary_crossentropy
+from keras.metrics import Mean
+from keras.layers import Layer, Dense, Input, Conv2D, Flatten
+from keras import Model
+
+from syngular.layers.TensorDense import TensorDense
+
+class Sampling(Layer):
+
+    def call(self, inputs):
+        z_mean, z_log_var = inputs
+
+        batch = tf.shape(z_mean)[0]
+        dim = tf.shape(z_mean)[1]
+
+        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
+
+        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
+
+class Encoder(Layer):
+
+    def __init__(self, latent_dim=(8,8), intermediate_dim=(16,16), name="encoder", **kwargs):
+        super(Encoder, self).__init__(name=name, **kwargs)
+
+        self.latent_dim = latent_dim
+        self.intermediate_dim = intermediate_dim
+
+        self.latent_bond_dim = tuple(d//4 for d in self.latent_dim[:-1])
+        self.intermediate_bond_dim = tuple(d//4 for d in self.intermediate_dim[:-1])
+
+        self.conv1 = Conv2D(32, 3, activation="relu", strides=2, padding="same")
+        self.conv2 = Conv2D(64, 3, activation="relu", strides=2, padding="same")
+
+        self.tt_dense_proj = Dense(64) #TensorDense(None, self.intermediate_dim, self.intermediate_bond_dim)
+        self.tt_dense_mean = Dense(32) #TensorDense(None, self.latent_dim, self.latent_bond_dim)
+        self.tt_dense_log_var = Dense(32) #TensorDense(None, self.latent_dim, self.latent_bond_dim)
+
+        self.sampling = Sampling()
+
+    def call(self, inputs):
+        x = self.conv1(inputs)
+        x = self.conv2(x)
+        x = Flatten()(x)
+
+        x = self.tt_dense_proj(x)
+        z_mean = self.tt_dense_mean(x)
+        z_log_var = self.tt_dense_log_var(x)
+
+        z = self.sampling((z_mean, z_log_var))
+
+        return z_mean, z_log_var, z
+
+class Decoder(Layer):
+
+    def __init__(self, original_dim, intermediate_dim=(16,16), name="decoder", **kwargs):
+        super(Decoder, self).__init__(name=name, **kwargs)
+
+        self.original_dim = original_dim
+        self.intermediate_dim = intermediate_dim
+
+        self.original_bond_dim = tuple(d//4 for d in self.original_dim[:-1])
+        self.intermediate_bond_dim = tuple(d//4 for d in self.intermediate_dim[:-1])
+
+        self.tt_dense_proj = Dense(64)#TensorDense(None, self.intermediate_dim, self.intermediate_bond_dim)
+        self.tt_dense_output = Dense(28*28) #TensorDense(None, self.original_dim, self.original_bond_dim) #, activation=tf.nn.sigmoid)
+
+    def call(self, inputs):
+        x = self.tt_dense_proj(inputs)
+
+        return self.tt_dense_output(x)
+
+class TensorAutoEncoder(Model):
+
+    def __init__(self, original_dim,  intermediate_dim=(16,16), latent_dim=(8,8), name="tensorautoencoder", **kwargs):
+        super(TensorAutoEncoder, self).__init__(name=name, **kwargs)
+
+        self.original_dim = original_dim
+        self.encoder = Encoder(latent_dim=latent_dim, intermediate_dim=intermediate_dim)
+        self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)
+
+        self.total_loss_tracker = Mean(name="total_loss")
+        self.reconstruction_loss_tracker = Mean(name="reconstruction_loss")
+        self.kl_loss_tracker = Mean(name="kl_loss")
+
+    @property
+    def metrics(self):
+        return [
+            self.total_loss_tracker,
+            self.reconstruction_loss_tracker,
+            self.kl_loss_tracker,
+        ]
+
+    def call(self, inputs):
+        z_mean, z_log_var, z = self.encoder(inputs)
+        reconstructed = self.decoder(z)
+
+        kl_loss = -0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
+
+        self.add_loss(kl_loss)
+        return reconstructed
+
+    def train_step(self, inputs):
+        with tf.GradientTape() as tape:
+            z_mean, z_log_var, z = self.encoder(inputs)
+
+            reconstruction = self.decoder(z)
+
+            print("train shape", inputs.shape, reconstruction.shape)
+            tf.reshape(inputs, (tf.shape(inputs)[0],28*28))
+
+            # reconstruction_loss = tf.reduce_mean(tf.reduce_sum(binary_crossentropy(inputs, reconstruction), axis=(1, 2)))
+            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
+            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
+
+            total_loss = kl_loss #reconstruction_loss + kl_loss
+
+        grads = tape.gradient(total_loss, self.trainable_weights)
+
+        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
+
+        self.total_loss_tracker.update_state(total_loss)
+        # self.reconstruction_loss_tracker.update_state(reconstruction_loss)
+        self.kl_loss_tracker.update_state(kl_loss)
+
+        return {
+            "loss": self.total_loss_tracker.result(),
+            # "reconstruction_loss": self.reconstruction_loss_tracker.result(),
+            "kl_loss": self.kl_loss_tracker.result(),
+        }
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+
+    (x_train, _), (x_test, _) = mnist.load_data()
+    mnist_digits = np.concatenate([x_train, x_test], axis=0)
+    mnist_digits = np.expand_dims(mnist_digits, -1).astype("float32") / 255
+
+    tae = TensorAutoEncoder(original_dim=(28,28), intermediate_dim=(16,16), latent_dim=(8,8))
+    tae.compile(optimizer="adam")
+    tae.fit(mnist_digits, epochs=30, batch_size=128, verbose=0)
+
+    def plot_latent_space(tae, n=30, figsize=15):
+        # display a n*n 2D manifold of digits
+        digit_size = 28
+        scale = 1.0
+        figure = np.zeros((digit_size * n, digit_size * n))
+        # linearly spaced coordinates corresponding to the 2D plot
+        # of digit classes in the latent space
+        grid_x = np.linspace(-scale, scale, n)
+        grid_y = np.linspace(-scale, scale, n)[::-1]
+
+        for i, yi in enumerate(grid_y):
+            for j, xi in enumerate(grid_x):
+                z_sample = np.array([[xi, yi]])
+                x_decoded = tae.decoder(z_sample)
+                digit = x_decoded[0].reshape(digit_size, digit_size)
+                figure[
+                    i * digit_size : (i + 1) * digit_size,
+                    j * digit_size : (j + 1) * digit_size,
+                ] = digit
+
+        plt.figure(figsize=(figsize, figsize))
+        start_range = digit_size // 2
+        end_range = n * digit_size + start_range
+        pixel_range = np.arange(start_range, end_range, digit_size)
+        sample_range_x = np.round(grid_x, 1)
+        sample_range_y = np.round(grid_y, 1)
+        plt.xticks(pixel_range, sample_range_x)
+        plt.yticks(pixel_range, sample_range_y)
+        plt.xlabel("z[0]")
+        plt.ylabel("z[1]")
+        plt.imshow(figure, cmap="Greys_r")
+        plt.show()
+
+
+    plot_latent_space(tae)

layers/TensorDense.py

@@ -1,13 +1,132 @@
 from functools import reduce
+from itertools import zip_longest
+from numpy.core.einsumfunc import einsum
 
 import tensorflow as tf
 from tensorflow.keras.layers import Layer
 
+from math import floor, ceil, sqrt
 import tensornetwork as tn
 import numpy as np
 
+def unfold_dim(shape):
+    return reduce(lambda x, y: x*y, shape)
+
 class TensorDense(Layer):
 
+    def __init__(self, tt_input_shape, tt_output_shape, tt_bond_shape, activation=tf.nn.relu):
+        super(TensorDense, self).__init__()
+
+        self.tt_input_shape = tt_input_shape
+        self.tt_output_shape = tt_output_shape
+        self.tt_bond_shape = tt_bond_shape
+
+        self.tt_input_shape_unfold = unfold_dim(self.tt_input_shape) if self.tt_input_shape != None else None
+        self.tt_output_shape_unfold = unfold_dim(self.tt_output_shape)
+
+        if (self.tt_input_shape != None and (len(self.tt_input_shape) != len(self.tt_output_shape) or len(self.tt_input_shape) != len(self.tt_bond_shape)+1)) or len(self.tt_output_shape) != len(self.tt_bond_shape)+1:
+            raise Exception(f"Incompatible shapes. Cannot create TensorDense with {len(self.tt_input_shape)} {len(self.tt_output_shape)} and {len(self.tt_bond_shape)} ")
+
+        self.cores = []
+        self.cores_number = len(self.tt_input_shape) if self.tt_input_shape != None else len(self.tt_output_shape)
+
+        self.activation = activation
+
+    def build(self, tt_input_shape):
+
+        if self.tt_input_shape == None:
+
+            roots = int(np.power(tt_input_shape[1:], 1/self.cores_number))
+
+            self.tt_input_shape = tuple([roots] * self.cores_number)
+
+            print("reshaped", tt_input_shape)
+
+        print("Input shape", self.tt_input_shape)
+
+        self.bias = tf.Variable(tf.zeros(shape=self.tt_output_shape), name="bias", trainable=True)
+
+        last_idx = self.cores_number-1        
+
+        # Creating the first core of the weight MPS
+        self.cores.append(self.add_weight(
+            shape = (self.tt_input_shape[0], self.tt_output_shape[0], self.tt_bond_shape[0]),
+            name = "core1",
+            initializer = "random_normal",
+            trainable = True
+        ))
+
+        # Creating the intermediate cores of the weight MPS
+        for idx in range(1, self.cores_number-1):
+            self.cores.append(self.add_weight(
+                shape = (self.tt_input_shape[idx], self.tt_output_shape[idx], self.tt_bond_shape[idx-1], self.tt_bond_shape[idx]),
+                name = f"core{idx}",
+                initializer="random_normal",
+                trainable=True
+            ))
+
+        # Creating the last core of the weight MPS
+        self.cores.append(self.add_weight(
+            shape = (self.tt_input_shape[last_idx], self.tt_output_shape[last_idx], self.tt_bond_shape[last_idx-1]),
+            name = f"core{self.cores_number}",
+            initializer = "random_normal",
+            trainable = True
+        ))
+
+    def call(self, inputs):
+
+        def process(input, bias):
+            input_reshaped = tf.reshape(input,self.tt_input_shape)
+
+            last_idx = self.cores_number-1
+            einsum_structure = []
+
+            cores = [tn.Node(core, backend="tensorflow").tensor for core in self.cores]
+            x = tn.Node(input_reshaped, backend="tensorflow")
+
+
+            einsum_structure = []
+            einsum_structure.append(list(range(1, self.cores_number+1)))
+            einsum_structure.append([1, -(self.cores_number+1), 2*self.cores_number+1])
+
+            for idx in range(1, last_idx):
+                einsum_structure.append([idx+1, -(self.cores_number+idx+1), 2*self.cores_number+idx+1])
+
+            einsum_structure.append([last_idx+1, -(self.cores_number+last_idx+1), 2*self.cores_number+last_idx-1+1])
+
+            print(einsum_structure)
+
+            result = tn.ncon(
+                tensors = [x.tensor] + cores,
+                network_structure = einsum_structure,
+                backend = "tensorflow"
+            )
+            # print(type(self.cores[0].numpy()))
+
+            # cores = list(map(lambda core: tf.convert_to_tensor(core), self.cores))
+            # print(type(cores[0]))
+
+            # einsum_structure.append(cores[0])
+            # einsum_structure.append([0, self.cores_number, 2*self.cores_number])
+
+            # for idx in range(1, last_idx):
+            #     einsum_structure.append(cores[idx])
+            #     einsum_structure.append([idx, self.cores_number+idx, 2*self.cores_number+idx])
+
+            # einsum_structure.append(cores[self.cores_number-1])
+            # einsum_structure.append([last_idx, self.cores_number+last_idx, 2*self.cores_number+last_idx-1])
+
+            # # print(list(range(self.cores_number, 2*self.cores_number)))
+            # einsum_structure.append(list(range(self.cores_number+1, 2*self.cores_number)))
+            # result = np.einsum(*einsum_structure)
+
+            return result+self.bias
+
+        result = tf.vectorized_map(lambda vec: process(vec, self.bias), inputs)
+        return self.activation(tf.reshape(result, (-1, self.tt_output_shape_unfold)))
+
+class TDense(Layer):
+
     def __init__(self, units, cores_number, bond_dim=2, shape=None) -> None:
         super(TensorDense, self).__init__()
 
@@ -23,36 +142,36 @@ def __init__(self, units, cores_number, bond_dim=2, shape=None) -> None:
 
         self.cores = []
 
-    def build(self, input_shape):
+    def build(self, tt_input_shape):
 
         self.bias = tf.Variable(tf.zeros(shape=self.shape), name="bias", trainable=True)
 
         # self.shape_input = []
 
         self.cores.append(self.add_weight(
-            shape = (input_shape[1], self.shape[0], self.bond_dim,),
+            shape = (tt_input_shape[1], self.shape[0], self.bond_dim,),
             name = "core_1",
             initializer = 'random_normal',
             trainable = True
         ))
-        # self.shape_input.append(input_shape[1])
+        # self.shape_input.append(tt_input_shape[1])
 
         for i in range(1, self.cores_number-1):
             self.cores.append(self.add_weight(
-                shape = (input_shape[1], self.shape[i], self.bond_dim, self.bond_dim,),
+                shape = (tt_input_shape[1], self.shape[i], self.bond_dim, self.bond_dim,),
                 name = "core_"+str(i),
                 initializer = 'random_normal',
                 trainable = True
             ))
-            # self.shape_input.append(input_shape[1])
+            # self.shape_input.append(tt_input_shape[1])
 
         self.cores.append(self.add_weight(
-            shape = (input_shape[1], self.shape[-1], self.bond_dim,),
+            shape = (tt_input_shape[1], self.shape[-1], self.bond_dim,),
             name = "core_"+str(self.cores_number),
             initializer = 'random_normal',
             trainable = True
         ))
-        # self.shape_input.append(input_shape[1])
+        # self.shape_input.append(tt_input_shape[1])
         # self.shape_input = tuple(self.shape_input)
 
     def call(self, inputs):
@@ -63,8 +182,13 @@ def process(input, cores, bias):
             # padding = tf.convert_to_tensor(np.zeros((reduction-unfold.shape[0]), dtype="float32"))
             # input = tf.reshape(tf.concat(values=[input, padding], axis=0), self.shape_input)
 
-            input = [input, input]
-            input = tf.reshape(input, (2,2))
+            # input = [input, input]
+            print("Core shape", self.cores[0].shape)
+            print("Input shape", input.shape)
+            print("Input reshape", (floor(sqrt(input.shape[0])),ceil(sqrt(input.shape[0]))))
+            input = tf.reshape(input, (floor(sqrt(input.shape[0])),ceil(sqrt(input.shape[0]))))
+
+            print(input.shape)
 
             mx = self.cores_number
 
@@ -74,6 +198,7 @@ def process(input, cores, bias):
             links = [[i, -i, "bond"+str(i-1), "bond"+str(i)] for i in range(2, mx)]
 
             # print([list(range(1,mx+1)), [1, -1, "bond"+str(1)], *links, [mx, -mx, "bond"+str(mx-1)]])
+            print( [list(range(1,mx+1)), [1, -1, "bond"+str(1)], *links, [mx, -mx, "bond"+str(mx)]])
 
             result = tn.ncon(
                 tensors = [x.tensor] + cores,