pytorch前向inference

inference

import torch
import numpy as np
from PIL import Image
from torchvision import transforms as T
from models import ShuffleNetV2

transforms = T.Compose([
   T.Resize(size=(112, 112)),
   T.ToTensor(),
   T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

##数据按要求读取 
img_path='demo.jpg'
src = Image.open(path)
src = transforms(src)
src = src.unsqueeze(dim=0)


##选择cuda，网络和模型加载
model_path = 'demo.pth'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
model=ShuffleNetV2(scale=0.5, in_channels=3, num_classes=2)
#model= torch.nn.DataParallel(model)
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint ['state_dict'])
model = model.to(device)
model.eval()

##模型前向输出结果
with torch.no_grad():         #1
     out = model.forward(src.to(device))
#or
out = model.forward(src)      #2
out=out