pytorch模型的多batch推理,这里用的是一个多标签分类模型,所以最后用的不是softmax而是sigmod。
import cv2
import time
import models
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
# TODO 分类训练的标签类型,顺序很重要,一定和模型训练时保持一致!!!
CLASSNAMES = ['1', '2’, ‘3', '4']
class TestDataset(Dataset):
def __init__(self, imgpath):
self.imgpath = imgpath
self.image_names = list(os.listdir(imgpath))
def __len__(self):
return len(self.image_names)
def __getitem__(self, index):
path = f"{self.imgpath}/{self.image_names[index]}"
image = Image.open(path)
transform = transforms.Compose([
# transforms.ToPILImage(),
transforms.Resize((256, 256)),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = transform(image)
return {
"id": torch.tensor(index, dtype=torch.float32),
'image': torch.tensor(image, dtype=torch.float32),
}
if __name__ == "__main__":
start_time = time.time()
# # initialize the computation device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#intialize the model
model = models.model(pretrained=False, requires_grad=False).to(device)
# load the model checkpoint
checkpoint = torch.load('/workspace/code/multi-label_image_classification/outputs/model_1117.pth')
# load model weights state_dict
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
'''多batch推理'''
batchsize = 8
numworkers = 4
inputpath = "/workspace/code//input/Images"
test_data = TestDataset(inputpath)
test_loader = DataLoader(test_data, batch_size=batchsize, shuffle=False, sampler=None, batch_sampler=None, num_workers=numworkers)
print(len(test_data), len(test_loader))
out_path = "/workspace/code/multi-label_image_classification/resimg1"
os.makedirs(out_path, exist_ok=True)
# 分类结果txt
txt = open("resimg_batch.txt", "w")
for counter, data in enumerate(test_loader):
id = data['id'].to(device)
image = data['image'].to(device)
outputs = model(image)
for i, output in enumerate(outputs):
# imgname = test_data.image_names[counter*batchsize+i]
imgname = test_data.image_names[int(id[i])] # int(id[i]) == counter*batchsize+i
output = torch.sigmoid(output)
output = output.detach().cpu()
end_time = time.time()
print("cost times:", end_time - start_time)