(network, weights)
| 40 | |
| 41 | |
| 42 | def populate_network(network, weights): |
| 43 | # Configure the network layers based on the weights provided. |
| 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() |
| 75 | conv1 = network.add_convolution( |
| 76 | input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b |
| 77 | ) |
| 78 | conv1.stride = (1, 1) |
| 79 | |
| 80 | pool1 = network.add_pooling(input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2)) |
| 81 | pool1.stride = (2, 2) |
| 82 | |
| 83 | conv2_w = weights["conv2.weight"].numpy() |
| 84 | conv2_b = weights["conv2.bias"].numpy() |
| 85 | conv2 = network.add_convolution(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b) |
| 86 | conv2.stride = (1, 1) |
| 87 | |
| 88 | pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2)) |
| 89 | pool2.stride = (2, 2) |
| 90 | |
| 91 | fc1_w = weights["fc1.weight"].numpy() |
| 92 | fc1_b = weights["fc1.bias"].numpy() |
| 93 | fc1 = add_matmul_as_fc(network, pool2.get_output(0), 500, fc1_w, fc1_b) |
| 94 | |
| 95 | relu1 = network.add_activation(input=fc1.get_output(0), type=trt.ActivationType.RELU) |
| 96 | |
| 97 | fc2_w = weights["fc2.weight"].numpy() |
| 98 | fc2_b = weights["fc2.bias"].numpy() |
| 99 | fc2 = add_matmul_as_fc(network, relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b) |
no test coverage detected