38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import torch
|
|
from torch_geometric.data import Data
|
|
|
|
|
|
class BatchToData(object):
|
|
def __init__(self, transforms=None):
|
|
super(BatchToData, self).__init__()
|
|
self.transforms = transforms if transforms else lambda x: x
|
|
|
|
def __call__(self, batch_dict):
|
|
# Convert to torch_geometric.data.Data type
|
|
|
|
batch_pos = batch_dict['pos']
|
|
batch_norm = batch_dict.get('norm', None)
|
|
batch_y = batch_dict.get('y', None)
|
|
batch_y_c = batch_dict.get('y_c', None)
|
|
|
|
batch_size, num_points, _ = batch_pos.shape # (batch_size, num_points, 3)
|
|
|
|
batch_size, N, _ = batch_pos.shape # (batch_size, num_points, 3)
|
|
pos = batch_pos.view(batch_size * N, -1)
|
|
norm = batch_norm.view(batch_size * N, -1) if batch_norm is not None else batch_norm
|
|
|
|
batch_y_l = batch_y.view(batch_size * N, -1) if batch_y is not None else batch_y
|
|
batch_y_c = batch_y_c.view(batch_size * N, -1) if batch_y_c is not None else batch_y_c
|
|
|
|
batch = torch.zeros((batch_size, num_points), device=pos.device, dtype=torch.long)
|
|
for i in range(batch_size):
|
|
batch[i] = i
|
|
batch = batch.view(-1)
|
|
|
|
data = Data()
|
|
data.norm, data.pos, data.batch, data.yl, data.yc = norm, pos, batch, batch_y_l, batch_y_c
|
|
|
|
data = self.transforms(data)
|
|
|
|
return data
|