| 24 | |
| 25 | |
| 26 | def dump_static_graph(model, graph_name, shape): |
| 27 | model.eval() |
| 28 | |
| 29 | data = mge.Tensor(np.ones(shape, dtype=np.uint8)) |
| 30 | |
| 31 | @jit.trace(capture_as_const=True) |
| 32 | def pred_func(data): |
| 33 | out = data.astype(np.float32) |
| 34 | |
| 35 | output_h, output_w = 224, 224 |
| 36 | # resize |
| 37 | print(shape) |
| 38 | M = mge.tensor(np.array([[1,0,0], [0,1,0], [0,0,1]], dtype=np.float32)) |
| 39 | M_shape = F.concat([data.shape[0],M.shape]) |
| 40 | M = F.broadcast_to(M, M_shape) |
| 41 | out = F.vision.warp_perspective(out, M, (output_h, output_w), format='NHWC') |
| 42 | # mean |
| 43 | _mean = mge.Tensor(np.array([103.530, 116.280, 123.675], dtype=np.float32)) |
| 44 | out = F.sub(out, _mean) |
| 45 | # div |
| 46 | _div = mge.Tensor(np.array([57.375, 57.120, 58.395], dtype=np.float32)) |
| 47 | out = F.div(out, _div) |
| 48 | # dimshuffile |
| 49 | out = F.transpose(out, (0,3,1,2)) |
| 50 | |
| 51 | outputs = model(out) |
| 52 | return outputs |
| 53 | |
| 54 | pred_func(data) |
| 55 | pred_func.dump( |
| 56 | graph_name, |
| 57 | arg_names=["data"], |
| 58 | optimize_for_inference=True, |
| 59 | enable_fuse_conv_bias_nonlinearity=True, |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def main(): |