Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Implement memory estimation for AMDGPUs.

Browse Source 
Stable.
This commit is contained in:
Seunghoon Lee
2023-04-26 17:44:32 +09:00
parent df0e89be48
commit d2d5011bd3
10 changed files with 47 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/dml/__init__.py
+7 -7
View File
@@ -3,27 +3,27 @@ import torch_directml
import modules.dml.hijack
from modules.dml.optimizer.unknown import UnknownOptimizer
from .optimizer.unknown import UnknownOptimizer
class DirectML():
def get_optimizer(self, device: torch.device):
def get_optimizer(device: torch.device):
assert(device.type == 'privateuseone')
try:
device_name = torch_directml.device_name(device.index)
if 'NVIDIA' in device_name or 'GeForce' in device_name:
from modules.dml.optimizer.nvidia import nVidiaOptimizer as optimizer
from .optimizer.nvidia import nVidiaOptimizer as optimizer
elif 'AMD' in device_name or 'Radeon' in device_name:
from modules.dml.optimizer.amd import AMDOptimizer as optimizer
from .optimizer.amd import AMDOptimizer as optimizer
elif 'Intel' in device_name:
from modules.dml.optimizer.intel import IntelOptimizer as optimizer
from .optimizer.intel import IntelOptimizer as optimizer
else:
return UnknownOptimizer
return optimizer
except:
return UnknownOptimizer
def memory_stats(self, device: torch.device):
optimizer = self.get_optimizer(device)
def memory_stats(device: torch.device):
optimizer = DirectML.get_optimizer(device)
return optimizer.memory_stats(device.index)
# Alternative of torch.cuda for DirectML.