1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
|
import numpy as np
import unittest
import caffe2.python.fakelowp.init_shared_libs # noqa
from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace
from caffe2.python.onnx.onnxifi import onnxifi_caffe2_net
from caffe2.python.fakelowp.test_utils import print_test_debug_info
import datetime
from hypothesis import given, settings
import hypothesis.strategies as st
import caffe2.python.serialized_test.serialized_test_util as serial
core.GlobalInit(["caffe2", "--caffe2_log_level=-3", "--glow_global_fp16=1"])
class TestBatchMatMul(serial.SerializedTestCase):
@given(
C=st.integers(min_value=1, max_value=10),
M=st.integers(min_value=1, max_value=50),
K=st.integers(min_value=1, max_value=512),
N=st.integers(min_value=1, max_value=50),
rand_seed=st.integers(0, 65534),
trans_a=st.booleans(),
trans_b=st.booleans(),
run_ints=st.booleans()
)
@settings(deadline=datetime.timedelta(seconds=10))
def test_batch_matmul(self, M, K, N, C, rand_seed, trans_a, trans_b, run_ints):
np.random.seed(rand_seed)
workspace.ResetWorkspace()
batch_dims = [C]
if run_ints:
X = np.random.randint(low=1, high=3, size=((C, M, K))).astype(np.float32)
else:
X = 100 * (np.random.rand(*(batch_dims + [M, K])).astype(np.float32) - 0.5)
if trans_a:
X = X.swapaxes(-1, -2)
if run_ints:
Y = np.random.randint(low=1, high=3, size=((C, K, N))).astype(np.float32)
else:
Y = 100 * (np.random.rand(*(batch_dims + [K, N])).astype(np.float32) - 0.5)
if trans_b:
Y = Y.swapaxes(-1, -2)
pred_net = caffe2_pb2.NetDef()
pred_net.name = "pred"
pred_net.external_input.extend(["X", "Y"])
pred_net.external_output.append("out")
pred_net.op.add().CopyFrom(
core.CreateOperator(
'BatchMatMul', ['X', 'Y'], 'out', trans_a=trans_a, trans_b=trans_b
)
)
pred_net_ref = core.Net("pred_net_ref")
# Reference updated to fp16 with fp32 accumulation
pred_net_ref.BatchMatMulFP16Acc32Fake(
["X", "Y"], ['out'], trans_a=trans_a, trans_b=trans_b)
print("dims", batch_dims, X.shape, Y.shape)
pred_net_onnxified = onnxifi_caffe2_net(pred_net,
{"X": X.shape, "Y": Y.shape},
debug=True,
adjust_batch=False,
use_onnx=False)
num_onnxified_ops = sum(
1 if o.type == "Onnxifi" else 0 for o in pred_net_onnxified.op)
np.testing.assert_equal(num_onnxified_ops, 1)
workspace.FeedBlob("X", X)
workspace.FeedBlob("Y", Y)
workspace.CreateNet(pred_net_onnxified)
workspace.CreateNet(pred_net_ref)
# Run Glow net
workspace.RunNet(pred_net_onnxified.name)
out_glow = workspace.FetchBlob('out')
# Run caffe2 net
workspace.RunNet(pred_net_ref)
out_c2_fakefp16 = workspace.FetchBlob('out')
diff = np.abs(out_c2_fakefp16 - out_glow)
if not np.allclose(out_glow, out_c2_fakefp16):
print_test_debug_info("bmm", {
"seed": rand_seed,
"m": M, "k": K,
"n": N, "X": X.shape, "Y": Y.shape,
"trans_a": trans_a,
"trans_b": trans_b,
"run_ints": run_ints,
"out_glow": out_glow,
"out_c2_fakefp16": out_c2_fakefp16,
"diff": diff
})
assert(0)
if __name__ == "__main__":
unittest.main()
|