HuggingFace Trainer底层用PyTorch还是TensorFlow？为何出现TF报错？

根据HuggingFace Transformers v4.27.2官方文档，Trainer类是“为大多数标准场景提供功能完备的PyTorch训练API”。但实际使用时，出现了如下TensorFlow相关提示信息：

tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.

请问：

1. Trainer类底层实际使用PyTorch还是TensorFlow？
2. 能否切换为仅使用PyTorch？在TrainingArguments中未找到相关参数。
3. 为何脚本会输出TensorFlow相关信息，Trainer不应仅使用PyTorch吗？

附上使用的源代码：

from transformers import GPT2Tokenizer
from transformers import GPT2LMHeadModel
from transformers import TextDataset
from transformers import DataCollatorForLanguageModeling
from transformers import Trainer
from transformers import TrainingArguments

import torch

# Load the GPT-2 tokenizer and LM head model
tokenizer    = GPT2Tokenizer.from_pretrained('gpt2')
lmhead_model = GPT2LMHeadModel.from_pretrained('gpt2')

# Load the training dataset and divide blocksize
train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path='./datasets/tinyshakespeare.txt',
    block_size=64
)

# Create a data collator for preprocessing batches
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

# Defining the training arguments
training_args = TrainingArguments(
    output_dir='./models/tinyshakespeare', # output directory for checkpoints
    overwrite_output_dir=True,             # overwrite any existing content

    per_device_train_batch_size=4,         # sample batch size for training
    dataloader_num_workers=1,              # number of workers for dataloader
    max_steps=100,                         # maximum number of training steps
    save_steps=50,                         # after # steps checkpoints are saved
    save_total_limit=5,                    # maximum number of checkpoints to save

    prediction_loss_only=True,             # only compute loss during prediction
    learning_rate=3e-4,                    # learning rate
    fp16=False,                            # use 16-bit (mixed) precision

    optim='adamw_torch',                   # define the optimizer for training
    lr_scheduler_type='linear',            # define the learning rate scheduler

    logging_steps=5,                       # after # steps logs are printed
    report_to='none',                      # report to wandb, tensorboard, etc.
)

if __name__ == '__main__':
    torch.multiprocessing.freeze_support()

    trainer = Trainer(
        model=lmhead_model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
    )

    trainer.train()

1. Trainer类底层实际使用的框架 #

Trainer类底层完全基于PyTorch实现，官方文档描述准确。你的代码中加载的GPT2LMHeadModel是PyTorch版本模型，TrainingArguments指定的optim='adamw_torch'是PyTorch原生优化器，整个训练流程全程由PyTorch驱动。

2. 如何确保仅使用PyTorch #

不需要在TrainingArguments中额外设置参数，满足以下两点即可完全使用PyTorch：

• 加载的是PyTorch版本模型（默认情况下，from_pretrained会自动匹配当前环境框架，安装PyTorch后优先加载PyTorch模型）
• 代码中未显式导入TensorFlow相关模块

如果想彻底避免TensorFlow被加载，可选择两种方式：

• 卸载环境中的TensorFlow包：执行pip uninstall tensorflow
• 在脚本最开头添加环境变量，阻止Transformers加载TensorFlow模块：

  import os
  os.environ["TRANSFORMERS_NO_TF"] = "1"
  

3. 为何会出现TensorFlow的提示信息 #

出现该提示的核心原因是你的Python环境中安装了TensorFlow包。Transformers库初始化时会自动检测环境中可用的框架（包括TensorFlow和PyTorch），这个检测过程会触发TensorFlow的初始化逻辑，进而输出这条CPU优化相关的日志。即便你没有在代码中显式调用TensorFlow，只要环境中存在该包，就可能触发这个加载行为。

内容的提问来源于stack exchange，提问作者AlanSTACK