为此，我将使用来自谷歌地图的地图图片组成的数据集，并根据它们所包含的地形特征对它们进行分类。我会写另一个关于我如何使用它的故事(简而言之: 为了确定无人机飞越或降落的安全区域)。但是现在，我只想使用一些训练数据来对这些地图图片进行分类。
下面的代码片段来自于一个 Jupyter Notebook。您可以将它们组合在一起构建您自己的 Python 脚本。
搞定训练集
PyTorch 希望数据按照文件夹组织，每个类有一个文件夹。大多数其他 PyTorch 教程和示例都希望您用一个训练和验证文件夹进一步组织它，然后在它们内部使用类文件夹。但是我认为这是非常麻烦的，必须从每个类中选择一定数量的图像，然后将它们从训练文件夹移动到验证文件夹。而且由于大多数人会通过选择一个连续的文件组来实现这一点，因此在这个选择中可能存在很多偏差。
因此，这里有一个更好的方法来将数据集分割成动态的训练和测试集，就像 Python 开发人员习惯于从 SKLearn 中使用的那样。但是首先，让我们导入模块：
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
import matplotlib.pyplot as pltimport numpy as np
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models接下来我们将定义训练/验证数据集加载器，使用 SubsetRandomSampler 进行分割：
data_dir = '/data/train'def load_split_train_test(datadir, valid_size = .2):
train_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
]) test_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
]) train_data = datasets.ImageFolder(datadir,
transform=train_transforms)
test_data = datasets.ImageFolder(datadir,
transform=test_transforms) num_train = len(train_data)
indices = list(range(num_train))
split = int(np.floor(valid_size * num_train))
np.random.shuffle(indices)
from torch.utils.data.sampler import SubsetRandomSampler
train_idx, test_idx = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_idx)
test_sampler = SubsetRandomSampler(test_idx)
trainloader = torch.utils.data.DataLoader(train_data,
sampler=train_sampler, batch_size=64)
testloader = torch.utils.data.DataLoader(test_data,
sampler=test_sampler, batch_size=64)
return trainloader, testloadertrainloader, testloader = load_split_train_test(data_dir, .2)
print(trainloader.dataset.classes)接下来，我们将决定我们是否使用 GPU。我假设如果你这样做，你有一个安装有 GPU 驱动的电脑，否则代码运行速度将至少慢10倍。但是一般情况下，检查 GPU 的可用性是个好主意。
我们还将加载一个预训练模型：
device = torch.device("cuda" if torch.cuda.is_available()
else "cpu")
model = models.resnet50(pretrained=True)
print(model)输出模型将向您展示 ResNet 模型的层架构。这可能超出了我或你的理解，但是看看这些深层隐藏的东西还是很有趣的。
这取决于您选择什么模型，并且可能根据您的特定数据集选择不同的模型。
现在我们进入深层神经网络的有趣部分。首先，我们必须冻结预训练的层。然后，我们重新定义最终的全连接层，我们将用我们的图像训练的层。我们还创建了标准(损失函数) ，并选择合适的优化器(本例中为 Adam)和学习率。
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Sequential(nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(512, 10),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.003)
model.to(device)现在，让我们来训练我们的模型！在这个例子中只有一个 epoch，但是在大多数情况下你需要更多。从代码中可以看出，基本过程相当直观：加载图像的批处理并执行前向训练。然后计算损失函数，并使用优化器应用梯度下降法进行反向传播。
PyTorch 就是这么简单。下面的大部分代码每10个 batch 处理损失并计算准确性，因此您可以在训练运行时获得参数的更新。在验证期间，不要忘记将模型设置为 eval（)模式，然后在完成之后返回 train()模式。
epochs = 1
steps = 0
running_loss = 0
print_every = 10
train_losses, test_losses = [], []for epoch in range(epochs):
for inputs, labels in trainloader:
steps += 1
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
logps = model.forward(inputs)
loss = criterion(logps, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if steps % print_every == 0:
test_loss = 0
accuracy = 0
model.eval（)
with torch.no_grad():
for inputs, labels in testloader:
inputs, labels = inputs.to(device),
labels.to(device)
logps = model.forward(inputs)
batch_loss = criterion(logps, labels)
test_loss += batch_loss.item()
ps = torch.exp(logps)
top_p, top_class = ps.topk(1, dim=1)
equals =
top_class == labels.view(*top_class.shape)
accuracy +=
torch.mean(equals.type(torch.FloatTensor)).item()
train_losses.append(running_loss/len(trainloader))
test_losses.append(test_loss/len(testloader))
print(f"Epoch {epoch+1}/{epochs}.. "
f"Train loss: {running_loss/print_every:.3f}.. "
f"Test loss: {test_loss/len(testloader):.3f}.. "
f"Test accuracy: {accuracy/len(testloader):.3f}")
running_loss = 0
model.train()
torch.save(model, 'aerialmodel.pth')然后... 等待几分钟(或者更长时间，取决于数据集的大小和设置的 epoch 的数目)之后，训练就结束了，模型保存下来以便以后进行预测！
现在还有一件事情你可以做，那就是计算训练和验证的损失：
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)
plt.show()
正如您可以看到的，在我的一个特定的例子中，验证损失(这是我们感兴趣的)在第一个 epoch 结束时趋于平缓，甚至开始上升趋势，所以可能1个 epoch 就足够了。训练损失，正如预期的那样，是很低的。
现在进入第二部分。目前您对模型进行了训练，保存了它，并且需要在应用程序中使用它。为此，您需要能够对图像执行简单的推理。这样我就可以使用一些例子：
data_dir = '/datadrive/FastAI/data/aerial_photos/train'test_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
])然后我们再次检查 GPU 的可用性，加载模型并将其放入评估模式(这样参数就不会改变) ：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model=torch.load('aerialmodel.pth')
model.eval（)预测特定图像类别的函数非常简单。注意，它需要一个 Pillow 图像，而不是一个文件路径。
def predict_image(image):
image_tensor = test_transforms(image).float()
image_tensor = image_tensor.unsqueeze_(0)
input = Variable(image_tensor)
input = input.to(device)
output = model(input)
index = output.data.cpu().numpy().argmax()
return index现在为了更容易测试，我还创建了一个函数，可以从数据集文件夹中随机选取一些图片：
def get_random_images(num):
data = datasets.ImageFolder(data_dir, transform=test_transforms)
classes = data.classes
indices = list(range(len(data)))
np.random.shuffle(indices)
idx = indices[:num]
from torch.utils.data.sampler import SubsetRandomSampler
sampler = SubsetRandomSampler(idx)
loader = torch.utils.data.DataLoader(data,
sampler=sampler, batch_size=num)
dataiter = iter(loader)
images, labels = dataiter.next()
return images, labels最后，为了演示预测函数，我得到了随机的图像样本，对它们进行预测并显示结果：
to_pil = transforms.ToPILImage()
images, labels = get_random_images(5)
fig=plt.figure(figsize=(10,10))
for ii in range(len(images)):
image = to_pil(images[ii])
index = predict_image(image)
sub = fig.add_subplot(1, len(images), ii+1)
res = int(labels[ii]) == index
sub.set_title(str(classes[index]) + ":" + str(res))
plt.axis('off')
plt.imshow(image)
plt.show()这里有一个关于 Google 地图的预测的例子。标签是预测的类，我还显示了它是否是一个正确的预测。
差不多就是这样了。请继续在您的数据集上尝试。只要您正确地输入并处理您的图像，这段代码应该可以正常工作。
