(
network: trt.INetworkDefinition,
weights: OrderedDict,
input: trt.ITensor,
out_channel: int,
ksize: int,
stride: int,
group: int,
layer_name: str,
)
| 44 | |
| 45 | |
| 46 | def Conv( |
| 47 | network: trt.INetworkDefinition, |
| 48 | weights: OrderedDict, |
| 49 | input: trt.ITensor, |
| 50 | out_channel: int, |
| 51 | ksize: int, |
| 52 | stride: int, |
| 53 | group: int, |
| 54 | layer_name: str, |
| 55 | ) -> trt.ILayer: |
| 56 | padding = ksize // 2 |
| 57 | if ksize > 3: |
| 58 | padding -= 1 |
| 59 | conv_w = trtweight(weights[layer_name + ".conv.weight"]) |
| 60 | conv_b = trtweight(weights[layer_name + ".conv.bias"]) |
| 61 | |
| 62 | conv = network.add_convolution_nd( |
| 63 | input, num_output_maps=out_channel, kernel_shape=trt.DimsHW(ksize, ksize), kernel=conv_w, bias=conv_b |
| 64 | ) |
| 65 | assert conv, "Add convolution_nd layer failed" |
| 66 | conv.stride_nd = trt.DimsHW(stride, stride) |
| 67 | conv.padding_nd = trt.DimsHW(padding, padding) |
| 68 | conv.num_groups = group |
| 69 | |
| 70 | sigmoid = network.add_activation(conv.get_output(0), trt.ActivationType.SIGMOID) |
| 71 | assert sigmoid, "Add activation layer failed" |
| 72 | dot_product = network.add_elementwise(conv.get_output(0), sigmoid.get_output(0), trt.ElementWiseOperation.PROD) |
| 73 | assert dot_product, "Add elementwise layer failed" |
| 74 | return dot_product |
| 75 | |
| 76 | |
| 77 | def Bottleneck( |
no test coverage detected