mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 04:53:47 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
#!/usr/bin/env python
|
|
import os
|
|
os.environ['USE_TF'] = '0' # prevent transformers from importing tensorflow
|
|
import unittest
|
|
from tinygrad import Tensor
|
|
import numpy as np
|
|
import torch
|
|
|
|
def get_question_samp(bsz, seq_len, vocab_size, seed):
|
|
np.random.seed(seed)
|
|
in_ids = np.random.randint(vocab_size, size=(bsz, seq_len), dtype=np.int32)
|
|
mask = np.random.choice([True, False], size=(bsz, seq_len))
|
|
seg_ids = np.random.randint(2, size=(bsz, seq_len), dtype=np.int32) # type_vocab_size
|
|
return in_ids, mask, seg_ids
|
|
|
|
def set_equal_weights(mdl, torch_mdl):
|
|
from tinygrad.nn.state import get_state_dict
|
|
state, torch_state = get_state_dict(mdl), torch_mdl.state_dict()
|
|
assert len(state) == len(torch_state)
|
|
for k, v in state.items():
|
|
assert k in torch_state
|
|
torch_state[k].copy_(torch.from_numpy(v.numpy()))
|
|
torch_mdl.eval()
|
|
|
|
class TestBert(unittest.TestCase):
|
|
def test_questions(self):
|
|
from extra.models.bert import BertForQuestionAnswering
|
|
from transformers import BertForQuestionAnswering as TorchBertForQuestionAnswering
|
|
from transformers import BertConfig
|
|
|
|
# small
|
|
config = {
|
|
'vocab_size':24, 'hidden_size':2, 'num_hidden_layers':2, 'num_attention_heads':2,
|
|
'intermediate_size':32, 'hidden_dropout_prob':0.1, 'attention_probs_dropout_prob':0.1,
|
|
'max_position_embeddings':512, 'type_vocab_size':2
|
|
}
|
|
|
|
# Create in tinygrad
|
|
Tensor.manual_seed(1337)
|
|
mdl = BertForQuestionAnswering(**config)
|
|
|
|
# Create in torch
|
|
with torch.no_grad():
|
|
torch_mdl = TorchBertForQuestionAnswering(BertConfig(**config))
|
|
|
|
set_equal_weights(mdl, torch_mdl)
|
|
|
|
seeds = (1337, 3141)
|
|
bsz, seq_len = 1, 16
|
|
for seed in seeds:
|
|
in_ids, mask, seg_ids = get_question_samp(bsz, seq_len, config['vocab_size'], seed)
|
|
out = mdl(Tensor(in_ids), Tensor(mask), Tensor(seg_ids))
|
|
torch_out = torch_mdl.forward(torch.from_numpy(in_ids).long(), torch.from_numpy(mask), torch.from_numpy(seg_ids).long())[:2]
|
|
torch_out = torch.cat(torch_out).unsqueeze(2)
|
|
np.testing.assert_allclose(out.numpy(), torch_out.detach().numpy(), atol=5e-4, rtol=5e-4)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|