Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixes to prep for weights_only default flip #2514

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
+27 −2
Draft

Fixes to prep for weights_only default flip #2514

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e022789
Fixes to prep for weights_only default flip
mikaylagawarecki Oct 17, 2024
df57893
fix os import
mikaylagawarecki Oct 17, 2024
eada452
Remove unnecessary import
mikaylagawarecki Oct 18, 2024
bedde13
Remove weights_only
mikaylagawarecki Oct 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion torchbenchmark/models/functorch_maml_omniglot/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
from pathlib import Path
from typing import Tuple

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
Expand Down Expand Up @@ -73,7 +75,14 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
self.model = net

root = str(Path(__file__).parent.parent)
self.meta_inputs = torch.load(f"{root}/maml_omniglot/batch.pt")
with torch.serialization.safe_globals([
np.core.multiarray._reconstruct,
np.ndarray,
np.dtype,
np.dtypes.Float32DType,
np.dtypes.Int64DType,
]):
self.meta_inputs = torch.load(f"{root}/maml_omniglot/batch.pt")
self.meta_inputs = tuple(
[torch.from_numpy(i).to(self.device) for i in self.meta_inputs]
)
Expand Down
10 changes: 9 additions & 1 deletion torchbenchmark/models/maml_omniglot/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from typing import Tuple

import higher
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
Expand Down Expand Up @@ -79,7 +80,14 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
self.model = net

root = str(Path(__file__).parent)
self.meta_inputs = torch.load(f"{root}/batch.pt")
with torch.serialization.safe_globals([
np.core.multiarray._reconstruct,
np.ndarray,
np.dtype,
np.dtypes.Float32DType,
np.dtypes.Int64DType,
]):
self.meta_inputs = torch.load(f"{root}/batch.pt")
self.meta_inputs = tuple(
[torch.from_numpy(i).to(self.device) for i in self.meta_inputs]
)
Expand Down
8 changes: 8 additions & 0 deletions torchbenchmark/models/opacus_cifar10/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from typing import Tuple

import os

import torch
import torch.nn as nn
import torch.optim as optim
Expand Down Expand Up @@ -29,7 +31,13 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
)

self.model = models.resnet18(num_classes=10)
prev_wo_envvar = os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", None)
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have access to the torch.load call within opacus here, updating opacus also won't help because torchbench doesn't depend on a nightly of opacus but rather stable release

os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
self.model = ModuleValidator.fix(self.model)
if prev_wo_envvar is None:
del os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"]
else:
os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = prev_wo_envvar
self.model = self.model.to(device)

# Cifar10 images are 32x32 and have 10 classes
Expand Down
Loading