mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Move opt_einsum import back out of class __init__

Browse Source
This commit is contained in:
Ross Wightman 2025-01-27 14:03:25 -08:00 committed by Ross Wightman
parent 9ab5464e4d
commit 67ef6f0a92
1 changed files with 16 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
timm/optim/kron.py
View File

@ -14,7 +14,15 @@ from typing import Any, Callable, Dict, Optional, Tuple, Union
import numpy as np import numpy as np
import torch import torch
try:
# NOTE opt_einsum needed to avoid blowing up memory with einsum ops
import opt_einsum
import torch.backends.opt_einsum
torch.backends.opt_einsum.enabled = True
torch.backends.opt_einsum.strategy = "auto-hq"
has_opt_einsum = True
except ImportError:
has_opt_einsum = False
try: try:
torch._dynamo.config.cache_size_limit = 1_000_000 torch._dynamo.config.cache_size_limit = 1_000_000
@ -92,14 +100,8 @@ class Kron(torch.optim.Optimizer):
flatten_dim: bool = False, flatten_dim: bool = False,
deterministic: bool = False, deterministic: bool = False,
): ):
try: if not has_opt_einsum:
# NOTE opt_einsum needed to avoid blowing up memory with einsum ops warnings.warn("It is highly recommended to have 'opt_einsum' installed for this optimizer.")
import opt_einsum
opt_einsum.enabled = True
opt_einsum.strategy = "auto-hq"
import torch.backends.opt_einsum
except ImportError:
warnings.warn("It is highly recommended to have 'opt_einsum' installed for this optimizer." )
if not 0.0 <= lr: if not 0.0 <= lr:
raise ValueError(f"Invalid learning rate: {lr}") raise ValueError(f"Invalid learning rate: {lr}")