Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Enable Auto-Mixed-Precision & TensorCore for AlexNet on GPU

This commit is contained in:
lishanni513 2021-06-16 19:12:15 +08:00
parent 1baebef648
commit bddfd9ac65
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/_extends/graph_kernel/expanders/conv2d.py
View File

@ -21,7 +21,7 @@ from ._utils import Expander, ExpanderInfoValidator as VLD
M_ALIGN = 16
N_ALIGN = 16
K_ALIGN = 8
K_LIMIT = 4096
K_LIMIT = 800
MNK_LIMIT = 3 * (10 ** 10)
N0_CHANNEL_ALIGN = 16
N1_CHANNEL_ALIGN = 16

4
model_zoo/official/cv/alexnet/train.py
View File

@ -58,6 +58,7 @@ def train_alexnet():
context.set_context(save_graphs=False)
if device_target == "GPU":
context.set_context(enable_graph_kernel=True)
context.set_context(graph_kernel_flags="--enable_cluster_ops=MatMul")
device_num = get_device_num()
if config.dataset_name == "cifar10":
@ -124,7 +125,8 @@ def train_alexnet():
model = Model(network, loss_fn=loss, optimizer=opt, metrics=metrics, amp_level="O2", keep_batchnorm_fp32=False,
loss_scale_manager=loss_scale_manager)
elif device_target == "GPU":
model = Model(network, loss_fn=loss, optimizer=opt, metrics=metrics, loss_scale_manager=loss_scale_manager)
model = Model(network, loss_fn=loss, optimizer=opt, metrics=metrics, amp_level="O2",
loss_scale_manager=loss_scale_manager)
else:
raise ValueError("Unsupported platform.")