Commit 90576890 authored by rusty1s's avatar rusty1s
Browse files

added old gpu impl

parent dd9778c8
import time import time
import sys
import torch import torch
from torch_spline_conv.functions.ffi import spline_basis_forward from torch.autograd import Variable
from torch_spline_conv import spline_conv
n, d = 9999999, 5 sys.path.insert(0, '../../pytorch_geometric')
pseudo = torch.FloatTensor(n, d).uniform_(0, 1)
from torch_geometric.nn.modules import SplineConv # noqa
n = 600
m_in = 64
m_out = 64
d = 4
x = torch.FloatTensor(n, m_in).uniform_(-1, 1)
row = torch.arange(0, n).view(-1, 1).repeat(1, n).view(-1).long()
col = torch.arange(0, n).repeat(n).long()
edge_index = torch.stack([row, col], dim=0)
pseudo = torch.FloatTensor(n * n, d).uniform_(0, 1)
kernel_size = torch.LongTensor(d).fill_(5) kernel_size = torch.LongTensor(d).fill_(5)
is_open_spline = torch.ByteTensor(d).fill_(1) is_open_spline = torch.ByteTensor(d).fill_(1)
K = kernel_size.prod() K = kernel_size.prod()
weight = torch.FloatTensor(K, m_in, m_out).uniform_(-1, 1)
t = time.perf_counter() # t = time.perf_counter()
basis, index = spline_basis_forward(1, pseudo, kernel_size, is_open_spline, K) # out = spline_conv(x, edge_index, pseudo, weight, kernel_size, is_open_spline)
t = time.perf_counter() - t # t = time.perf_counter() - t
print('CPU:', t) # print('CPU:', t)
x = x.cuda()
edge_index = edge_index.cuda()
pseudo = pseudo.cuda() pseudo = pseudo.cuda()
weight = weight.cuda()
kernel_size = kernel_size.cuda() kernel_size = kernel_size.cuda()
is_open_spline = is_open_spline.cuda() is_open_spline = is_open_spline.cuda()
out = spline_conv(x, edge_index, pseudo, weight, kernel_size, is_open_spline)
torch.cuda.synchronize() torch.cuda.synchronize()
t = time.perf_counter() t = time.perf_counter()
basis, index = spline_basis_forward(1, pseudo, kernel_size, is_open_spline, K) out = spline_conv(x, edge_index, pseudo, weight, kernel_size, is_open_spline)
torch.cuda.synchronize() torch.cuda.synchronize()
t = time.perf_counter() - t t = time.perf_counter() - t
print('GPU:', t) print('GPU:', t)
conv = SplineConv(m_in, m_out, d, kernel_size, is_open_spline.long()).cuda()
adj = {'indices': edge_index, 'values': pseudo, 'size': torch.Size([n, n, d])}
x = Variable(x)
conv(adj, x)
torch.cuda.synchronize()
t = time.perf_counter()
conv(adj, x)
torch.cuda.synchronize()
t = time.perf_counter() - t
print('GPU old:', t)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment