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

dml autocast

Browse Source
This commit is contained in:
Vladimir Mandic
2023-05-14 13:24:59 -04:00
parent 618a1703ae
commit 5134471bc8
4 changed files with 88 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/dml/__init__.py
+19 -2
View File
@@ -1,12 +1,15 @@
# pylint: disable=no-member,no-self-argument
# pylint: disable=no-member,no-self-argument,no-method-argument
import torch
import torch_directml # pylint: disable=import-error
import modules.dml.hijack
import modules.dml.amp as amp
from .optimizer.unknown import UnknownOptimizer
class DirectML():
_is_autocast_enabled = False
_autocast_dtype = torch.float16
def get_optimizer(device: torch.device):
assert device.type == 'privateuseone'
try:
@@ -27,5 +30,19 @@ class DirectML():
optimizer = DirectML.get_optimizer(device)
return optimizer.memory_stats(device.index)
def get_autocast_gpu_dtype():
return DirectML._autocast_dtype
def set_autocast_gpu_dtype(dtype):
DirectML._autocast_dtype = dtype
def is_autocast_enabled():
return DirectML._is_autocast_enabled
def set_autocast_enabled(enabled: bool):
DirectML._is_autocast_enabled = enabled
# Alternative of torch.cuda for DirectML.
DirectML.amp = amp
torch.dml = DirectML