SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Set weights_only in torch.load (#36991)

Browse Source
This commit is contained in:
cyyever
2025-03-27 22:55:50 +08:00
committed by GitHub
parent de77f5b1ec
commit 41a0e58e5b
28 changed files with 64 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/flax/vision/run_image_classification.py
View File

@@ -36,12 +36,12 @@ import optax
# for dataset and preprocessing
import torch
import torchvision
import torchvision.transforms as transforms
from flax import jax_utils
from flax.jax_utils import pad_shard_unpad, unreplicate
from flax.training import train_state
from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
from huggingface_hub import HfApi
from torchvision import transforms
from tqdm import tqdm
import transformers

2
examples/legacy/multiple_choice/utils_multiple_choice.py
View File

@@ -113,7 +113,7 @@ if is_torch_available():
with FileLock(lock_path):
if os.path.exists(cached_features_file) and not overwrite_cache:
logger.info(f"Loading features from cached file {cached_features_file}")
self.features = torch.load(cached_features_file)
self.features = torch.load(cached_features_file, weights_only=True)
else:
logger.info(f"Creating features from dataset file at {data_dir}")
label_list = processor.get_labels()

2
examples/legacy/pytorch-lightning/run_glue.py
View File

@@ -81,7 +81,7 @@ class GLUETransformer(BaseTransformer):
cached_features_file = self._feature_file(mode)
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
features = torch.load(cached_features_file, weights_only=True)
all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)

4
examples/legacy/pytorch-lightning/run_ner.py
View File

@@ -63,7 +63,7 @@ class NERTransformer(BaseTransformer):
cached_features_file = self._feature_file(mode)
if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
features = torch.load(cached_features_file, weights_only=True)
else:
logger.info("Creating features from dataset file at %s", args.data_dir)
examples = self.token_classification_task.read_examples_from_file(args.data_dir, mode)
@@ -89,7 +89,7 @@ class NERTransformer(BaseTransformer):
"Load datasets. Called after prepare data."
cached_features_file = self._feature_file(mode)
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
features = torch.load(cached_features_file, weights_only=True)
all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
if features[0].token_type_ids is not None:

6
examples/legacy/question-answering/run_squad.py
View File

@@ -105,8 +105,8 @@ def train(args, train_dataset, model, tokenizer):
os.path.join(args.model_name_or_path, "scheduler.pt")
):
# Load in optimizer and scheduler states
optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))
optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt"), weights_only=True))
scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt"), weights_only=True))
if args.fp16:
try:
@@ -417,7 +417,7 @@ def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=Fal
# Init features and dataset from cache if it exists
if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Loading features from cached file %s", cached_features_file)
features_and_dataset = torch.load(cached_features_file)
features_and_dataset = torch.load(cached_features_file, weights_only=True)
features, dataset, examples = (
features_and_dataset["features"],
features_and_dataset["dataset"],

2
examples/legacy/run_swag.py
View File

@@ -244,7 +244,7 @@ def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=Fal
)
if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
features = torch.load(cached_features_file, weights_only=True)
else:
logger.info("Creating features from dataset file at %s", input_file)
examples = read_swag_examples(input_file)

2
examples/legacy/seq2seq/convert_model_to_fp16.py
View File

@@ -22,7 +22,7 @@ from tqdm import tqdm
def convert(src_path: str, map_location: str = "cpu", save_path: Union[str, None] = None) -> None:
"""Convert a pytorch_model.bin or model.pt file to torch.float16 for faster downloads, less disk space."""
state_dict = torch.load(src_path, map_location=map_location)
state_dict = torch.load(src_path, map_location=map_location, weights_only=True)
for k, v in tqdm(state_dict.items()):
if not isinstance(v, torch.Tensor):
raise TypeError("FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin")

2
examples/legacy/token-classification/utils_ner.py
View File

@@ -242,7 +242,7 @@ if is_torch_available():
with FileLock(lock_path):
if os.path.exists(cached_features_file) and not overwrite_cache:
logger.info(f"Loading features from cached file {cached_features_file}")
self.features = torch.load(cached_features_file)
self.features = torch.load(cached_features_file, weights_only=True)
else:
logger.info(f"Creating features from dataset file at {data_dir}")
examples = token_classification_task.read_examples_from_file(data_dir, mode)