update

Author: SAAS R7 User1

xnn/base.py

@@ -98,8 +98,13 @@ def call(self, inputs, training=False):
 
         return output
 
+    @tf.function
     def predict(self, x):
-        return self.apply(tf.cast(x, tf.float32), training=False).numpy()
+        return self.apply(tf.cast(x, tf.float32), training=False)
+    
+    @tf.function
+    def evaluate(self, x, y, training=False):
+        return self.loss_fn(y, self.apply(tf.cast(x, tf.float32), training=training))
 
     @tf.function
     def train_step_init(self, inputs, labels):
@@ -113,13 +118,12 @@ def get_active_subnets(self):
         if self.bn_flag:
             beta = self.output_layer.output_weights.numpy() * self.output_layer.subnet_swicher.numpy()
         else:
-            subnet_norm = [self.subnet_blocks.subnets[i].subnet_bn.moving_variance.numpy()[0] ** 0.5 for i in range(self.subnet_num)]
+            subnet_norm = [self.subnet_blocks.subnets[i].subnet_norm.numpy()[0] for i in range(self.subnet_num)]
             beta = self.output_layer.output_weights.numpy() * np.array([subnet_norm]).reshape([-1, 1]) * self.output_layer.subnet_swicher.numpy()
 
         subnets_scale = (np.abs(beta) / np.sum(np.abs(beta))).reshape([-1])
         sorted_index = np.argsort(subnets_scale)
         active_index = sorted_index[subnets_scale[sorted_index].cumsum()>self.beta_threshold][::-1]
-        # active_index = sorted_index[subnets_scale[sorted_index]>self.beta_threshold][::-1]
         return active_index, beta, subnets_scale
 
     def fit(self, train_x, train_y):
@@ -150,8 +154,8 @@ def fit(self, train_x, train_y):
                 batch_yy = tr_y[offset:(offset + self.batch_size)]
                 self.train_step_init(tf.cast(batch_xx, tf.float32), batch_yy)
 
-            self.err_train.append(self.loss_fn(tr_y, self.apply(tf.cast(tr_x, tf.float32), training=True)).numpy())
-            self.err_val.append(self.loss_fn(val_y, self.apply(tf.cast(val_x, tf.float32), training=True)).numpy())
+            self.err_train.append(self.evaluate(tr_x, tr_y, training=True))
+            self.err_val.append(self.evaluate(val_x, val_y, training=True))
             if self.verbose & (epoch % 1 == 0):
                 print("Training epoch: %d, train loss: %0.5f, val loss: %0.5f" %
                       (epoch + 1, self.err_train[-1], self.err_val[-1]))
@@ -188,12 +192,13 @@ def fit(self, train_x, train_y):
                 batch_yy = tr_y[offset:(offset + self.batch_size)]
                 self.train_step_finetune(tf.cast(batch_xx, tf.float32), batch_yy)
 
-            self.err_train.append(self.loss_fn(tr_y, self.apply(tf.cast(tr_x, tf.float32), training=True)).numpy())
-            self.err_val.append(self.loss_fn(val_y, self.apply(tf.cast(val_x, tf.float32), training=True)).numpy())
+            self.err_train.append(self.evaluate(tr_x, tr_y, training=True))
+            self.err_val.append(self.evaluate(val_x, val_y, training=True))
             if self.verbose & (epoch % 1 == 0):
                 print("Tuning epoch: %d, train loss: %0.5f, val loss: %0.5f" %
                       (epoch + 1, self.err_train[-1], self.err_val[-1]))
 
+        self.evaluate(train_x, train_y, training=True)
         # record the key values in the network
         self.subnet_input_min = []
         self.subnet_input_max = []
@@ -240,9 +245,10 @@ def visualize(self, folder="./results", name="demo", dummy_name=None, save_eps=F
                 np.min(subnets_outputs), np.max(subnets_outputs), 6), 2)
             ax1.set_yticks(yint)
             ax1.set_yticklabels(["{0: .2f}".format(j) for j in yint])
-            legend_style = mlines.Line2D([], [], color='black', marker='o', linewidth=0.0, markersize=6,
-                                         label='Scale: ' + str(np.round(100 * subnets_scale[indice], 1)) + "%")
-            plt.legend(handles=[legend_style], fontsize=18)
+            ax1.set_ylim([np.min(subnets_outputs) - (np.max(subnets_outputs) - np.min(subnets_outputs))*0.1, 
+                      np.max(subnets_outputs) + (np.max(subnets_outputs) - np.min(subnets_outputs))*0.25])
+            ax1.text(0.25, 0.9,'Scale: ' + str(np.round(100 * subnets_scale[indice], 1)) + "%",
+                  fontsize=16,  horizontalalignment='center', verticalalignment='center', transform=ax1.transAxes)
 
             ax2 = f.add_subplot(np.int(max_ids), 2, i * 2 + 2)
             ax2.bar(np.arange(input_size), coef_index.T[indice, :input_size])

xnn/gamnet.py

