pytorch
1
2
3
4
5
6import unittest
7from caffe2.proto import caffe2_pb2
8from caffe2.python import core, dyndep, workspace
9
10dyndep.InitOpsLibrary("@/caffe2/caffe2/contrib/prof:cuda_profile_ops")
11
12
13class CudaProfileOpsTest(unittest.TestCase):
14@unittest.skipIf(workspace.NumCudaDevices() < 1, "Need at least 1 GPU")
15def test_run(self):
16net = core.Net("net")
17net.CudaProfileInitialize([], [], output="/tmp/cuda_profile_test")
18net.CudaProfileStart([], [])
19with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
20net.ConstantFill([], ["out"], shape=[1, 3, 244, 244])
21net.CudaProfileStop([], [])
22
23workspace.CreateNet(net)
24workspace.RunNet(net)
25