計圖(Jittor) 1.1版本：新增骨干網絡、JIT功能升級、支持多卡訓練
深度學習框架—計圖（Jittor），Jittor的新版本V1.1上線了。主要變化包括：
? 增加了大量骨干網絡的支持，增強了輔助轉換腳本的能力，降低用戶開發和移植模型的難度。
? JIT（動態編譯）功能升級，可支持高性能的自定義算子開發，并降低了用戶開發自定義算子的難度。
? 新增分布式功能，用戶無需修改代碼，只需要修改啟動命令，單卡版本的訓練程序可以直接無縫部署到多卡甚至多機上。
此外，Jittor還新增支持了大量神經網絡算子，完善了對深度神經網絡開發的支持。
新增骨干網絡
Jittor 1.1版本在模型庫中新增了大量基礎骨干網絡，用戶可以基于這些基礎骨干網絡搭建自己的深度學習模型。這些骨干網絡參數與PyTorch參數格式兼容，可以相互加載調用，方便用戶學習和遷移。
更新的骨干網絡包括：
import jittor.models
from jittor.models import
alexnet,
googlenet,
inception_v3,
mnasnet,
mobilenet_v2,
resnext101_32x8d,
shufflenetv2,
squeezenet,
wide_resnet101_2

更多主干網絡支持請參考

print(dir(jittor.models))
下面是對AlexNet、VGG、ResNet、Wide ResNet和SqueezeNet等多種骨干網絡，在不同的batch size下，Jittor與PyTorch的單次前向性能對比，圖1中橫軸為不同Batch size，縱軸為FPS(每秒處理的圖像數)。實驗環境為：GPU為1080ti，顯存11GB，CPU為i7-6850K，內存32GB，使用32位浮點數計算。

可以看出Jittor在這些常用骨干網絡上的速度非常可觀，大多數網絡Jittor會有速度的提升。
增強輔助轉換腳本
此次更新還提供了簡單的輔助轉換腳本，支持從PyTorch轉換模型代碼，暫不支持模型以外代碼，下面是一個簡單的實例：
from jittor.utils.pytorch_converter import convert
pytorch_code = “”"
class Model(nn.Module):
def init(self):
super().init()
self.conv1 = nn.Conv2d(1, 10, 3)
self.conv2 = nn.Conv2d(10, 32, 3)
self.fc = nn.Linear(1200, 100)

def forward(self, x):x = self.conv1(x)x = self.conv2(x)x = x.view(x.size(0), -1)x = self.fc(x)return x

“”"
jittor_code = convert(pytorch_code)
print("## Generate Jittor code:", jittor_code)
exec(jittor_code)
model = Model()
print("## Jittor model:", model)
上面的腳本將輸出計圖模型代碼，以及一個可用的模型：

生成的計圖代碼:

import jittor as jt
from jittor import init
from jittor import nn

class Model(nn.Module):

def __init__(self):super().__init__()self.conv1 = nn.Conv(1, 10, 3)self.conv2 = nn.Conv(10, 32, 3)self.fc = nn.Linear(1200, 100)def execute(self, x):x = self.conv1(x)x = self.conv2(x)x = x.view((x.shape[0], (- 1)))x = self.fc(x)return x

計圖模型:

