search

glistering96 / AlphaRouter

2 stars 0 forks source link

Mixed precision training #9

Closed glistering96 closed 1 year ago

glistering96 commented 1 year ago

Apply the torch auto cast mixed precision for faster training

glistering96 commented 1 year ago

params

'env_type' : 'tsp',
'num_nodes' : 50,
'epochs' : 1000,
'num_episode' : 1024, (= batch size)
'qkv_dim' :  32,

Result

image

with amp: 2.98 min without amp: 3.839 min

glistering96 commented 1 year ago

Closing the issue as PR is accepted