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} """