Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
FastMoE
Commits
d9ca437a
Commit
d9ca437a
authored
Dec 07, 2020
by
Jiezhong Qiu
Browse files
update test code
parent
39996fef
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
4 additions
and
4 deletions
+4
-4
pytorch/topk.py
pytorch/topk.py
+4
-4
No files found.
pytorch/topk.py
View file @
d9ca437a
...
...
@@ -12,25 +12,25 @@ output = torch.rand(16, 512, 512).cuda()
torch
.
cuda
.
synchronize
()
start
=
time
.
time
()
_
,
pred
=
output
.
topk
(
k
=
1
,
dim
=-
1
,
largest
=
True
,
sorted
=
Tru
e
)
_
,
pred
=
output
.
topk
(
k
=
1
,
dim
=-
1
,
largest
=
True
,
sorted
=
Fals
e
)
torch
.
cuda
.
synchronize
()
print
(
"torch.top1 Time :{}"
.
format
(
time
.
time
()
-
start
))
torch
.
cuda
.
synchronize
()
start
=
time
.
time
()
_
,
pred_
=
my_topk
(
output
,
k
=
1
,
inplace
=
Fals
e
)
_
,
pred_
=
my_topk
(
output
,
k
=
1
,
inplace
=
Tru
e
)
torch
.
cuda
.
synchronize
()
print
(
"my top1 Time :{}"
.
format
(
time
.
time
()
-
start
))
torch
.
cuda
.
synchronize
()
start
=
time
.
time
()
_
,
pred
=
output
.
topk
(
k
=
2
,
dim
=-
1
,
largest
=
True
,
sorted
=
Tru
e
)
_
,
pred
=
output
.
topk
(
k
=
2
,
dim
=-
1
,
largest
=
True
,
sorted
=
Fals
e
)
torch
.
cuda
.
synchronize
()
print
(
"torch.top2 Time :{}"
.
format
(
time
.
time
()
-
start
))
torch
.
cuda
.
synchronize
()
start
=
time
.
time
()
_
,
pred_
=
my_topk
(
output
,
k
=
2
,
inplace
=
Fals
e
)
_
,
pred_
=
my_topk
(
output
,
k
=
2
,
inplace
=
Tru
e
)
torch
.
cuda
.
synchronize
()
print
(
"my top2 Time :{}"
.
format
(
time
.
time
()
-
start
))
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment