def test_list_backbone_names(): assert len(list_backbone_names()) > 0 def tvm_enabled(): try: tvm = try_import_tvm() return True except: return False @pytest.mark.slow @pytest.mark.parametrize('name', list_backbone_names()) def test_get_backbone(name, ctx): with tempfile.TemporaryDirectory() as root, ctx: model_cls, cfg, tokenizer, local_params_path, _ = get_backbone( name, root=root) net = model_cls.from_cfg(cfg) net.load_parameters(local_params_path) net.hybridize() num_params, num_fixed_params = count_parameters(net.collect_params()) assert num_params > 0 # Test for model export + save if 'gpt2' in name: pytest.skip('Skipping GPT-2 test') batch_size = 1 sequence_length = 4
def test_list_backbone_names(): assert len(list_backbone_names()) > 0
import json import mxnet as mx from gluonnlp.models import list_backbone_names, get_backbone mx.npx.set_np() batch_size = 1 sequence_length = 32 all_possible_ops = [] for name in list_backbone_names(): model_cls, cfg, tokenizer, local_params_path, others = get_backbone( model_name=name) net = model_cls.from_cfg(cfg) net.initialize() net.hybridize() print('Save the architecture of {} to {}.json'.format(name, name)) inputs = mx.np.random.randint(0, 10, (batch_size, sequence_length)) token_types = mx.np.random.randint(0, 2, (batch_size, sequence_length)) valid_length = mx.np.random.randint(1, 10, (batch_size, )) if 'roberta' in name or 'xlmr' in name: out = net(inputs, valid_length) else: out = net(inputs, token_types, valid_length) sym = net._cached_graph[1] sym.save('{}.json'.format(name), remove_amp_cast=True) all_ops = set() with open('{}.json'.format(name), 'r') as f: sym_info = json.load(f) for ele in sym_info['nodes']: all_ops.add(ele['op']) with open('{}_all_ops.json'.format(name), 'w') as f: json.dump(list(all_ops), f)