模型验证

  • 示例代码:
  1. import torch
  2. import torchvision
  3. from PIL import Image
  4. from torch import nn
  6. image_path = "image/dog3.jpg"
  7. img = Image.open(image_path)
  8. # print(img)
  9. img = img.convert("RGB")
  11. transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
  12.                                             torchvision.transforms.ToTensor()])
  13. img = transform(img)
  15. class Lcy(nn.Module):
  16.     def __init__(self):
  17.         super(Lcy,self).__init__()
  18.         self.model = nn.Sequential(
  19.             nn.Conv2d(3,32,5,1,2),
  20.             nn.MaxPool2d(2),
  21.             nn.Conv2d(32,32,5,1,2),
  22.             nn.MaxPool2d(2),
  23.             nn.Conv2d(32,64,5,1,2),
  24.             nn.MaxPool2d(2),
  25.             nn.Flatten(),
  26.             nn.Linear(64*4*4,64),
  27.             nn.Linear(64,10)
  28.         )
  30.     def forward(self,input):
  31.         output = self.model(input)
  32.         return output
  34. # model = torch.load("lcy.pth")
  35. model = torch.load("model_dir/l_29_gpu.pth")
  36. # print(model)
  37. # print(img.type)
  38. img = torch.reshape(img,(1,3,32,32))
  39. img = img.cuda()
  40. model.eval()
  41. with torch.no_grad():
  42.     output = model(img)
  43. # print(output)
  44. print(output.argmax(1))
  46. args = ['飞机','汽车','鸟','猫','鹿','狗','青蛙','马','船','卡车']
  47. if output.argmax(1) == 0:
  48.     print("这张图片是{}".format(args[0]))
  49. elif output.argmax(1) == 1:
  50.     print("这张图片是{}".format(args[1]))
  51. elif output.argmax(1) == 2:
  52.     print("这张图片是{}".format(args[2]))
  53. elif output.argmax(1) == 3:
  54.     print("这张图片是{}".format(args[3]))
  55. elif output.argmax(1) == 4:
  56.     print("这张图片是{}".format(args[4]))
  57. elif output.argmax(1) == 5:
  58.     print("这张图片是{}".format(args[5]))
  59. elif output.argmax(1) == 6:
  60.     print("这张图片是{}".format(args[6]))
  61. elif output.argmax(1) == 7:
  62.     print("这张图片是{}".format(args[7]))
  63. elif output.argmax(1) == 8:
  64.     print("这张图片是{}".format(args[8]))
  65. elif output.argmax(1) == 9:
  66.     print("这张图片是{}".format(args[9]))

15.jpg