(net, input, outputs, w, b)
| 44 | input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE) |
| 45 | |
| 46 | def add_matmul_as_fc(net, input, outputs, w, b): |
| 47 | assert len(input.shape) >= 3 |
| 48 | m = 1 if len(input.shape) == 3 else input.shape[0] |
| 49 | k = int(np.prod(input.shape) / m) |
| 50 | assert np.prod(input.shape) == m * k |
| 51 | n = int(w.size / k) |
| 52 | assert w.size == n * k |
| 53 | assert b.size == n |
| 54 | |
| 55 | input_reshape = net.add_shuffle(input) |
| 56 | input_reshape.reshape_dims = trt.Dims2(m, k) |
| 57 | |
| 58 | filter_const = net.add_constant(trt.Dims2(n, k), w) |
| 59 | mm = net.add_matrix_multiply( |
| 60 | input_reshape.get_output(0), |
| 61 | trt.MatrixOperation.NONE, |
| 62 | filter_const.get_output(0), |
| 63 | trt.MatrixOperation.TRANSPOSE, |
| 64 | ) |
| 65 | |
| 66 | bias_const = net.add_constant(trt.Dims2(1, n), b) |
| 67 | bias_add = net.add_elementwise(mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM) |
| 68 | |
| 69 | output_reshape = net.add_shuffle(bias_add.get_output(0)) |
| 70 | output_reshape.reshape_dims = trt.Dims4(m, n, 1, 1) |
| 71 | return output_reshape |
| 72 | |
| 73 | conv1_w = weights["conv1.weight"].numpy() |
| 74 | conv1_b = weights["conv1.bias"].numpy() |
no test coverage detected