Exemplo n.º 1
0
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")
Exemplo n.º 2
0
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}
    """