前言

在训练神经网络的过程中需要用到很多的工具，最重要的是数据处理、可视化和GPU加速。本章主要介绍PyTorch在这些方面常用的工具模块，合理使用这些工具可以极大地提高编程效率。

由于内容较多，本文分成了五篇文章（1）数据处理（2）预训练模型（3）TensorBoard（4）Visdom（5）CUDA与小结。

整体结构如下：

• 1 数据处理
  • 1.1 Dataset
  • 1.2 DataLoader
• 2 预训练模型
• 3 可视化工具
• 3.1 TensorBoard
• 3.2 Visdom
• 4 使用GPU加速：CUDA
• 5 小结

全文链接：

1. PyTorch中常用的工具（1）数据处理
2. PyTorch常用工具（2）预训练模型
3. PyTorch中常用的工具（3）TensorBoard
4. PyTorch中常用的工具（4）Visdom
5. PyTorch中常用的工具（5）使用GPU加速：CUDA

2 预训练模型

除了加载数据，并对数据进行预处理之外，torchvision还提供了深度学习中各种经典的网络结构以及预训练模型。这些模型封装在torchvision.models中，包括经典的分类模型：VGG、ResNet、DenseNet及MobileNet等，语义分割模型：FCN及DeepLabV3等，目标检测模型：Faster RCNN以及实例分割模型：Mask RCNN等。读者可以通过下述代码使用这些已经封装好的网络结构与模型，也可以在此基础上根据需求对网络结构进行修改：

from torchvision import models
# 仅使用网络结构，参数权重随机初始化
mobilenet_v2 = models.mobilenet_v2()
# 加载预训练权重
deeplab = models.segmentation.deeplabv3_resnet50(pretrained=True)

下面使用torchvision中预训练好的实例分割模型Mask RCNN进行一次简单的实例分割：

In: from torchvision import modelsfrom torchvision import transforms as Tfrom torch import nnfrom PIL import Imageimport numpy as npimport randomimport cv2# 加载预训练好的模型，不存在的话会自动下载# 预训练好的模型保存在 ~/.torch/models/下面detection = models.detection.maskrcnn_resnet50_fpn(pretrained=True)detection.eval()def predict(img_path, threshold):# 数据预处理，标准化至[-1, 1]，规定均值和标准差img = Image.open(img_path)transform = T.Compose([T.ToTensor(),T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])])img = transform(img)# 对图像进行预测pred = detection([img])# 对预测结果进行后处理：得到mask与bboxscore = list(pred[0]['scores'].detach().numpy())t = [score.index(x) for x in score if x > threshold][-1]mask = (pred[0]['masks'] > 0.5).squeeze().detach().cpu().numpy()pred_boxes = [[(i[0], i[1]), (i[2], i[3])] \for i in list(pred[0]['boxes'].detach().numpy())]pred_masks = mask[:t+1]boxes = pred_boxes[:t+1]return pred_masks, boxes

Transforms中涵盖了大部分对Tensor和PIL Image的常用处理，这些已在上文提到，本节不再详细介绍。需要注意的是转换分为两步，第一步：构建转换操作，例如transf = transforms.Normalize(mean=x, std=y)；第二步：执行转换操作，例如output = transf(input)。另外还可以将多个处理操作用Compose拼接起来，构成一个处理转换流程。

In: # 随机颜色，以便可视化def color(image):colours = [[0, 255, 255], [0, 0, 255], [255, 0, 0]]R = np.zeros_like(image).astype(np.uint8)G = np.zeros_like(image).astype(np.uint8)B = np.zeros_like(image).astype(np.uint8)R[image==1], G[image==1], B[image==1] = colours[random.randrange(0,3)]color_mask = np.stack([R,G,B],axis=2)return color_mask

In: # 对mask与bounding box进行可视化def result(img_path, threshold=0.9, rect_th=1, text_size=1, text_th=2):masks, boxes = predict(img_path, threshold)img = cv2.imread(img_path)img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)for i in range(len(masks)):color_mask = color(masks[i])img = cv2.addWeighted(img, 1, color_mask, 0.5, 0)cv2.rectangle(img, boxes[i][0], boxes[i][1], color=(255,0,0), thickness=rect_th)return img

In: from matplotlib import pyplot as pltimg=result('data/demo.jpg')plt.figure(figsize=(10, 10))plt.axis('off')img_result = plt.imshow(img)

上述代码完成了一个简单的实例分割任务。如上图所示，Mask RCNN能够分割出该图像中的部分实例，读者可考虑对预训练模型进行微调，以适应不同场景下的不同任务。注意：上述代码均在CPU上进行，速度较慢，读者可以考虑将数据与模型转移至GPU上，具体操作可以参考第4节。