File size: 2,843 Bytes
913ec88 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
import torch.nn as nn
import torch.nn.functional as F
from model.common import (MLP, Custom1x1Subm3d, ResidualBlock, UBlock)
import functools
from collections import OrderedDict
import spconv.pytorch as spconv
from spconv.pytorch.modules import SparseModule
# this file only provides the 2 modules used in VQVAE
__all__ = ['Encoder', 'Decoder',]
class Encoder(nn.Module):
def __init__(self, input_dim, hidden_dim, spconv_channels, num_blocks):
super(Encoder, self).__init__()
block = ResidualBlock
norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(
input_dim, spconv_channels, kernel_size=3, padding=1, bias=False, indice_key='subm1'))
block_channels = [spconv_channels * (i + 1) for i in range(num_blocks)]
self.unet = UBlock(block_channels, norm_fn, 2, block, indice_key_id=1)
self.output_layer = spconv.SparseSequential(norm_fn(spconv_channels), nn.ReLU())
def forward(self, voxel_feats, voxel_coords, spatial_shape):
# spconv encode
batch_size = len(voxel_coords)
voxel_feats_total = []
voxel_coords_total = []
voxel_batch_id_total = []
for i in range(batch_size):
batch_col = torch.zeros((voxel_coords[i].shape[0], 1), device=voxel_coords[i].device, dtype=torch.int32)
voxel_batch_id_total.append(batch_col)
voxel_coords_total.append(voxel_coords[i].int())
voxel_coords[i] = torch.cat([batch_col, voxel_coords[i].int()], dim=1)
input = spconv.SparseConvTensor(voxel_feats[i], voxel_coords[i], spatial_shape, 1)
output = self.input_conv(input)
output = self.unet(output)
voxel_feats_total.append(output)
return voxel_feats_total, voxel_coords_total, voxel_batch_id_total
class Decoder(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
super(Decoder, self).__init__()
self.num_layers = num_layers
self.start_layer = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim)
)
self.layers = nn.ModuleList([
nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim)
) for _ in range(num_layers - 2)
])
self.final_layer = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, x):
x = self.start_layer(x)
for layer in self.layers:
x = layer(x)
return self.final_layer(x) |