兼容PyTorch，25倍加速，國產(chǎn)！

import torchvision.models as models_torchimport flowvision.models as models_flow
resnet101_torch = models_torch.resnet101(pretrained=True)resnet101_flow = models_flow.resnet101()
state_dict_torch = resnet101_torch.state_dict()state_dict_numpy = {key: value.detach().cpu().numpy() for key, value in state_dict_torch.items()}
resnet101_flow.load_state_dict(state_dict_numpy)

class?ResNet101Graph(oneflow.nn.Graph):    def __init__(self, input_shape, input_dtype=oneflow.float32):        super().__init__()        # 添加 ResNet101 nn.Module        self.model = ResNet101Module(input_shape, input_dtype)        self.loss_fn = ResNet101_loss_fn        # 添加 對應(yīng)的 Optimizer        of_sgd = torch.optim.SGD(self.model.parameters(), lr=1.0, momentum=0.0)        self.add_optimizer(of_sgd)        # 配置靜態(tài)圖的自動優(yōu)化選項        _config_graph(self)     def build(self, input):        # 類似 nn.Module 的 forward 方法，這里是構(gòu)圖，包括了構(gòu)建后向圖，所以叫 build        out = self.model(input)        loss = self.loss_fn(out)        # build 里面支持構(gòu)建后向圖        loss.backward()        return loss

resnet101_graph?=?ResNet101Graph((args.batch_size,?3,?img_shape[1],?img_shape[0]))

for?i?in?range(m):????loss?=?resnet101_graph(images)

def?_config_graph(graph):    if args.fp16:        # 打開 nn.Graph 的自動混合精度執(zhí)行        graph.config.enable_amp(True)     if args.conv_try_run:        # 打開 nn.Graph 的卷積的試跑優(yōu)化        graph.config.enable_cudnn_conv_heuristic_search_algo(False)     if args.fuse_add_to_output:        # 打開 nn.Graph 的add算子的融合        graph.config.allow_fuse_add_to_output(True)     if args.fuse_pad_to_conv:        # 打開 nn.Graph 的pad算子的融合????????graph.config.allow_fuse_pad_to_conv(True)

def?_config_graph(graph):    if args.tensorrt:        # 使用 TensorRT 后端執(zhí)行????????graph.config.enable_tensorrt(True)

class?ResNet101InferenceGraph(oneflow.nn.Graph):    def __init__(self):        super().__init__()                self.model = resnet101_graph.model        def build(self, input):        return self.model(input) inference_graph?=?ResNet101InferenceGraph()

unused_output?=?inference_graph(flow.zeros(1,?3,?224,?224))

flow.save(inference_graph,?"model")

docker?run?--rm?--runtime=nvidia?--network=host?-v$(pwd)/model:/models/resnet101/1?\??oneflowinc/oneflow-serving:nightly

docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \??oneflowinc/oneflow-serving:nightly?oneflow-serving?--model-store?/models?--enable-tensorrt?resnet101

curl?-v?localhost:8000/v2/health/ready

#/usr/bin/env python3 import numpy as npimport tritonclient.http as httpclientfrom PIL import Image  triton_client = httpclient.InferenceServerClient(url='127.0.0.1:8000') image?=?Image.open("image.jpg")image = image.resize((224, 224))image = np.asarray(image)image = image / 255image = np.expand_dims(image, axis=0)# Transpose NHWC to NCHWimage = np.transpose(image, axes=[0, 3, 1, 2])image = image.astype(np.float32) input?=?httpclient.InferInput('INPUT_0',?image.shape,?"FP32")input.set_data_from_numpy(image, binary_data=True)output_placeholder = httpclient.InferRequestedOutput('OUTPUT_0', binary_data=True, class_count=1)output?=?triton_client.infer("resnet101",?inputs=[input],?outputs=[output_placeholder]).as_numpy('OUTPUT_0')print(output)

$ python3 triton_client.py[b'3.630257:499']????#?class?id?為?499，值為?3.630257