save 한 모델 적용해서 이미지 분류값 예측하기!(pth 파일)

베스트 결과가 나온 모델을 저장해놓고 사용하는 법을 간략히 소개합니다!

 

Resnet18 예시)

    model = models.resnet18(pretrained=True)
    model.load_state_dict(torch.load('base_model.pth', map_location=torch.device('cpu')))

    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.5),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])

먼저 어떤 모델을 사용할거다 알려주고, 저장한 파일 값을 불러옵니다.

전처리가 필요하다면 transform을 이용해서 이미지 데이터를 전처리 합니다.

 

    image = Image.open(img)
    image = transform(img).unsqueeze(0)

    model.eval()
    with torch.no_grad():
        output = model(image)
    pred_index = output.argmax().item()
    return pred_index

그리고 나서 사용할 이미지 파일을 열고, transform을 통해 전처리 한 이후에

model.eval을 통해서  evaluation모드로 지정하고 모델에 이미지를 넣어줍니다.

나온 결과값 중, argmax 를 통해서 어떤 인덱스값을 결과로 주는지 받아오면 됩니다.