train_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
valid_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
# 构建MyDataset实例
train_data = DogCatDataset(data_dir=train_dir, transform=train_transform)
valid_data = DogCatDataset(data_dir=valid_dir, transform=valid_transform)
# 构建DataLoder
train_loader = DataLoader(dataset=train_data,
batch_size=BATCH_SIZE,
shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)
# ============================ step 2/5 模型 ============================
net = MyNet(classes=2)
net.initialize_weights()
net.to("cuda")
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
valid_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
path_model = "./3_best.pkl"
net = torch.load(path_model)
# net = torch.load(path_model, map_location="cpu")
test_dir = "./test_data"
test_data = DogCatDataset(data_dir=test_dir, transform=valid_transform)
valid_loader = DataLoader(dataset=test_data, batch_size=1)
for i, data in enumerate(valid_loader):
# print(valid_loader)
# forward
inputs, labels = data
# print(inputs, labels) #返回的是图片的像素矩阵 和label
outputs = net(inputs.cuda())
# print(outputs.data)
_, predicted = torch.max(outputs.data, 1) #1代表维度
# print(predicted.cpu().numpy()[0])#转化成numpy取出对应标签
"""
rmb_label = {"dog": 0, "cat": 1}
"""