predict文件篇

predict文件只包含一个predict.py文件，他说到底其实就是将train文件篇的微小版，将dataset.py和train.py融合在了一起，加载图片后进行predict。

x.1 predict.py代码迭代

x.1.1 AlexNet和5种花分类数据集

20230403：add os.environ[“CUDA_VISIBLE_DEVICES”]

# --- add path
import os, sys
project_path = os.path.dirname(os.path.dirname(__file__))
root_path = os.path.dirname(project_path)
sys.path.append(project_path)
# ---
import torch, os
from model.model import AlexNet
from torchvision.transforms import transforms
from PIL import Image
# 指定设备
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 会将实际上0, 2显卡设定为可见：即在该脚本中0 -> 0, 2 -> 1

def main():
    # parse some arguments
    classes = 5
    weight_path = "/home/yingmuzhi/_learning/pytorch/_pipeline/_example_AlexNet/data/output/model/model.pth"
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    # read img, remember the signal need to expand dimension(torch.unsqueeze) and do the same transform in train. You should alse put the image into GPU.
    signal = "/home/yingmuzhi/_learning/pytorch/_pipeline/_example_AlexNet/data/origin/flower_data/val/daisy/105806915_a9c13e2106_n.jpg"
    signal_transforms = transforms.Compose([
            transforms.Resize(256),
            transforms.RandomCrop((224, 224)),
            transforms.ToTensor(),
    ])
    signal = signal_transforms(Image.open(signal))
    signal = signal.unsqueeze(0)
    signal = signal.to(device)

    target = "0"
    
    # load model
    model = AlexNet(num_classes=classes)
    model.load_state_dict(torch.load(weight_path, map_location = "cpu")["model"])
    model = model.to(device)
    
    # test
    model.eval()
    with torch.no_grad():
        prediction = model(signal)
        prediction = torch.argmax(prediction[0])
        print("predition is {}, target is {}".format(target, prediction.item()))
        pass


if __name__=="__main__":
    main()

20230214：origin

# --- add path
import os, sys
project_path = os.path.dirname(os.path.dirname(__file__))
root_path = os.path.dirname(project_path)
sys.path.append(project_path)
# ---
import torch, os
from model.model import AlexNet
from torchvision.transforms import transforms
from PIL import Image

def main():
    # parse some arguments
    classes = 5
    weight_path = "/home/yingmuzhi/_learning/pytorch/_pipeline/_example_AlexNet/data/output/model/model.pth"
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    # read img, remember the signal need to expand dimension(torch.unsqueeze) and do the same transform in train. You should alse put the image into GPU.
    signal = "/home/yingmuzhi/_learning/pytorch/_pipeline/_example_AlexNet/data/origin/flower_data/val/daisy/105806915_a9c13e2106_n.jpg"
    signal_transforms = transforms.Compose([
            transforms.Resize(256),
            transforms.RandomCrop((224, 224)),
            transforms.ToTensor(),
    ])
    signal = signal_transforms(Image.open(signal))
    signal = signal.unsqueeze(0)
    signal = signal.to(device)

    target = "0"
    
    # load model
    model = AlexNet(num_classes=classes)
    model.load_state_dict(torch.load(weight_path, map_location = "cpu")["model"])
    model = model.to(device)
    
    # test
    model.eval()
    with torch.no_grad():
        prediction = model(signal)
        prediction = torch.argmax(prediction[0])
        print("predition is {}, target is {}".format(target, prediction.item()))
        pass


if __name__=="__main__":
    main()