import torchvision.models as Model
import torchvision.transforms as TT
import torch.multiprocessing as mp
class baseProcess:
max_size=128
def __init__(self,data_root) -> None:
self.Q=mp.Queue(maxsize=self.max_size)
self.model=Model.resnet18(True)
self.model.eval()
self.image_loader=image_loader
self.transform = TT.Compose([
TT.Resize(224),
TT.CenterCrop(224),
TT.ToTensor(),
TT.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
def update(self):
for i in range(self.image_loader.length()):
img,file_name=self.image_loader.getitem()
data=self.transform(img)
data=torch.unsqueeze(data,0)
with torch.no_grad():
output=self.model(data)
output=torch.softmax(output,dim=-1)
v,id=output[0].topk(1)
name=id2name[id.item()]
# print(name,v)
while self.Q.full():
time.sleep(2)
self.Q.put((img, name))
def start(self):
p = mp.Process(target=self.update, args=())
p.daemon = True
p.start()
return self
when I just use the model on CPU, the output is fine.
when I use self.model.cuda(), the output is always zeros on windows platform, but it is fine on linux platform.