深入了解 TabNet ：架构详解和分类代码实现

来源：Deephub Imba本文约3500字，建议阅读5分钟本文我们将深入研究称为 TabNet (Arik & Pfister (2019)) 的神经网络架构，该架构旨在可解释并与表格数据很好地配合使用。

def glu(x, n_units=None):
  """Generalized linear unit nonlinear activation."""
  return x[:, :n_units] * tf.nn.sigmoid(x[:, n_units:])

class FeatureBlock(tf.keras.Model):
  """
  Implementation of a FL->BN->GLU block
  """
  def __init__(
      self,
      feature_dim,
      apply_glu = True,
      bn_momentum = 0.9,
      fc = None,
      epsilon = 1e-5,
  ):
      super(FeatureBlock, self).__init__()
      self.apply_gpu = apply_glu
      self.feature_dim = feature_dim
      units = feature_dim * 2 if apply_glu else feature_dim # desired dimension gets multiplied by 2
                                                            # because GLU activation halves it

      self.fc = tf.keras.layers.Dense(units, use_bias=False) if fc is None else fc # shared layers can get re-used
      self.bn = tf.keras.layers.BatchNormalization(momentum=bn_momentum, epsilon=epsilon)

  def call(self, x, training = None):
      x = self.fc(x) # inputs passes through the FC layer
      x = self.bn(x, training=training) # FC layer output gets passed through the BN
      if self.apply_gpu:
          return glu(x, self.feature_dim) # GLU activation applied to BN output
      return x

class FeatureTransformer(tf.keras.Model):
  def __init__(
      self,
      feature_dim,
      fcs = [],
      n_total = 4,
      n_shared = 2,
      bn_momentum = 0.9,
  ):
      super(FeatureTransformer, self).__init__()
      self.n_total, self.n_shared = n_total, n_shared

      kwrgs = {
          "feature_dim": feature_dim,
          "bn_momentum": bn_momentum,
      }

      # build blocks
      self.blocks = []
      for n in range(n_total):
          # some shared blocks
          if fcs and n < len(fcs):
              self.blocks.append(FeatureBlock(**kwrgs, fc=fcs[n])) # Building shared blocks by providing FC layers
          # build new blocks
          else:
              self.blocks.append(FeatureBlock(**kwrgs)) # Step dependent blocks without the shared FC layers

  def call(self, x, training = None):
      # input passes through the first block
      x = self.blocks[0](x, training=training)
      # for the remaining blocks
      for n in range(1, self.n_total):
          # output from previous block gets multiplied by sqrt(0.5) and output of this block gets added
          x = x * tf.sqrt(0.5) + self.blocks[n](x, training=training)
      return x

  @property
  def shared_fcs(self):
      return [self.blocks[i].fc for i in range(self.n_shared)]

class AttentiveTransformer(tf.keras.Model):
  def __init__(self, feature_dim):
      super(AttentiveTransformer, self).__init__()
      self.block = FeatureBlock(
          feature_dim,
          apply_glu=False, # sparsemax instead of glu
      )

  def call(self, x, prior_scales, training=None):
      # Pass input trhough a FC-BN block
      x = self.block(x, training=training)
      # Pass the output through sparsemax activation
      return sparsemax(x * prior_scales)

def sparse_loss(at_mask):
  loss = tf.reduce_mean(
      tf.reduce_sum(tf.multiply(-at_mask, tf.math.log(at_mask + 1e-15)),
                    axis=1)
  )
   
  return loss

not_sparse_mask = np.array([[0.4, 0.5, 0.05, 0.05],
                    [0.2, 0.2, 0.5, 0.1]])

sparse_mask = np.array([[0.0, 0.0, 0.7, 0.3],
                    [0.0, 0.0, 1, 0.0]])

print('Loss for non-sparse attention mask:', sparse_loss(not_sparse_mask).numpy())
print('Loss for sparse attention mask:', sparse_loss(sparse_mask).numpy())

# Loss for non-sparse attention mask: 1.1166351874690217
# Loss for sparse attention mask: 0.3054321510274452

# Params after 1 hour of tuning
tabnet = TabNet(num_features = train_X_transformed.shape[1],
              output_dim = 128,
              feature_dim = 128,
              n_step = 2,
              relaxation_factor= 2.2,
              sparsity_coefficient=2.37e-07,
              n_shared = 2,
              bn_momentum = 0.9245)


# Early stopping based on validation loss    
cbs = [tf.keras.callbacks.EarlyStopping(
      monitor="val_loss", patience=30, restore_best_weights=True
  )]

# Optimiser
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01, clipnorm=10)

# Second loss in None because we also output the importances
loss = [tf.keras.losses.CategoricalCrossentropy(from_logits=False), None]

# Compile the model
tabnet.compile(optimizer,
              loss=loss)

# Train the model
tabnet.fit(train_ds,
          epochs=1000,
          validation_data=val_ds,
          callbacks=cbs,
          verbose=1,
          class_weight={
            0:1,
            1: 10
          })

Test ROC AUC 0.8505
Test PR AUC 0.464