TorchAdam warmup...[DONE]	222.59ms
TorchAdam...[DONE]	1,356.01ms
MyAdam warmup...[DONE]	119.15ms
MyAdam...[DONE]	1,192.89ms

19importtorch20importtorch.nnasnn21fromlabml\_helpers.deviceimportDeviceInfo22fromtorch.optimimportAdamasTorchAdam2324fromlabmlimportmonit25fromlabml\_nn.optimizers.adamimportAdamasMyAdam26fromlabml\_nn.optimizers.mnist\_experimentimportModel

29deftest():30device\_info=DeviceInfo(use\_cuda=True,cuda\_device=0)31print(device\_info)32inp=torch.randn((64,1,28,28),device=device\_info.device)33target=torch.ones(64,dtype=torch.long,device=device\_info.device)34loss\_func=nn.CrossEntropyLoss()35model=Model().to(device\_info.device)36my\_adam=MyAdam(model.parameters())37torch\_adam=TorchAdam(model.parameters())38loss=loss\_func(model(inp),target)39loss.backward()40withmonit.section('MyAdam warmup'):41foriinrange(100):42my\_adam.step()43withmonit.section('MyAdam'):44foriinrange(1000):45my\_adam.step()46withmonit.section('TorchAdam warmup'):47foriinrange(100):48torch\_adam.step()49withmonit.section('TorchAdam'):50foriinrange(1000):51torch\_adam.step()525354if\_\_name\_\_=='\_\_main\_\_':55test()