Model(
conv1: Conv(1, 10, (3, 3), (1, 1), (0, 0), (1, 1), groups=1, bias=float[10,])
conv2: Conv(10, 32, (3, 3), (1, 1), (0, 0), (1, 1), groups=1, bias=float[32,])
fc: Linear(1200, 100, float32[100,], None)
)
關于模型轉換腳本的詳細使用請參考計圖官方文檔。
JIT功能升級
Jittor 1.1版本的 JIT(動態編譯)功能的增強，主要體現在code算子的功能增強上。Code算子是一個基于高性能語言的動態編譯算子，允許用戶直接在Python中內聯C++/CUDA代碼，只需要寥寥數行代碼，就可以完成高性能的自定義算子開發，降低用戶開發自定義算子的難度。
下面的實例展示了如何使用code算子，使用數行代碼實現三維點云中十分常用的K近鄰查找。Code算子的設計和實現，讓用戶既可以享受到Python語言的便捷與易用性，又可以獲得高性能語言的性能。
a = jt.random((n,3))
b = jt.code([n, k], “int32”, [a],
cpu_header="#include “,
cpu_src=”""
using namespace std;
auto n=out_shape0, k=out_shape1;
// 使用openmp實現自動并行化
#pragma omp parallel for
for (int i=0; i<n; i++) {
// 存儲k近鄰的距離和下標
vector<pair<float,int>> id(n);
for (int j=0; j<n; j++) {
auto dx = @in0(i,0)-@in0(j,0);
auto dy = @in0(i,1)-@in0(j,1);
auto dz = @in0(i,2)-@in0(j,2);
id[j] = {dxdx+dydy+dz*dz, j};
}
// 使用c++算法庫的nth_element排序
nth_element(id.begin(),
id.begin()+k, id.end());
// 將下標輸出到計圖的變量中
for (int j=0; j<k; j++)
@out(i,j) = id[j].second;
}"""
)
將計圖使用code算子實現的K近鄰查找，和PyTorch的算子用時進行比較，速度對比如下（k=10,點云數量n=[100,1000,10000]）：

注：此處使用的K近鄰算法為暴力算法，還存在更優的算法實現，由于文章篇幅有限，此處僅用于展示Code算子的使用。
更新大幅度提升了code算子的易用性和可讀性，主要包含以下幾點：
? code 算子 可以有多個輸出
? code 算子 允許輸出動態大小的變量
? code 算子 內部可以寫注釋
? code 算子 可以通過@alias為input和outputs增加別名，增加代碼可讀性
具體文檔請參考help(jt.code)和提供的文檔，計圖提供了多種實例以供參考。目前內聯C++代碼支持CUDA和openmp，未來會加入更多語言和庫的支持。
分布式接口
計圖本次分布式更新主要基于MPI（Message Passing Interface），依賴OpenMPI，用戶可以使用如下命令安裝OpenMPI：
sudo apt install openmpi-bin openmpi-common libopenmpi-dev
OpenMPI安裝完成以后，用戶無需修改代碼，需要做的僅僅是修改啟動命令行，計圖就會用數據并行的方式自動完成并行操作。

單卡訓練代碼

python3.7 -m jittor.test.test_resnet

分布式多卡訓練代碼

mpirun -np 4 python3.7 -m jittor.test.test_resnet

指定特定顯卡的多卡訓練代碼

CUDA_VISIBLE_DEVICES=“2,3” mpirun -np 2 python3.7 -m jittor.test.test_resnet
這種便捷性的背后是計圖的分布式算子的支撐，計圖支持的mpi算子后端會使用nccl進行進一步的加速。計圖所有分布式算法的開發均在Python前端完成，這讓分布式算法的靈活度增強，開發分布式算法的難度也大大降低。下面的代碼是使用計圖實現分布式同步批歸一化層的實例代碼：
def execute(self, x):
if self.is_train:
xmean = jt.mean(x, dims=[0,2,3], keepdims=1)
x2mean = jt.mean(x*x, dims=[0,2,3], keepdims=1)
if self.sync and jt.mpi:
xmean = xmean.mpi_all_reduce(“mean”)
x2mean = x2mean.mpi_all_reduce(“mean”)

    xvar = x2mean-xmean*xmeannorm_x = (x-xmean)/jt.sqrt(xvar+self.eps)self.running_mean += (xmean.sum([0,2,3])-self.running_mean)*self.momentumself.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum
else:running_mean = self.running_mean.broadcast(x, [0,2,3])running_var = self.running_var.broadcast(x, [0,2,3])norm_x = (x-running_mean)/jt.sqrt(running_var+self.eps)
w = self.weight.broadcast(x, [0,2,3])
b = self.bias.broadcast(x, [0,2,3])
return norm_x * w + b

這次更新，開放了mpi算子的穩定接口，用戶可以自行使用mpi算子開發所需的自定義的分布式算法，相關文檔請參考help(jittor.mpi.ops)和計圖分布式教程。基于這些mpi算子接口，研發團隊已經集成了如下三種分布式相關的算法：
? 分布式數據并行加載
? 分布式優化器
? 分布式同步批歸一化層
其他更新
教程更新列表：
? 模型輔助轉換工具教程
? 計圖多卡分布式教程
? 利用Code算子簡化高性能算子開發
算子更新列表：
? group conv
? 三角函數，反三角函數，雙曲函數，反雙曲函數支持
? flatten
? view
? permute
? adapool
? PReLU
? LeakyReLU
? ReLU6
? ReflectionPad2d
? ZeroPad2d
? ConstantPad2d
? ReplicationPad2d
? PixelShuffle
? Upsample
損失函數更新列表：
? MSELoss
? BCELoss
? L1Loss
? BCEWithLogitsLoss
分布式相關算子：
? mpi_all_reduce
? mpi_reduce
? mpi_broadcast
? nccl_all_reduce
? nccl_reduce
? nccl_broadcast

總結

以上是生活随笔為你收集整理的计图(Jittor) 1.1版本：新增骨干网络、JIT功能升级、支持多卡训练的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。