Cian-H/symbolic_nn_tests
1
0
Fork 0
mirror of https://github.com/Cian-H/symbolic_nn_tests.git synced 2025-12-22 22:22:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added automatic enabling of tensor cores

Browse Source
This commit is contained in:
Cian Hughes
2024-06-10 12:58:19 +01:00
parent 08fe32302c
commit 888fe9f2e8
3 changed files with 401 additions and 399 deletions
Download Patch File Download Diff File Expand all files Collapse all files

791
poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

7
symbolic_nn_tests/__main__.py
View File

@@ -27,6 +27,13 @@ def main(
bool, typer.Option(help="Whether or not to log via Weights & Biases") bool, typer.Option(help="Whether or not to log via Weights & Biases")
] = True, ] = True,
): ):
import torch
# Enable tensor cores for compatible GPUs
for i in torch.cuda.device_count():
if torch.cuda.get_device_properties(i).major > 6:
torch.set_float32_matmul_precision("medium")
for i, n in enumerate(experiments, start=1): for i, n in enumerate(experiments, start=1):
j = n - 1 j = n - 1
experiment = EXPERIMENTS[j].run experiment = EXPERIMENTS[j].run

2
symbolic_nn_tests/experiment2/model.py
View File

@@ -49,7 +49,7 @@ def get_singleton_dataset():
from symbolic_nn_tests.experiment2.dataset import collate, pubchem from symbolic_nn_tests.experiment2.dataset import collate, pubchem
return create_dataset( return create_dataset(
dataset=pubchem, collate_fn=collate, batch_size=128, shuffle=True dataset=pubchem, collate_fn=collate, batch_size=512, shuffle=True
) )