AI部署：聊一聊深度學(xué)習(xí)中的模型權(quán)重

今天簡單聊聊模型權(quán)重，也就是我們俗稱的weight。

深度學(xué)習(xí)中，我們一直在訓(xùn)練模型，通過反向傳播求導(dǎo)更新模型的權(quán)重，最終得到一個(gè)泛化能力比較強(qiáng)的模型。同樣，如果我們不訓(xùn)練，僅僅隨機(jī)初始化權(quán)重，同樣能夠得到一個(gè)同樣大小的模型。雖然兩者大小一樣，不過兩者其中的權(quán)重信息分布相差會很大，一個(gè)腦子裝滿了知識、一個(gè)腦子都是水，差不多就這個(gè)意思。

所謂的AI模型部署階段，說白了就是將訓(xùn)練好的權(quán)重挪到另一個(gè)地方去跑。一般來說，權(quán)重信息以及權(quán)重分布基本不會變（可能會改變精度、也可能會合并一些權(quán)重）。

不過執(zhí)行模型操作（卷積、全連接、反卷積）的算子會變化，可能從Pytorch->TensorRT或者TensorFlow->TFLITE，也就是實(shí)現(xiàn)算子的方式變了，同一個(gè)卷積操作，在Pytorch框架中是一種實(shí)現(xiàn)，在TensorRT又是另一種時(shí)間，兩者的基本原理是一樣的，但是精度和速度不一樣，TensorRT可以借助Pytorch訓(xùn)練好的卷積的權(quán)重，實(shí)現(xiàn)與Pytorch中一樣的操作，不過可能更快些。

權(quán)重/Weight/CheckPoint

那么權(quán)重都有哪些呢？他們長什么樣？

這還真不好描述...其實(shí)就是一堆數(shù)據(jù)。對的，我們千辛萬苦不斷調(diào)優(yōu)訓(xùn)練出來的權(quán)重，就是一堆數(shù)據(jù)而已。也就是這個(gè)神奇的數(shù)據(jù)，搭配各種神經(jīng)網(wǎng)絡(luò)的算子，就可以實(shí)現(xiàn)各種檢測、分類、識別的任務(wù)。

例如上圖，我們用Netron這個(gè)工具去查看某個(gè)ONNX模型的第一個(gè)卷積權(quán)重。很顯然這個(gè)卷積只有一個(gè)W權(quán)重，沒有偏置b。而這個(gè)卷積的權(quán)重值的維度是[64,3,7,7]，也就是輸入通道3、輸出通道64、卷積核大小7x7。

再仔細(xì)看，其實(shí)這個(gè)權(quán)重的數(shù)值范圍相差還是很大，最大的也就0.1的級別。但是最小的呢，肉眼看了下（其實(shí)應(yīng)該統(tǒng)計(jì)一波），最小的竟然有1e-10級別。

一般我們訓(xùn)練的時(shí)候，輸入權(quán)重都是0-1，當(dāng)然也有0-255的情況，但不論是0-1還是0-255，只要不溢出精度上限和下限，就沒啥問題。對于FP32來說，1e-10是小case，但是對于FP16來說就不一定了。

我們知道FP16的普遍精度是~5.96e?8 (6.10e?5) … 65504，具體的精度細(xì)節(jié)先不說，但是可以很明顯的看到，上述的1e-10的精度，已經(jīng)溢出了FP16的精度下限。如果一個(gè)模型中的權(quán)重分布大部分都處在溢出邊緣的話，那么模型轉(zhuǎn)換完FP16精度的模型指標(biāo)可能會大大下降。

除了FP16，當(dāng)然還有很多其他精度(TF32、BF16、IN8)，這里暫且不談，不過有篇討論各種精度的文章可以先了解下：https://moocaholic.medium.com/fp64-fp32-fp16-bfloat16-tf32-and-other-members-of-the-zoo-a1ca7897d407

話說回來，我們該如何統(tǒng)計(jì)該層的權(quán)重信息呢？利用Pytorch中原生的代碼就可以實(shí)現(xiàn)：

# 假設(shè)v是某一層conv的權(quán)重，我們可以簡單通過以下命令查看到該權(quán)重的分布
v.max()
tensor(0.8559)
v.min()
tensor(-0.9568)
v.abs()
tensor([[0.0314, 0.0045, 0.0182,  ..., 0.0309, 0.0204, 0.0345],
        [0.0295, 0.0486, 0.0746,  ..., 0.0363, 0.0262, 0.0108],
        [0.0328, 0.0582, 0.0149,  ..., 0.0932, 0.0444, 0.0221],
        ...,
        [0.0337, 0.0518, 0.0280,  ..., 0.0174, 0.0078, 0.0010],
        [0.0022, 0.0297, 0.0167,  ..., 0.0472, 0.0006, 0.0128],
        [0.0631, 0.0144, 0.0232,  ..., 0.0072, 0.0704, 0.0479]])
