a oyd= @sddlZddlZddlZddlZddlmZddlmZddlmZddZ ddZ Gdd d ej Z d d Z ed kred gdZee edS)N)cuda) BertTokenizercCs>td}||dd|dd}|d}|d}|d}|||fS)Nbert-base-chineseTpt)padding truncation max_lengthreturn_tensors input_idstoken_type_idsattention_mask)rfrom_pretrained)max_lentext tokenizerr r r r3/mnt/chromeos/MyFiles/work/nlp_tsa/TSA/inference.pyencoder s rcCs|||txtd|\}}}||||||}}}||||}tjj| }|ddWdS1s0YdS)Nir) toevaltorchno_gradrnn functionalsoftmaxdetachcpunumpytolist)modeldevicerr r r Zout_putprobsrrrpredicts  " r#cs$eZdZfddZddZZS)BertClassificationModelcsHtt|d}tj||_|jD] }d|_q*t dd|_ dS)NrTi) superr$__init__ transformers BertModelr bert parameters requires_gradrLineardense)selfZpretrained_weightsparam __class__rrr'(s z BertClassificationModel.__init__cCs&|j|||d}|d}||}|S)N)r r r r)r*r.)r/r r r Z bert_outputZbert_cls_hidden_stateZ linear_outputrrrforward0s zBertClassificationModel.forward)__name__ __module__ __qualname__r'r3 __classcell__rrr1rr$'s r$c Csttrdnd}d}tj|t|d}tjt|t|gtd}t |D]T\}}t |D]B\}}||kr~d|||<q`|d|} t ||| } | |||<q`qP|S)NrrzTSA/bert_model.pkl) map_location)dtyperu 是否包含 ) rr!r is_availableloadnpzeroslenfloat enumerater#) topicsr! load_pathr matrixiZi_textjZj_texttestoutputsrrrinference_matrix6s  rH__main__yes)u{在本次报告中我将介绍分布式并行加速算法模型架构内存和计算优化以及集群架构等关键技术uW在现代机器学习任务中大模型训练已成为解决复杂问题的重要手段ul首先分布式并行加速策略包括数据并行模型并行流水线并行和张量并行等四种方式uH选择合适的集群架构是实现大模型的分布式训练的关键uf这些策略帮助我们将训练数据和模型分布到多个设备上以加速大模型训练过程)osrr<r(rtorch.nnrrrrr#Moduler$rHr4printrArrrrs