(self, num_hidden, num_classes)
| 386 | """A subclass model based small MLP.""" |
| 387 | |
| 388 | def __init__(self, num_hidden, num_classes): |
| 389 | super(_SmallSubclassMLP, self).__init__() |
| 390 | self.layer_a = keras.layers.Dense(num_hidden, activation='relu') |
| 391 | activation = 'sigmoid' if num_classes == 1 else 'softmax' |
| 392 | self.layer_b = keras.layers.Dense(num_classes, activation=activation) |
| 393 | |
| 394 | def call(self, inputs, **kwargs): |
| 395 | x = self.layer_a(inputs) |