【模型推理】加载pth文件进行模型推理

import torch
from torchvision import datasets, models, transforms
import cv2
import numpy as np
from PIL import Image
def pth_push(img):
    model=#这里调用你的模型
    model_path=''#pth权重文件地址
    device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')#cpu or gpu
    model.load_state_dict(torch.load(self.model_path, map_location=device))#加载pth文件
    model  = model.eval()
    transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
])#对图片进行resize并转换成tensor
    inputs = transform(img)
    inputs=torch.unsqueeze(inputs, 0)#添加一个维度
    inputs = inputs.to(device)#把图片也转成相应的设备cuda or cpu
    #进行推理
    outputs = model(inputs)
    #根据自己要解决的问题进行解码
    outputs1=outputs.tolist()
    outputs1 = torch.from_numpy(np.array(outputs1))
    outputs_softmax = torch.softmax(outputs1, dim=1).numpy()[:, 1].tolist()[0]
if __name__ == "__main__":  
    i='.jpg' 
    image = Image.open(i)
    pth_push(image)

猜你喜欢

转载自blog.csdn.net/qq_44992785/article/details/129181775