摘要：這篇文章跟大家分享下THOR的實踐應用。THOR算法的部分內容當前已經在MindSpore中開源
本文分享自華為雲社區《MindSpore 自研高階優化器源碼分析和實踐應用》，原文作者：HWCloudAI 。
這篇文章跟大家分享下THOR的實踐應用。THOR算法的部分內容當前已經在MindSpore中開源，源碼位置：
https://gitee.com/mindspore/mindspore/blob/master/mindspore/nn/optim/thor.py
MindSpore中使用THOR訓練網絡非常簡單，下面用四行代碼先來帶大家看一下怎么使用。
from mindspore.nn.optim import THOR #引用二階優化器 #創建網絡 net = Net() #調用優化器 opt = THOR(net, lr, Tensor(damping), config.momentum, config.weight_decay, config.loss_scale, config.batch_size, split_indices=split_indices) #增加計算圖提升性能 model = ConvertModelUtils().convert_to_thor_model(model=model, network=net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False, frequency=config.frequency) #訓練網絡 model.train(config.epoch_size, dataset, callbacks=cb, sink_size=dataset.get_dataset_size(), dataset_sink_mode=True)
- 導入二階優化器THOR所需要的包
- 第一行代碼常規創建網絡
- 第二行代碼定義我們使用的優化器THOR
- 第三行代碼是為了增加計算圖從而使THOR達到更優性能
- 第四行代碼訓練網絡
我們再具體展開介紹下。首先導入MindSpore所需的二階優化器的包，位於 mindspore.nn.optim
然后創建你所需的網絡；接着定義THOR優化器，傳入網絡信息和THOR所需的超參信息（如學習率，正則化項系數等）；
再調用 convert_to_thor_model函數，該函數是通過增加計算圖使THOR達到更優性能，什么意思呢，本身網絡運行的時候是一張計算圖，THOR中會使用過時的二階信息，通過額外增加一張計算圖，兩張計算圖分別執行更新二階矩陣和不更新二階矩陣的操作從而達到更優性能（PS. MindSpore支持動靜態圖，在這里為了更好的性能使用的是靜態圖模式，對這塊內容比較感興趣的同學，可以點這個鏈接：https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/white_paper/MindSpore_white_paper.pdf）；
最后，調用model.train就可以開始訓練啦。簡單介紹了下怎么使用，接下來我們來看下它的源碼。
源碼分析
init 函數用於THOR的初始化，需要傳入THOR所需的超參和網絡結構，THOR支持GPU和Ascend，分別為class THOR_GPU(Optimizer)和 class THOR_Ascend(Optimizer)，這兩個類之間的主要差別是算子不同。下面我們以 class THOR_Ascend(Optimizer)為例，來分析一下。
class THOR_Ascend(Optimizer): def __init__(self, net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32, decay_filter=lambda x: x.name not in [], split_indices=None): params = filter(lambda x: x.requires_grad, net.get_parameters()) super(THOR_Ascend, self).__init__(learning_rate, params, weight_decay, loss_scale) if isinstance(momentum, float) and momentum < 0.0: raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum)) self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum") self.params = self.parameters self.moments = self.params.clone(prefix="moments", init='zeros') self.hyper_map = C.HyperMap() self.opt = P.ApplyMomentum() self.net = net self.matrix_A_cov = ParameterTuple(filter(lambda x: 'matrix_A' in x.name, net.get_parameters())) self.matrix_G_cov = ParameterTuple(filter(lambda x: 'matrix_G' in x.name, net.get_parameters())) ...
MindSpore中所有優化器都繼承了 class Optimizer，該基類中定義了一些基本函數（如獲取學習率，梯度縮放等）。THOR初始化時將傳進去的超參定義為類屬性方便調用，並且定義了后續計算會使用到的算子。
也就是說初始化函數的作用就是定義THOR計算所需要用到的算子和變量（Parameter，Tensor等）。
重點介紹下 self.matrix_A_cov , self.matrix_G_cov 。這兩個變量是計算二階梯度所需要的信息，分別為每層輸入 的協方差矩陣 和每層輸出的一階導數 的協方差矩陣 ，其中 已經在運行時的前向過程和反向過程中保存下來。
我們再來看下創建THOR時的入參：
- net：本次訓練建立的模型；
- learning_rate：學習率超參；
- damping：二階矩陣中加的正則化項的超參；
- momentum：動量超參；
- weight_decay：權值衰減，用於防止過擬合，默認值為0.0，即不使用權值衰減；loss_scale：用於縮放訓練過程中的loss，防止梯度越界，默認值為1.0，即不使用縮放；batch_size：當前訓練一個step所使用的數據量，默認為32；
- decay_filter：選擇對哪些層做weight decay，當weight_decay>0時起作用；split_indices：這個參數的作用是用於加速allreduce過程。
- _get_Ainv_Ginv_Amax_Gmax_list函數用於計算協方差矩陣A/G的逆，並返回求完逆后的矩陣。具體過程是遍歷模型所有層，按層處理，對每一層的協方差矩陣加上正則化項，然后對矩陣進行cholesky分解從而來求逆。當前開源代碼THOR中支持全連接層和卷積層的處理。
def _get_Ainv_Ginv_Amax_Gmax_list(self, gradients, damping_step, matrix_a_allreduce, matrix_g_allreduce, matrix_a_max_allreduce, matrix_g_max_allreduce): """get matrixA inverse list, matrixG inverse list, matrixA_max list, matrixG_max list""" for i in range(len(self.params)): thor_layer_count = self.weight_fim_idx_map[i] conv_layer_count = self.weight_conv_idx_map[i] layer_type = self.weight_layerType_idx_map[i] if layer_type in [Conv, FC, Embedding]: g = gradients[i] matrix_A = self.matrix_A_cov[thor_layer_count] matrix_G = self.matrix_G_cov[thor_layer_count] matrix_A = F.depend(matrix_A, g) matrix_G = F.depend(matrix_G, g) A_shape = self.shape(matrix_A) A_eye = self.eye(A_shape[0], A_shape[0], mstype.float32) G_shape = self.shape(matrix_G) G_eye = self.eye(G_shape[0], G_shape[0], mstype.float32) if layer_type == Conv: ... elif layer_type == FC: matrix_A = matrix_A + damping * A_eye matrix_A_inv = self.cholesky(matrix_A) matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv)
- _get_second_gradients函數用於計算最終參數更新方向，在論文中參數更新方向公式為
,所以代碼實際實現的方式為
，代碼如下
def _get_second_gradients(self, new_grads, damping_step, gradients): """get second gradients for thor""" params_len = len(self.params) for i in range(params_len): ... else: ... elif layer_type == FC: temp_a = self.matrix_A_cov[thor_layer_count] temp_g = self.matrix_G_cov[thor_layer_count] temp_a = self.cast(temp_a, mstype.float16) temp_g = self.cast(temp_g, mstype.float16) g = self.cast(g, mstype.float16) g = self.matmul(temp_g, g) g = self.matmul(g, temp_a) g = self.cast(g, mstype.float32)
construct函數是在網絡訓練過程中會實際執行的內容，該函數中包含了上述兩個函數_get_Ainv_Ginv_Amax_Gmax_list和_get_second_gradients的調用，該函數完成了二階矩陣的計算和梯度更新方向的調整。
def construct(self, gradients): params = self.params moments = self.moments damping_step = self.gather(self.damping, self.cov_step, self.axis) damping_step = self.cast(damping_step, mstype.float32) if self.thor: matrix_A_allreduce = () matrix_G_allreduce = () matrix_A_max_allreduce = () matrix_G_max_allreduce = () matrix_A_allreduce, matrix_G_allreduce, matrix_A_max_allreduce, matrix_G_max_allreduce = \ self._get_Ainv_Ginv_Amax_Gmax_list(gradients, damping_step, matrix_A_allreduce, matrix_G_allreduce, matrix_A_max_allreduce, matrix_G_max_allreduce) #計算A/G的逆 ... new_grads = () for i in range(len(self.params)): ... if self.conv_layer_count > 0:#有卷積層時的處理 ... else: #都是全連接層時的處理 if layer_type == Embedding: ... elif layer_type == FC: temp_a = matrix_A_allreduce[thor_layer_count] temp_g = matrix_G_allreduce[thor_layer_count] fake_A = self.assign(self.matrix_A_cov[thor_layer_count], temp_a) fake_G = self.assign(self.matrix_G_cov[thor_layer_count], temp_g) g = F.depend(g, fake_A)#確保執行順序 g = F.depend(g, fake_G) temp_a = self.cast(temp_a, mstype.float16) temp_g = self.cast(temp_g, mstype.float16) g = self.cast(g, mstype.float16) g = self.matmul(temp_g, g) g = self.matmul(g, temp_a)#將一階方向變為二階方向 g = self.cast(g, mstype.float32) elif layer_type == LayerNorm: g = self._process_layernorm(damping_step, g) new_grads = new_grads + (g,) gradients = new_grads #計算后得到的更新方向 else: #該分支表示使用過時二階信息更新參數 new_grads = () gradients = self._get_second_gradients(new_grads, damping_step, gradients) #調用_get_second_gradients函數計算方向 ...
THOR的實踐應用
在這一節中跟大家分享下THOR的實踐應用，舉了兩個例子分別為ResNet50和BERT，這兩個例子的代碼也已開源，鏈接如下：ResNet50:https://gitee.com/mindspore/mindspore/blob/master/model_zoo/official/cv/resnet/train.pyBERT:https://gitee.com/mindspore/mindspore/blob/master/model_zoo/official/nlp/bert/run_pretrain.py
ResNet50[1]
優化器的調用方式與文中開頭提到的一致，在這個例子中把具體訓練過程給展開了。
首先創建了網絡訓練需要的訓練集和網絡定義為ResNet50；隨后設置THOR所需要用到的超參策略，其他超參值設定可去該目錄下的src/config.py中修改；接着創建THOR優化器，並傳入設置的超參值；然后轉換模型保存二階所需信息；最后就可以訓練網絡了。
from mindspore.nn.optim import Momentum, THOR #引用二階優化器 from src.resnet import resnet50 as resnet from mindspore.train.model import Model ... if __name__ == '__main__': ... #創建網絡訓練過程中的訓練集 dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=True, repeat_num=1, batch_size=config.batch_size, target=target, distribute=args_opt.run_distribute) step_size = dataset.get_dataset_size() #創建resnet50模型 net = resnet(class_num=config.class_num) ... # init lr if cfg.optimizer == "Thor": #設置超參值 from src.lr_generator import get_thor_lr lr = get_thor_lr(0, config.lr_init, config.lr_decay, config.lr_end_epoch, step_size, decay_epochs=39) # define loss, model if target == "Ascend": if args_opt.dataset == "imagenet2012": if not config.use_label_smooth: config.label_smooth_factor = 0.0 loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=config.label_smooth_factor, num_classes=config.class_num) else: loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False) #高層抽象，集成網絡模型的訓練和測試 model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False) if cfg.optimizer == "Thor" and args_opt.dataset == "imagenet2012": from src.lr_generator import get_thor_damping #設置超參damping damping = get_thor_damping(0, config.damping_init, config.damping_decay, 70, step_size) #用於通信時的並行加速 split_indices = [26, 53] #創建THOR優化器 opt = THOR(net, lr, Tensor(damping), config.momentum, config.weight_decay, config.loss_scale, config.batch_size, split_indices=split_indices) #增加計算圖提升性能 model = ConvertModelUtils().convert_to_thor_model(model=model, network=net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False, frequency=config.frequency) ... #訓練網絡 model.train(config.epoch_size - config.pretrain_epoch_size, dataset, callbacks=cb, sink_size=dataset.get_dataset_size(), dataset_sink_mode=dataset_sink_mode)
最后輸入
即可運行腳本啦。
BERT[2]
BERT中步驟與ResNet50差不多。首先創建了網絡訓練需要的訓練集和網絡定義為BERT；隨后設置THOR所需要用到的超參策略，其他超參值設定可去該目錄下的src/config.py中修改；優化器創建時傳入BERT設定的超參值，本例中創建時傳入了：
表示做weight decay操作時排除LN層和FC中的bias參數；然后轉換模型保存二階所需信息；最后就可以訓練網絡了。
from mindspore.nn.optim import Lamb, Momentum, AdamWeightDecay, THOR #引用二階優化器 from src import BertNetworkWithLoss ... def _get_optimizer(args_opt, network): """get bert optimizer, support Lamb, Momentum, AdamWeightDecay.""" if cfg.optimizer == 'Lamb': ... elif cfg.optimizer == "Thor": from src.utils import get_bert_thor_lr, get_bert_thor_damping #設置lr和damping的超參值 lr = get_bert_thor_lr(cfg.Thor.lr_max, cfg.Thor.lr_min, cfg.Thor.lr_power, cfg.Thor.lr_total_steps) damping = get_bert_thor_damping(cfg.Thor.damping_max, cfg.Thor.damping_min, cfg.Thor.damping_power, cfg.Thor.damping_total_steps) split_indices = None #設置並行加速方式 if bert_net_cfg.num_hidden_layers == 12: if bert_net_cfg.use_relative_positions: split_indices = [29, 58, 87, 116, 145, 174, 203, 217] else: split_indices = [28, 55, 82, 109, 136, 163, 190, 205] elif bert_net_cfg.num_hidden_layers == 24: if bert_net_cfg.use_relative_positions: split_indices = [30, 90, 150, 210, 270, 330, 390, 421] else: split_indices = [38, 93, 148, 203, 258, 313, 368, 397] #創建優化器 optimizer = THOR(network, lr, damping, cfg.Thor.momentum, cfg.Thor.weight_decay, cfg.Thor.loss_scale, cfg.batch_size, decay_filter=lambda x: 'layernorm' not in x.name.lower() and 'bias' not in x.name.lower(), split_indices=split_indices) ... return optimizer def run_pretrain(): ... #創建數據集 ds = create_bert_dataset(device_num, rank, args_opt.do_shuffle, args_opt.data_dir, args_opt.schema_dir) #網絡和損失函數創建 net_with_loss = BertNetworkWithLoss(bert_net_cfg, True) ... #加載初始checkpoint if args_opt.load_checkpoint_path: param_dict = load_checkpoint(args_opt.load_checkpoint_path) load_param_into_net(net_with_loss, param_dict) #動態loss縮放 if args_opt.enable_lossscale == "true": ... #固定loss縮放值 else: #反向過程梯度計算過程創建 net_with_grads = BertTrainOneStepCell(net_with_loss, optimizer=optimizer) #創建網絡 model = Model(net_with_grads) #增加計算圖提升性能 model = ConvertModelUtils().convert_to_thor_model(model, network=net_with_grads, optimizer=optimizer, frequency=cfg.Thor.frequency) #網絡訓練 model.train(new_repeat_count, ds, callbacks=callback, dataset_sink_mode=(args_opt.enable_data_sink == "true"), sink_size=args_opt.data_sink_steps) if __name__ == '__main__': set_seed(0)
最后輸入
即可運行腳本啦.至此高階優化器系列的內容就結束啦，該系列總共有三篇文章分別從優化器的背景，MindSpore自研優化器的介紹和MindSpore 高階優化器THOR 的源碼分析&實踐應用這三個內容來跟大家分享，如有不足之處歡迎大家批評指正。同時也歡迎大家到MindSpore開源社區中一起玩耍。
參考文獻：
[1]He K, Zhang X, Ren S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 770-778.
[2]Devlin J, Chang M W, Lee K, et al. Bert: Pre-training of deep bidirectional transformers for language understanding[J]. arXiv preprint arXiv:1810.04805, 2018.