@@ -70,7 +70,7 @@ def __init__(self, input_num, input_dummy_num=0, subnet_arch=[10, 6], task="Regr
                                      subnet_arch=subnet_arch,
                                      task=task,
                                      proj_method="gam",
-                                     activation_func=tf.tanh,
+                                     activation_func=activation_func,
                                      bn_flag=True,
                                      lr_bp=lr_bp,
                                      l1_proj=0,

xnn/layers.py

@@ -51,6 +51,7 @@ def build(self, input_shape=None):
                                             initializer=self.kernel_iniializer,
                                             trainable=self.trainable,
                                             regularizer=tf.keras.regularizers.l1(self.l1_proj))
+        self.built = True
 
     def call(self, inputs, training=False):
         output = tf.matmul(inputs, self.proj_weights)
@@ -59,19 +60,22 @@ def call(self, inputs, training=False):
 
 class Subnetwork(tf.keras.layers.Layer):
 
-    def __init__(self, subnet_arch=[10, 6], activation_func=tf.tanh, smooth_lambda=0.0, bn_flag=False):
+    def __init__(self, subnet_arch=[10, 6], activation_func=tf.tanh, smooth_lambda=0.0, bn_flag=False, subnet_id=0):
         super(Subnetwork, self).__init__()
         self.dense = []
         self.subnet_arch = subnet_arch
         self.activation_func = activation_func
         self.smooth_lambda = smooth_lambda
         self.bn_flag = bn_flag
+        self.subnet_id = subnet_id
 
     def build(self, input_shape=None):
         for nodes in self.subnet_arch:
             self.dense.append(layers.Dense(nodes, activation=self.activation_func))
-        self.output_layer = layers.Dense(1, activation=tf.identity)
-        self.subnet_bn = BatchNormalization(momentum=0.0, epsilon=1e-10, center=False, scale=False)
+        self.output_layer = layers.Dense(1, activation=self.activation_func)
+        self.moving_mean = self.add_weight(name="mean"+str(self.subnet_id), shape=[1], initializer=tf.zeros_initializer(),trainable=False)
+        self.moving_norm = self.add_weight(name="norm"+str(self.subnet_id), shape=[1], initializer=tf.ones_initializer(),trainable=False)
+        self.built = True      
 
     def call(self, inputs, training=False):
         with tf.GradientTape() as t1:
@@ -85,12 +89,21 @@ def call(self, inputs, training=False):
             self.grad1 = t2.gradient(self.output_original, inputs)
         self.grad2 = t1.gradient(self.grad1, inputs)
 
+        if training:
+            mean, norm = tf.reduce_mean(self.output_original, 0), tf.maximum(tf.math.reduce_std(self.output_original, 0), 1e-10) 
+            self.subnet_mean = mean
+            self.subnet_norm = norm
+            self.moving_mean.assign(mean)
+            self.moving_norm.assign(norm)
+        else:
+            self.subnet_mean = self.moving_mean
+            self.subnet_norm = self.moving_norm
+
         if self.bn_flag:
-            output = self.subnet_bn(self.output_original, training=training)
+            output = (self.output_original - self.subnet_mean) / (self.subnet_norm)
         else:
-            _ = self.subnet_bn(self.output_original, training=training)
             output = self.output_original
-        self.smooth_penalty = tf.reduce_mean(tf.square(self.grad2)) / tf.sqrt(self.subnet_bn.moving_variance)
+        self.smooth_penalty = tf.reduce_mean(tf.square(self.grad2)) / self.subnet_norm
         return output
 
 
@@ -110,7 +123,8 @@ def build(self, input_shape=None):
             self.subnets.append(Subnetwork(self.subnet_arch,
                                            self.activation_func,
                                            self.smooth_lambda,
-                                           self.bn_flag))
+                                           self.bn_flag,
+                                           subnet_id=i))
         self.built = True
 
     def call(self, inputs, training=False):
@@ -149,6 +163,7 @@ def build(self, input_shape=None):
                                            shape=[1],
                                            initializer=tf.zeros_initializer(),
                                            trainable=True)
+        self.built = True
 
     def call(self, inputs, training=False):
         output = (tf.matmul(inputs, self.subnet_swicher * self.output_weights) + self.output_bias)

xnn/sosxnn.py

@@ -91,7 +91,7 @@ def __init__(self, input_num, subnet_num, input_dummy_num=0, subnet_arch=[10, 6]
                                      subnet_arch=subnet_arch,
                                      task=task,
                                      proj_method="orthogonal",
-                                     activation_func=tf.tanh,
+                                     activation_func=activation_func,
                                      bn_flag=True,
                                      lr_bp=lr_bp,
                                      l1_proj=l1_proj,

xnn/xnn.py

@@ -85,7 +85,7 @@ def __init__(self, input_num, subnet_num, input_dummy_num=0, subnet_arch=[10, 6]
                                   subnet_arch=subnet_arch,
                                   task=task,
                                   proj_method="comb",
-                                  activation_func=tf.tanh,
+                                  activation_func=activation_func,
                                   bn_flag=False,
                                   lr_bp=lr_bp,
                                   l1_proj=l1_proj,