v.abs().min() # 可以看到權(quán)重絕對值的最小值是1e-10級別
tensor(2.0123e-10)
v.abs().max()
tensor(0.9568)
torch.histc(v.abs()) # 這里統(tǒng)計(jì)權(quán)重的分布，分為100份，最小最大分別是[-0.9558,0.8559]
tensor([3.3473e+06, 3.2437e+06, 3.0395e+06, 2.7606e+06, 2.4251e+06, 2.0610e+06,
        1.6921e+06, 1.3480e+06, 1.0352e+06, 7.7072e+05, 5.5376e+05, 3.8780e+05,
        2.6351e+05, 1.7617e+05, 1.1414e+05, 7.3327e+04, 4.7053e+04, 3.0016e+04,
        1.9576e+04, 1.3106e+04, 9.1220e+03, 6.4780e+03, 4.6940e+03, 3.5140e+03,
        2.8330e+03, 2.2040e+03, 1.7220e+03, 1.4020e+03, 1.1130e+03, 1.0200e+03,
        8.2400e+02, 7.0600e+02, 5.7900e+02, 4.6400e+02, 4.1600e+02, 3.3400e+02,
        3.0700e+02, 2.4100e+02, 2.3200e+02, 1.9000e+02, 1.5600e+02, 1.1900e+02,
        1.0800e+02, 9.9000e+01, 6.9000e+01, 5.2000e+01, 4.9000e+01, 2.2000e+01,
        1.8000e+01, 2.8000e+01, 1.2000e+01, 1.3000e+01, 8.0000e+00, 3.0000e+00,
        4.0000e+00, 3.0000e+00, 1.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00,
        1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 2.0000e+00,
        0.0000e+00, 2.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00,
        2.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00])

這樣看如果覺著不是很直觀，那么也可以自己畫圖或者通過Tensorboard來時(shí)候看。

那么看權(quán)重分布有什么用呢？

肯定是有用處的，訓(xùn)練和部署的時(shí)候權(quán)重分布可以作為模型是否正常，精度是否保持的一個(gè)重要信息。不過這里先不展開說了。

有權(quán)重，所以重點(diǎn)關(guān)照

在模型訓(xùn)練過程中，有很多需要通過反向傳播更新的權(quán)重，常見的有：

• 卷積層
• 全連接層
• 批處理化層(BN層、或者各種其他LN、IN、GN)
• transformer-encoder層
• DCN層

這些層一般都是神經(jīng)網(wǎng)絡(luò)的核心部分，當(dāng)然都是有參數(shù)的，一定會參與模型的反向傳播更新，是我們在訓(xùn)練模型時(shí)候需要注意的重要參數(shù)。

