雷达图像识别

从小就不爱喝酸奶
228 阅读3分钟

一起养成写作习惯！这是我参与「掘金日新计划 · 4 月更文挑战」的第9天，点击查看活动详情。

目标

使用通过雷达图预测天气，最近看了一个使用卷积神经网络的示例，然而它是使用tensorflow 1.x 版本的api写的。和现在经常使用的深度学习框架使用差距过于巨大，为此尝试用更加用户友好的api实现相关的功能。

数据集的读取

数据集被存放在一个目录中，且各种类型的图片储存在这之中的子目录中。 这里分为四种天气，每种天气属于一个目录。

train

  • nowind-norain
    • xx.png
  • nowind-rain
    • xx.png

Pytorch 为我们提供了直接读取该种数据集的方式

train_transforms = torchvision.transforms.Compose([
    transforms.ToTensor()
])
dataset = ImageFolder(root=path, transform=train_transforms)

只需要把 path 设置成前述说明的那种目录即可。可以直接读取出所有的图片，并设定好类别。比起自行实现的一张张的读取目测能快 50 倍以上。

数据迭代器

类似的 pytorch 为我们实现了一个好用的迭代器，不需要我们自行实现并打乱数据了。

train_data = DataLoader(dataset, batch_size=100, shuffle=True)

参数的意思比较直观。

优化器

pytorch 有 Adam，SGD 等常用优化器可供使用。这里是分类任务。使用交叉熵损失函数搭配Adam 优化器使用。

loss_function = nn.CrossEntropyLoss()
adam_optimizer = optim.Adam(net.parameters(), lr=0.001, eps=1e-3)

模型的定义

定义一个卷积神经网络, 注意我们的数据是 3 x 100 x 100，最后分成4类的。


class Reshape(nn.Module):
    def __init__(self, shape):
        super(Reshape, self).__init__()
        self.shape = shape

    def forward(self, x):
        return x.view(self.shape)


net = nn.Sequential(
    nn.Conv2d(3, 32, kernel_size=5, stride=1, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(32, 64, kernel_size=5, stride=1, padding="same"),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(64, 128, kernel_size=3, stride=1, padding="same"),
    # nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(128, 128, kernel_size=3, stride=1, padding="same"),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    Reshape((-1, 6 * 6 * 128)),
    nn.Linear(6 * 6 * 128, 1024),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(1024, 512),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(512, 4),
)

稍微测试一下我们的网络输出

from torch.utils.data import DataLoader

import torch
images = torch.randn(2, 3, 100, 100)

# get shape of output
print(net(images).shape)

结果的形状是

torch.Size([2, 4])

模型的参数初始化

采用方差为0.1的正态分布对模型进行权重的初始化，由于实际上可能的需要初始化的只有线性层和卷积层，因此可以直接像下面这样对特定的层使用特定的方法进行初始化。

def init_weights(m):
    if type(m) == nn.Linear:
        torch.nn.init.normal_(m.weight, std=0.01)
        m.bias.data.fill_(0)
    elif type(m) == nn.Conv2d:
        torch.nn.init.normal_(m.weight, std=0.01)
        m.bias.data.fill_(0)
        
       
net.apply(init_weights)

不过这里其实也可以直接合并书写，因为实际上线性层等价于1x1的卷积层，有着近似的接口以及参数。

训练模型

采用 adam 优化器，交叉熵损失函数。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
loss_function = nn.CrossEntropyLoss()
adam_optimizer = optim.Adam(net.parameters(), lr=0.001, eps=1e-3)


# 保存模型的路径
model_path = "./model.pth"
# 训练次数
epochs = 4

for epoch in range(epochs):  
    for i, data in enumerate(train_data, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        # 梯度置 0 
        adam_optimizer.zero_grad()

        output = net(inputs)
        loss = loss_function(output, labels)
        loss.backward()
        adam_optimizer.step()

        # loss.item()
np.save(model_path, net)

保存模型以便于后续在现有基础上继续训练。

这里的参数也比较的清晰，和其他的框架不同的地方主要可能在于梯度运算。它的梯度会一致保留，需要我们自行清零。

至此主要流程的完成。

温馨提示：如果没有gpu可以使用Google的colab进行训练哦。

avatar
文章
阅读
粉丝