如何使用Skorch预测单个图像?

时间:2019-12-03 15:55:58

标签: python neural-network deep-learning pytorch skorch

我刚刚用Skorch创建了一个神经网络来检测照片上的飞机,并使用形状为(40000, 64, 64, 3)的火车数据集对其进行了训练。

然后我用(15000, 64, 64, 3)的测试数据集进行了测试。

module = nn.Sequential(
    nn.Conv2d(3, 64, 3),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(64, 64, 3),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(64, 64, 3),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(6 * 6 * 64, 256),
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 2),
    nn.Softmax(),
)

early_stopping = EarlyStopping(monitor='valid_loss', lower_is_better=True)
net = NeuralNetClassifier(
    module,
    max_epochs=20,
    lr=1e-4,
    callbacks=[early_stopping],
    # Shuffle training data on each epoch
    iterator_train__shuffle=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
    optimizer=optim.Adam
)
net.fit(
    train_images_balanced.transpose((0, 3, 1, 2)).astype(np.float32),
    train_labels_balanced
)

现在我需要在512 * 512图片上对其进行测试,因此我有了一个(30, 512, 512, 3)的新数据集。
因此,我使用了一个滑动窗口代码,该代码使我可以将图片分成64 * 64的部分。

def sliding_window(image, stepSize, windowSize):
# slide a window across the image
for y in range(0, image.shape[0], stepSize):
    for x in range(0, image.shape[1], stepSize):
        # yield the current window
        yield (x, y, image[y:y + windowSize[1], x:x + windowSize[0]])

现在,我想能够预测每个64 * 64图像是否都包含飞机,但是我不知道该怎么做,因为net.predict()将数据集作为参数(arg:dim 4)

1 个答案:

答案 0 :(得分:0)

  

net.predict()将数据集作为参数(arg:暗4)

net.predict接受多种数据格式,其中包括数据集。但是,对于您来说,似乎最好是接受火炬张量或numpy数组-而且可以!只需将您的64x64块传递给net.predict,就像这样:

# (n, 512, 512, 3)
X = my_data
# (n, 4096, 64, 64, 3)
X = sliding_window(X, 64, 64)
# (n * 4096, 64, 64, 3)
X = X.reshape(-1, 64, 64, 3)
y = net.predict(X)