# 加载模型
model = models.fishnet150()
# 加载模型文件的权重数据到model模型
checkpoint = torch.load('checkpoints/pre-trainedModels/fishnet150_ckpt.tar') # 加载指定的模型文件
model.load_state_dict(checkpoint['state_dict']) # 导入模型中的权重数据
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 设置计算硬件,CPU 或 GPU
torch.cuda.set_device(device)
# torch.backends.cudnn.benchmark = True
model = model.module.to(device)
model.eval()
# 准备transform
img_size = 224
ratio = 224.0 / float(img_size)
normalize = transforms.Normalize(mean=[0.14300402, 0.1434545, 0.14277956],
std=[0.10050353, 0.100842826, 0.10034215])
transform = transforms.Compose([transforms.Resize(int(224 * ratio)), transforms.CenterCrop(img_size),
transforms.ToTensor(), normalize])
# 读取单张人脸RGB图像进行预测
img = Image.open("face1.jpg")
img = img.convert('RGB')
img = img.resize((224, 224))
img = transform(img)
img = np.array(img) # 形状为(3, 224, 224)
img = np.expand_dims(img, 0) # 此时形状为(1, 3, 224, 224)
input = torch.tensor(img, dtype=torch.float32, device=device)
output = model(input) # output形状为[1, 1000]
soft_output = torch.softmax(output, dim=-1) # 对output中最后一个维度上的元素进行softmax处理,得到的形状依然是[1, 1000]
preds = soft_output.to('cpu').detach().numpy() # 转换为numpy数组,形状为(1, 1000)