Spaces:
Sleeping
Sleeping
File size: 1,181 Bytes
78d2329 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | import torch
import time
from layer import KNNAttention, TransformerBlock, PlainPointTransformer, SubsampleBlock
device = torch.device('cuda')
def test():
bs = 4
npts = 1024
len_xyz = 3
feat_dims = 64
num_classes = 23
coord = torch.rand(bs * npts, len_xyz).cuda()
feat = torch.rand(bs * npts, feat_dims).cuda()
offset = [npts * i for i in range(1, bs + 1)]
offset = torch.tensor(offset).cuda()
# data_dict = dict(
# coord = coord,
# feat = feat,
# offset = offset
# )
# model = PointTransformerSeg26().cuda()
# model = KNNAttention(feat_dims, num_samples=16).cuda()
# model = TransformerBlock(feat_dims).cuda()
# model = PlainPointTransformer(feat_dims, num_blocks=2).cuda()
model = SubsampleBlock(feat_dims, feat_dims).cuda()
print(model)
# count time
# count = 100
# torch.cuda.synchronize()
# start = time.time()
# for _ in range(count):
# out = model((coord, feat, offset))
# torch.cuda.synchronize()
# print(time.time() - start)
out = model((coord, feat, offset))
print(out[0].shape)
# print(out.shape)
test()
|