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)
【模型推理】加载pth文件进行模型推理
猜你喜欢
转载自blog.csdn.net/qq_44992785/article/details/129181775
今日推荐
周排行