File size: 2,677 Bytes
92f0e98 |
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 79 80 |
import torch.nn as nn
class LinearBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super(LinearBlock, self).__init__()
self.linear_block = nn.Sequential(
nn.Linear(in_channels, out_channels),
nn.BatchNorm1d(out_channels),
nn.LeakyReLU()
)
def forward(self, x):
x = self.linear_block(x)
return x
class FracClassifier(nn.Module):
def __init__(
self,
encoder_channels,
classifier_channels=(256, 128, 64),
final_channels=2,
linear_kernel=4096,
p = 0.
):
super(FracClassifier, self).__init__()
# resnet gives
self.linear_kernel = linear_kernel
self.initial_conv = nn.Conv3d(encoder_channels, 256, kernel_size=3,stride=1, padding=1)
self.bn_init = nn.InstanceNorm3d(256, affine=True)
self.drop_1 = nn.Dropout(p=p)
self.initial_conv1 = nn.Conv3d(256, 128, kernel_size=3,stride=1, padding=1)
self.bn_init1 = nn.InstanceNorm3d(128, affine=True)
self.drop_2 = nn.Dropout(p=p)
self.initial_conv2 = nn.Conv3d(128, 64, kernel_size=3,stride=1, padding=1)
self.bn_init2 = nn.InstanceNorm3d(64, affine=True)
self.drop_3 = nn.Dropout(p=p)
self.initial_conv3 = nn.Conv3d(64, 8, kernel_size=3,stride=1, padding=1)
self.bn_init3 = nn.InstanceNorm3d(8, affine=True)
self.drop_4 = nn.Dropout(p=p)
self.vector_shape = encoder_channels
self.layer1 = LinearBlock(self.linear_kernel, classifier_channels[0])
self.drop_5 = nn.Dropout(p=p)
self.layer2 = LinearBlock(classifier_channels[0], classifier_channels[1])
self.drop_6 = nn.Dropout(p=p)
self.layer3 = LinearBlock(classifier_channels[1], classifier_channels[2])
self.drop_7 = nn.Dropout(p=p)
self.final_dense = nn.Linear(classifier_channels[2], final_channels)
# initialization of weights dropped as pytorch already uses them (now)
def forward(self, x):
x = self.initial_conv(x)
x = self.drop_1(self.bn_init(x))
x = self.initial_conv1(x)
x = self.drop_2(self.bn_init1(x))
x = self.initial_conv2(x)
x = self.drop_3(self.bn_init2(x))
x = self.initial_conv3(x)
x = self.drop_4(self.bn_init3(x))
x = x.view(x.shape[0], -1)
x = self.drop_5(self.layer1(x))
x = self.drop_6(self.layer2(x))
x = self.drop_7(self.layer3(x))
x = self.final_dense(x)
return x |