# Pytorch中conv層的部分代碼，可以看到參數(shù)的維度等信息
self._reversed_padding_repeated_twice = _reverse_repeat_tuple(self.padding, 2)
if transposed:
    self.weight = Parameter(torch.Tensor(
        in_channels, out_channels // groups, *kernel_size))
else:
    self.weight = Parameter(torch.Tensor(
        out_channels, in_channels // groups, *kernel_size))
if bias:
    self.bias = Parameter(torch.Tensor(out_channels))

也有不參與反向傳播，但也會隨著訓(xùn)練一起更新的參數(shù)。比較常見的就是BN層中的running_mean和running_std：

# 截取了Pytorch中BN層的部分代碼
def __init__(
    self,
    num_features: int,
    eps: float = 1e-5,
    momentum: float = 0.1,
    affine: bool = True,
    track_running_stats: bool = True
) -> None:
    super(_NormBase, self).__init__()
    self.num_features = num_features
    self.eps = eps
    self.momentum = momentum
    self.affine = affine
    self.track_running_stats = track_running_stats
    if self.affine:
        self.weight = Parameter(torch.Tensor(num_features))
        self.bias = Parameter(torch.Tensor(num_features))
    else:
        self.register_parameter('weight', None)
        self.register_parameter('bias', None)
    if self.track_running_stats:
        # 可以看到在使用track_running_stats時(shí)，BN層會更新這三個(gè)參數(shù)
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
    else:
        self.register_parameter('running_mean', None)
        self.register_parameter('running_var', None)
        self.register_parameter('num_batches_tracked', None)
    self.reset_parameters()

可以看到上述代碼的注冊區(qū)別，對于BN層中的權(quán)重和偏置使用的是register_parameter，而對于running_mean和running_var則使用register_buffer，那么這兩者有什么區(qū)別呢，那就是注冊為buffer的參數(shù)往往不會參與反向傳播的計(jì)算，但仍然會在模型訓(xùn)練的時(shí)候更新，所以也需要認(rèn)真對待。

關(guān)于BN層，轉(zhuǎn)換模型和訓(xùn)練模型的時(shí)候會有暗坑，需要注意一下。

剛才描述的這些層都是有參數(shù)的，那么還有一些沒有參數(shù)的層有哪些呢？當(dāng)然有，我們的網(wǎng)絡(luò)中其實(shí)有很多op，僅僅是做一些維度變換、索引取值或者上/下采樣的操作，例如：

• Reshape
• Squeeze
• Unsqueeze
• Split
• Transpose
• Gather

等等等等，這些操作沒有參數(shù)僅僅是對上一層傳遞過來的張量進(jìn)行維度變換，用于實(shí)現(xiàn)一些”炫技“的操作。至于這些炫技嗎，有些很有用有些就有些無聊了。

開個(gè)玩笑，其實(shí)有時(shí)候在通過Pytorch轉(zhuǎn)換為ONNX的時(shí)候，偶爾會發(fā)生一些轉(zhuǎn)換詭異的情況。比如一個(gè)簡單的reshape會四分五裂為gather+slip+concat，這種操作相當(dāng)于復(fù)雜化了，不過一般來說這種情況可以使用ONNX-SIMPLIFY去優(yōu)化掉，當(dāng)然遇到較為復(fù)雜的就需要自行優(yōu)化了。

哦對了，對于這些變形類的操作算子，其實(shí)有些是有參數(shù)的，例如下圖的reshap:

不過這些都是小問題，大部分情況我們可以通過改模型或者換結(jié)構(gòu)解決，而且成本也不高。但是還會有一些其他復(fù)雜的問題，可能就需要我們重點(diǎn)研究下了。

提取權(quán)重

想要將訓(xùn)練好的模型從這個(gè)平臺部署至另一個(gè)平臺，那么首要的就是轉(zhuǎn)移權(quán)重。不過實(shí)際中大部分的轉(zhuǎn)換器都幫我們做好了（比如onnx-TensorRT），不用我們自己操心！

onnx-TensorRT：https://github.com/onnx/onnx-tensorrt

不過如果想要對模型權(quán)重的有個(gè)整體認(rèn)知的話，還是建議自己親手試一試。

Caffe2Pytorch

先簡單說下Caffe和Pytorch之間的權(quán)重轉(zhuǎn)換。這里推薦一個(gè)開源倉庫Caffe-python（https://github.com/marvis/pytorch-caffe），已經(jīng)幫我們寫好了提取Caffemodel權(quán)重和根據(jù)prototxt構(gòu)建對應(yīng)Pytorch模型結(jié)構(gòu)的過程，不需要我們重復(fù)造輪子。

model = caffe_pb2.NetParameter()
print('Loading caffemodel: ' + caffemodel)
with open(caffemodel, 'rb') as fp:
    model.ParseFromString(fp.read())

caffe_pb2就是caffemodel格式的protobuf結(jié)構(gòu)，具體的可以看上方老潘提供的庫，總之就是定義了一些Caffe模型的結(jié)構(gòu)。

而提取到模型權(quán)重后，通過prototxt中的模型信息，挨個(gè)從caffemodel的protobuf權(quán)重中找，然后復(fù)制權(quán)重到Pytorch端，仔細(xì)看這句caffe_weight = torch.from_numpy(caffe_weight).view_as(self.models[lname].weight)，其中self.models[lname]就是已經(jīng)搭建好的對應(yīng)Pytorch的卷積層，這里取weight之后通過self.models[lname].weight.data.copy_(caffe_weight)將caffe的權(quán)重放到Pytorch中。

很簡單吧。

if ltype in ['Convolution', 'Deconvolution']:
    print('load weights %s' % lname)
    convolution_param = layer['convolution_param']
    bias = True
    if 'bias_term' in convolution_param and convolution_param['bias_term'] == 'false':
        bias = False
    # weight_blob = lmap[lname].blobs[0]
    # print('caffe weight shape', weight_blob.num, weight_blob.channels, weight_blob.height, weight_blob.width)
    caffe_weight = np.array(lmap[lname].blobs[0].data)
    caffe_weight = torch.from_numpy(caffe_weight).view_as(self.models[lname].weight)
    # print("caffe_weight", caffe_weight.view(1,-1)[0][0:10])
    self.models[lname].weight.data.copy_(caffe_weight)
    if bias and len(lmap[lname].blobs) > 1:
        self.models[lname].bias.data.copy_(torch.from_numpy(np.array(lmap[lname].blobs[1].data)))
        print("convlution %s has bias" % lname)

Pytorch2TensorRT

先舉個(gè)簡單的例子，一般我們使用Pytorch模型進(jìn)行訓(xùn)練。訓(xùn)練得到的權(quán)重，我們一般都會使用torch.save()保存為.pth的格式。

PTH是Pytorch使用python中內(nèi)置模塊pickle來保存和讀取，我們使用netron看一下pth長什么樣。。

可以看到只有模型中有參數(shù)權(quán)重的表示，并不包含模型結(jié)構(gòu)。不過我們可以通過.py的模型結(jié)構(gòu)一一加載.pth的權(quán)重到我們模型中即可。

看一下我們讀取.pth后，state_dict的key。這些key也就對應(yīng)著我們在構(gòu)建模型時(shí)候注冊每一層的權(quán)重名稱和權(quán)重信息（也包括維度和類型等）。

當(dāng)然這個(gè)pth也可以包含其他字符段{'epoch': 190, 'state_dict': OrderedDict([('conv1.weight', tensor([[...，比如訓(xùn)練到多少個(gè)epoch，學(xué)習(xí)率啥的。

對于pth，我們可以通過以下代碼將其提取出來，存放為TensorRT的權(quán)重格式。

def extract_weight(args):
    # Load model
    state_dict = torch.load(args.weight)
    with open(args.save_path, "w") as f:
        f.write("{}\n".format(len(state_dict.keys())))
        for k, v in state_dict.items():
            vr = v.reshape(-1).cpu().numpy()
            f.write("{} {} ".format(k, len(vr)))
            for vv in vr:
                f.write(" ")
                f.write(struct.pack(">f", float(vv)).hex())
            f.write("\n")

需要注意，這里的TensorRT權(quán)重格式指的是在build之前的權(quán)重，TensorRT僅僅是拿來去構(gòu)建整個(gè)網(wǎng)絡(luò)，將每個(gè)解析到的層的權(quán)重傳遞進(jìn)去，然后通過TensorRT的network去build好engine。

// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
    std::cout << "Loading weights: " << file << std::endl;
    std::map<std::string, Weights> weightMap;

    // Open weights file
    std::ifstream input(file);
    assert(input.is_open() && "Unable to load weight file.");

    // Read number of weight blobs
    int32_t count;
    input >> count;
    assert(count > 0 && "Invalid weight map file.");

    while (count--)
    {
        Weights wt{DataType::kFLOAT, nullptr, 0};
        uint32_t size;

        // Read name and type of blob
        std::string name;
        input >> name >> std::dec >> size;
        wt.type = DataType::kFLOAT;

        // Load blob
        uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
        for (uint32_t x = 0, y = size; x < y; ++x)
        {
            input >> std::hex >> val[x];
        }
        wt.values = val;
        wt.count = size;
        weightMap[name] = wt;
    }
    std::cout << "Finished Load weights: " << file << std::endl;
    return weightMap;
}

那么被TensorRT優(yōu)化后？模型又長什么樣子呢？我們的權(quán)重放哪兒了呢？

肯定在build好后的engine里頭，不過這些權(quán)重因?yàn)門ensorRT的優(yōu)化，可能已經(jīng)被合并/移除/merge了。

模型參數(shù)的學(xué)問還是很多，近期也有很多相關(guān)的研究，比如參數(shù)重參化，是相當(dāng)solid的工作，在很多訓(xùn)練和部署場景中經(jīng)常會用到。

后記

先說這些吧，比較基礎(chǔ)，也偏向于底層些。神經(jīng)網(wǎng)絡(luò)雖然一直被認(rèn)為是黑盒，那是因?yàn)闆]有確定的理論證明。但是訓(xùn)練好的模型權(quán)重我們是可以看到的，模型的基本結(jié)構(gòu)我們也是可以知道的，雖然無法證明模型為什么起作用？為什么work？但通過結(jié)構(gòu)和權(quán)重分布這些先驗(yàn)知識，我們也可以大概地對模型進(jìn)行了解，也更好地進(jìn)行部署。

至于神經(jīng)網(wǎng)絡(luò)的可解釋性，這就有點(diǎn)玄學(xué)了，我不清楚這里也就不多說了~

如果覺得有用，就請分享到朋友圈吧！