File size: 923 Bytes
854f0d0 216282e 854f0d0 |
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 |
import torch
def generate_grid(n_vox, interval):
"""
generate grid
if 3D volume, grid[:,:,x,y,z] = (x,y,z)
:param n_vox:
:param interval:
:return:
"""
with torch.no_grad():
# Create voxel grid
grid_range = [torch.arange(0, n_vox[axis], interval) for axis in range(3)]
grid = torch.stack(torch.meshgrid(grid_range[0], grid_range[1], grid_range[2], indexing="ij")) # 3 dx dy dz
# ! don't create tensor on gpu; imbalanced gpu memory in ddp mode
grid = grid.unsqueeze(0).type(torch.float32) # 1 3 dx dy dz
return grid
if __name__ == "__main__":
import torch.nn.functional as F
grid = generate_grid([5, 6, 8], 1)
pts = 2 * torch.tensor([1, 2, 3]) / (torch.tensor([5, 6, 8]) - 1) - 1
pts = pts.view(1, 1, 1, 1, 3)
pts = torch.flip(pts, dims=[-1])
sampled = F.grid_sample(grid, pts, mode='nearest')
print(sampled)
|