def __init__(self, backbone, units=768, layer_norm_eps=1E-12, dropout_prob=0.1,
activation='tanh', weight_initializer=None, bias_initializer=None,
use_segmentation=True):
super().__init__()
self.backbone = backbone
self.use_segmentation = use_segmentation
self.start_scores = nn.Dense(1, flatten=False,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer)
self.end_scores = nn.HybridSequential()
self.end_scores.add(nn.Dense(units, flatten=False,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer))
self.end_scores.add(get_activation(activation))
self.end_scores.add(nn.LayerNorm(epsilon=layer_norm_eps))
self.end_scores.add(nn.Dense(1, flatten=False,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer))
self.answerable_scores = nn.HybridSequential()
self.answerable_scores.add(nn.Dense(units, flatten=False,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer))
self.answerable_scores.add(get_activation(activation))
self.answerable_scores.add(nn.Dropout(dropout_prob))
self.answerable_scores.add(nn.Dense(2, flatten=False,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer))
def test_get_activation():
# Here we just test that the scripts are runnable. Should be revised to test for correctness
for act_type in ['leaky', 'identity', 'elu', 'gelu', 'gelu(tanh)', 'gelu(sigmoid)',
'relu', 'sigmoid', 'tanh', 'softrelu', 'softsign']:
act = get_activation(act_type)
act.hybridize()
_ = act(mx.np.random.normal(0, 1, (10, 10)))