Spaces:
Sleeping
Sleeping
File size: 3,555 Bytes
757ed1c |
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 75 76 77 78 |
# -*- coding: utf-8 -*-
#
# Developed by Haozhe Xie <[email protected]>
import torch
class Refiner(torch.nn.Module):
def __init__(self, cfg):
super(Refiner, self).__init__()
self.cfg = cfg
# Layer Definition
self.layer1 = torch.nn.Sequential(
torch.nn.Conv3d(1, 32, kernel_size=4, padding=2),
torch.nn.BatchNorm3d(32),
torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE),
torch.nn.MaxPool3d(kernel_size=2)
)
self.layer2 = torch.nn.Sequential(
torch.nn.Conv3d(32, 64, kernel_size=4, padding=2),
torch.nn.BatchNorm3d(64),
torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE),
torch.nn.MaxPool3d(kernel_size=2)
)
self.layer3 = torch.nn.Sequential(
torch.nn.Conv3d(64, 128, kernel_size=4, padding=2),
torch.nn.BatchNorm3d(128),
torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE),
torch.nn.MaxPool3d(kernel_size=2)
)
self.layer4 = torch.nn.Sequential(
torch.nn.Linear(8192, 2048),
torch.nn.ReLU()
)
self.layer5 = torch.nn.Sequential(
torch.nn.Linear(2048, 8192),
torch.nn.ReLU()
)
self.layer6 = torch.nn.Sequential(
torch.nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1),
torch.nn.BatchNorm3d(64),
torch.nn.ReLU()
)
self.layer7 = torch.nn.Sequential(
torch.nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1),
torch.nn.BatchNorm3d(32),
torch.nn.ReLU()
)
self.layer8 = torch.nn.Sequential(
torch.nn.ConvTranspose3d(32, 1, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1),
torch.nn.Sigmoid()
)
def forward(self, coarse_volumes):
volumes_32_l = coarse_volumes.view((-1, 1, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX))
# print(volumes_32_l.size()) # torch.Size([batch_size, 1, 32, 32, 32])
volumes_16_l = self.layer1(volumes_32_l)
# print(volumes_16_l.size()) # torch.Size([batch_size, 32, 16, 16, 16])
volumes_8_l = self.layer2(volumes_16_l)
# print(volumes_8_l.size()) # torch.Size([batch_size, 64, 8, 8, 8])
volumes_4_l = self.layer3(volumes_8_l)
# print(volumes_4_l.size()) # torch.Size([batch_size, 128, 4, 4, 4])
flatten_features = self.layer4(volumes_4_l.view(-1, 8192))
# print(flatten_features.size()) # torch.Size([batch_size, 2048])
flatten_features = self.layer5(flatten_features)
# print(flatten_features.size()) # torch.Size([batch_size, 8192])
volumes_4_r = volumes_4_l + flatten_features.view(-1, 128, 4, 4, 4)
# print(volumes_4_r.size()) # torch.Size([batch_size, 128, 4, 4, 4])
volumes_8_r = volumes_8_l + self.layer6(volumes_4_r)
# print(volumes_8_r.size()) # torch.Size([batch_size, 64, 8, 8, 8])
volumes_16_r = volumes_16_l + self.layer7(volumes_8_r)
# print(volumes_16_r.size()) # torch.Size([batch_size, 32, 16, 16, 16])
volumes_32_r = (volumes_32_l + self.layer8(volumes_16_r)) * 0.5
# print(volumes_32_r.size()) # torch.Size([batch_size, 1, 32, 32, 32])
return volumes_32_r.view((-1, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX))
|