Python tensorrt.float32方法代码示例

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def __call__(self, img):
        h, w = img.size(1), img.size(2)
        mask = np.ones((h, w), np.float32)
        y = np.random.randint(h)
        x = np.random.randint(w)

        y1 = np.clip(y - self.length // 2, 0, h)
        y2 = np.clip(y + self.length // 2, 0, h)
        x1 = np.clip(x - self.length // 2, 0, w)
        x2 = np.clip(x + self.length // 2, 0, w)

        mask[y1: y2, x1: x2] = 0.
        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img *= mask
        return img 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def __init__(self,
                 max_batchsize,
                 workspace,
                 dtype=trt.float32,
                 builder_config_fn=None,
                 net_post_fn=None,
                 input_names=None,
                 verbose=False):
        super().__init__()
        self.max_batchsize = max_batchsize
        self.workspace = workspace
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.built = False
        self.graph_pth = None
        self.refit_weight_dict = {}
        self.engine = None
        self.ctx = None
        self.output_shapes = None
        self.output_names = None
        self.need_refit = False
        self.verbose = verbose
        self.builder_config_fn = builder_config_fn
        self.input_names = input_names
        self.net_post_fn = net_post_fn 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def __call__(self, *args, **kw):
        if not self.training and not self.built:
            self.build_tensorrt(self.net, args)
            self.built = True
            self.need_refit = False
        if not self.training:
            if self.need_refit:
                self.refit_engine(self.net)
                self.need_refit = False
            assert all([a.is_cuda for a in args])
            torch.cuda.synchronize()
            # args = [a.detach().cpu().numpy() for a in args]
            output_dict = self.ctx.inference_async(*args)
            # for k,v in output_dict.items():
            #     output_dict[k] = torch.tensor(v, dtype=torch.float32, device=torch.device("cuda:0"))
            outputs = [None] * len(output_dict)
            for k, v in output_dict.items():
                outputs[self.output_names.index(k)] = v.view(
                    v.shape[0], *self.output_shapes[k])
            if len(outputs) == 1:
                return outputs[0]
            return tuple(outputs)
        else:
            self.need_refit = True
            return super().__call__(*args, **kw) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def torch_dtype_from_trt(dtype):
    if dtype == trt.int8:
        return torch.int8
    elif dtype == trt.int32:
        return torch.int32
    elif dtype == trt.float16:
        return torch.float16
    elif dtype == trt.float32:
        return torch.float32
    else:
        raise TypeError('%s is not supported by torch' % dtype) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def populate_network(network, weights):
    # Configure the network layers based on the weights provided.
    input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)

    # Set dummy weights for the kernel and bias weights in the conv1 layer. We
    # will refit the engine with the actual weights later.
    conv1_w = np.zeros((20,5,5), dtype=np.float32)
    conv1_b = np.zeros(20, dtype=np.float32)

    conv1 = network.add_convolution(input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b)
    conv1.name = "conv_1"
    conv1.stride = (1, 1)

    pool1 = network.add_pooling(input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2))
    pool1.stride = (2, 2)

    conv2_w = weights['conv2.weight'].numpy()
    conv2_b = weights['conv2.bias'].numpy()
    conv2 = network.add_convolution(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b)
    conv2.stride = (1, 1)

    pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
    pool2.stride = (2, 2)

    fc1_w = weights['fc1.weight'].numpy()
    fc1_b = weights['fc1.bias'].numpy()
    fc1 = network.add_fully_connected(input=pool2.get_output(0), num_outputs=500, kernel=fc1_w, bias=fc1_b)

    relu1 = network.add_activation(input=fc1.get_output(0), type=trt.ActivationType.RELU)

    fc2_w = weights['fc2.weight'].numpy()
    fc2_b = weights['fc2.bias'].numpy()
    fc2 = network.add_fully_connected(relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b)

    fc2.get_output(0).name = ModelData.OUTPUT_NAME
    network.mark_output(tensor=fc2.get_output(0)) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def torch_dtype_to_trt(dtype):
    if dtype == torch.int8:
        return trt.int8
    elif dtype == torch.int32:
        return trt.int32
    elif dtype == torch.float16:
        return trt.float16
    elif dtype == torch.float32:
        return trt.float32
    else:
        raise TypeError('%s is not supported by tensorrt' % dtype) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def convert_caffe_model_to_trt(caffe_weights_file, caffe_deploy_file, trt_model_filename,
                               output_tensor_name, output_data_type, max_workspace_size, max_batch_size):
    "Convert a pair of (caffe_weights_file,caffe_deploy_file) into a trt_model_file using the given parameters"

    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:

        if (output_data_type=='fp16'):
            if not builder.platform_has_fast_fp16:
                print('Warning: This platform is not optimized for fast fp16 mode')

            builder.fp16_mode = True
            print('Converting into fp16, max_batch_size={}'.format(max_batch_size))
        else:
            print('Converting into fp32 (default), max_batch_size={}'.format(max_batch_size))

        builder.max_workspace_size  = max_workspace_size
        builder.max_batch_size      = max_batch_size

        model_tensors       = parser.parse(deploy=caffe_deploy_file, model=caffe_weights_file, network=network, dtype=trt.float32)
        network.mark_output(model_tensors.find(output_tensor_name))

        trt_model_object    = builder.build_cuda_engine(network)

        try:
            serialized_trt_model = trt_model_object.serialize()
            with open(trt_model_filename, "wb") as trt_model_file:
                trt_model_file.write(serialized_trt_model)
        except:
            print('Error: cannot serialize or write TensorRT engine to file {}.'.format(trt_model_filename)) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def torch_dtype_to_trt(dtype):
    if trt_version() >= '7.0' and dtype == torch.bool:
        return trt.bool
    elif dtype == torch.int8:
        return trt.int8
    elif dtype == torch.int32:
        return trt.int32
    elif dtype == torch.float16:
        return trt.float16
    elif dtype == torch.float32:
        return trt.float32
    else:
        raise TypeError("%s is not supported by tensorrt" % dtype) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def torch_dtype_from_trt(dtype):
    if dtype == trt.int8:
        return torch.int8
    elif trt_version() >= '7.0' and dtype == trt.bool:
        return torch.bool
    elif dtype == trt.int32:
        return torch.int32
    elif dtype == trt.float16:
        return torch.float16
    elif dtype == trt.float32:
        return torch.float32
    else:
        raise TypeError("%s is not supported by torch" % dtype) 

# 需要导入模块: import tensorrt [as 别名]
# 或者: from tensorrt import float32 [as 别名]
def build_tensorrt(self, net, torch_inputs):
        if self.input_names is None:
            input_names = get_torch_forward_name(net.forward)
        else:
            input_names = self.input_names
        self.graph_pth = torch2trt.GraphModule(net, torch_inputs)
        self.output_names = []
        with trt.Builder(
                self.logger) as builder, builder.create_network() as trt_net:
            builder.max_workspace_size = self.workspace
            builder.max_batch_size = self.max_batchsize
            builder.fp16_mode = builder.platform_has_fast_fp16
            # builder.refittable = False
            if self.builder_config_fn is not None:
                self.builder_config_fn(builder)
            with torch2trt.trt_network(trt_net):
                inputs = []
                for i, arg in enumerate(torch_inputs):
                    name = input_names[i]
                    inp = trt_net.add_input(name=name,
                                            shape=arg.shape[1:],
                                            dtype=trt.float32)
                    inputs.append(inp)
                outputs = self.graph_pth(*inputs, verbose=self.verbose)
            self.refit_weight_dict = self.graph_pth.graph.refit_weight_dict
            if not isinstance(outputs, (list, tuple)):
                outputs = [outputs]
            for i, out in enumerate(outputs):
                name = "output{}".format(i)
                out.name = name
                self.output_names.append(name)
                trt_net.mark_output(tensor=out)
            self.builder = builder
            if self.net_post_fn is not None:
                self.net_post_fn(trt_net)
            self.engine = builder.build_cuda_engine(trt_net)
            self.ctx = self.engine.create_execution_context()
            self.ctx = torch2trt.TorchInferenceContext(self.ctx)
        # get output shapes
        outputs = self.graph_pth(*torch_inputs)
        if not isinstance(outputs, (list, tuple)):
            outputs = [outputs]
        self.output_shapes = {}
        for n, v in zip(self.output_names, outputs):
            self.output_shapes[n] = v.shape[1:]