class DenseNet(nn.Module):
def __init__(self):
super(DenseNet, self).__init__()
input_channels = 1
conv_channels = 32
down_structure = [2, 2, 2]
output_channels = 4 #4#2
act_fn = config["act_fn"]
norm_fn = config["norm_fn"]
self.features = nn.Sequential()
self.features.add_module(
"init_conv",
nn.Conv3d(input_channels,
conv_channels,
kernel_size=3,
stride=1,
padding=1,
bias=True))
self.features.add_module("init_norm", norm_fn(conv_channels))
self.features.add_module("init_act", act_fn())
self.dropblock = LinearScheduler(DropBlock3D(drop_prob=0.,
block_size=5),
start_value=0.,
stop_value=0.5,
nr_steps=5e3)
channels = conv_channels
self.features.add_module('drop_block',
DropBlock3D(drop_prob=0.1, block_size=5))
for i, num_layers in enumerate(down_structure):
for j in range(num_layers):
conv_layer = ConvBlock(channels)
self.features.add_module(
"block{}_layer{}".format(i + 1, j + 1), conv_layer)
channels = conv_layer.out_channels
# dowmsample
trans_layer = TransmitBlock(
channels, is_last_layer=(i == len(down_structure) - 1))
self.features.add_module("transition{}".format(i + 1), trans_layer)
channels = trans_layer.out_channels
self.classifier = nn.Linear(channels, output_channels)
def forward(self, x, **return_opts):
self.dropblock.step()
batch_size, _, z, h, w = x.size()
features = self.dropblock(self.features(x))
# print("features", features.size())
pooled = F.adaptive_avg_pool3d(features, 1).view(batch_size, -1)
# print("pooled", pooled.size())
scores = self.classifier(pooled)
# print("scored", scores.size())
if len(return_opts) == 0:
return scores
class ResNetCustom(ResNet):
def __init__(self,
block,
layers,
num_classes=1000,
drop_prob=0.,
block_size=5):
super(ResNet, self).__init__()
self.inplanes = 64
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.dropblock = LinearScheduler(DropBlock2D(drop_prob=drop_prob,
block_size=block_size),
start_value=0.,
stop_value=drop_prob,
nr_steps=5e3)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
self.dropblock.step() # increment number of iterations
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.dropblock(self.layer1(x))
x = self.dropblock(self.layer2(x))
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.shape[0], -1)
x = self.fc(x)
return x
class ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes=10, drop_prob=0., block_size=5):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = conv3x3(3, 64)
self.bn1 = nn.BatchNorm2d(64)
self.dropblock = LinearScheduler(
DropBlock2D(drop_prob=drop_prob, block_size=block_size, att=True),
start_value=0.,
stop_value=drop_prob,
nr_steps=5e4
)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512 * block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
self.dropblock.step()
out = F.relu(self.bn1(self.conv1(x)))
out = self.dropblock(self.layer1(out))
out1 = out
out = self.dropblock(self.layer2(out))
out2 = out
out = self.layer3(out)
out3 = out
out = self.layer4(out)
out4 = out
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return [out, out1.cpu().detach().numpy(),
out2.cpu().detach().numpy(),
out3.cpu().detach().numpy(),
out4.cpu().detach().numpy()]
class ResNet(nn.Module):
def __init__(self,
block,
num_blocks,
num_classes=10,
drop_prob=0.,
block_size=5,
nr_steps=5000):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = conv3x3(3, 64)
self.bn1 = nn.BatchNorm2d(64)
self.dropblock = LinearScheduler(DropBlock2D(drop_prob=drop_prob,
block_size=block_size,
att=True),
start_value=0.,
stop_value=drop_prob,
nr_steps=nr_steps)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512 * block.expansion, num_classes)
# for m in self.modules():
# if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
self.dropblock.step()
out = F.relu(self.bn1(self.conv1(x)))
out = self.dropblock(self.layer1(out))
out = self.dropblock(self.layer2(out))
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
def get_mask(self, x, prob):
# self.dropblock.step()
self.dropblock.dropblock.drop_prob = prob
out = F.relu(self.bn1(self.conv1(x)))
out, blk1 = self.dropblock.get_mask(self.layer1(out))
out, blk2 = self.dropblock.get_mask(self.layer2(out))
# out = self.layer3(out)
# out = self.layer4(out)
# out = F.avg_pool2d(out, 4)
# out = out.view(out.size(0), -1)
# out = self.linear(out)
return out, blk1, blk2
class FluxNet(nn.Module):
default_filters = (8, 16, 24, 32, 64)
def __init__(self,
in_channel=1,
filters=default_filters,
aspp_dilation_ratio=1,
symmetric=True,
use_flux_head=True,
use_skeleton_head=False):
super().__init__()
# encoding path
self.symmetric = symmetric
if self.symmetric:
self.layer1_E = nn.Sequential(
conv3d_bn_relu(in_planes=in_channel,
out_planes=filters[0],
kernel_size=(5, 5, 5),
stride=1,
padding=(2, 2, 2)),
conv3d_bn_relu(in_planes=filters[0],
out_planes=filters[0],
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)),
residual_block_3d(filters[0], filters[0], projection=False))
else:
self.layer1_E = nn.Sequential(
conv3d_bn_relu(in_planes=in_channel,
out_planes=filters[0],
kernel_size=(1, 5, 5),
stride=1,
padding=(0, 2, 2)),
conv3d_bn_relu(in_planes=filters[0],
out_planes=filters[0],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_2d(filters[0], filters[0], projection=False))
self.layer2_E = nn.Sequential(
conv3d_bn_relu(in_planes=filters[0],
out_planes=filters[1],
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)),
residual_block_3d(filters[1], filters[1], projection=False))
self.layer3_E = nn.Sequential(
conv3d_bn_relu(in_planes=filters[1],
out_planes=filters[2],
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)),
residual_block_3d(filters[2], filters[2], projection=False))
self.layer4_E = nn.Sequential(
conv3d_bn_relu(in_planes=filters[2],
out_planes=filters[3],
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)),
residual_block_3d(filters[3], filters[3], projection=False))
# center ASPP block
self.center = ASPP(
filters[3], filters[4],
[[2, int(2 * aspp_dilation_ratio),
int(2 * aspp_dilation_ratio)],
[3, int(3 * aspp_dilation_ratio),
int(3 * aspp_dilation_ratio)],
[5, int(5 * aspp_dilation_ratio),
int(5 * aspp_dilation_ratio)]])
# decoding path
self.layer1_D_flux = nn.Sequential(
conv3d_bn_relu(in_planes=filters[1],
out_planes=filters[0],
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0)),
residual_block_3d(filters[0], filters[0], projection=False),
conv3d_bn_non(in_planes=filters[0],
out_planes=3,
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)))
self.layer1_D_skeleton = nn.Sequential(
conv3d_bn_relu(in_planes=filters[1],
out_planes=filters[0],
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0)),
residual_block_3d(filters[0], filters[0], projection=False),
conv3d_bn_non(in_planes=filters[0],
out_planes=1,
kernel_size=(3, 3, 3),
stride=1,
padding=(1, 1, 1)))
self.layer2_D = nn.Sequential(
conv3d_bn_relu(in_planes=filters[2],
out_planes=filters[1],
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0)),
residual_block_3d(filters[1], filters[1], projection=False))
self.layer3_D = nn.Sequential(
conv3d_bn_relu(in_planes=filters[3],
out_planes=filters[2],
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0)),
residual_block_3d(filters[2], filters[2], projection=False))
self.layer4_D = nn.Sequential(
conv3d_bn_relu(in_planes=filters[4],
out_planes=filters[3],
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0)),
residual_block_3d(filters[3], filters[3], projection=False))
# downsample pooling
self.down = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
self.down_aniso = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
# upsampling
self.up = nn.Upsample(scale_factor=(2, 2, 2), mode='nearest')
self.up_aniso = nn.Upsample(scale_factor=(1, 2, 2), mode='nearest')
self.dropblock: LinearScheduler = None
self.use_flux_head = use_flux_head
self.use_skeleton_head = use_skeleton_head
#initialization
ortho_init(self)
def init_dropblock(self, start_value, stop_value, nr_steps, block_size):
self.dropblock = LinearScheduler(DropBlock3D(drop_prob=stop_value,
block_size=block_size),
start_value=start_value,
stop_value=stop_value,
nr_steps=nr_steps)
def forward(self,
x,
call_dropblock_step=False,
get_penultimate_layer=False):
if self.dropblock and call_dropblock_step:
self.dropblock.step()
# encoding path
z1 = self.layer1_E(x)
z1 = self.dropblock(z1) if self.dropblock else z1
x = self.down(z1) if self.symmetric else self.down_aniso(z1)
z2 = self.layer2_E(x)
z2 = self.dropblock(z2) if self.dropblock else z2
x = self.down(z2)
z3 = self.layer3_E(x)
x = self.down(z3)
z4 = self.layer4_E(x)
#center ASPP block
x = self.center(z4)
# decoding path
x = self.layer4_D(x)
x = self.up(x)
x = self.layer3_D(x)
x = x + z3
x = self.up(x)
x = self.layer2_D(x)
x = x + z2
x = self.up(x) if self.symmetric else self.up_aniso(x)
output = dict()
if self.use_flux_head:
flux = self.layer1_D_flux(x)
# TODO share the penultimate layer code for skeleton head and flux head
# for i, layer in enumerate(self.layer1_D_flux):
# x = layer(x)
# if get_penultimate_layer and i is 0:
# output['penultimate_layer'] = x
flux = torch.tanh(flux)
output['flux'] = flux
if self.use_skeleton_head:
skeleton = self.layer1_D_skeleton(x)
skeleton = torch.sigmoid(skeleton)
output['skeleton'] = skeleton
if not output:
raise ValueError("Neither flux or skeleton head was specified.")
return output
class ResNetDropblock(nn.Module):
# TNet with Non-Local-Block on the 4th stage (before the last conv of the 4th stage)
def __init__(self, output_class=7, model_path=None, resnet_type=34, drop_block=False, drop_prob=0.5, drop_pos=None, layer_depth=4, drop_block_size=7):
super(ResNetDropblock, self).__init__()
assert resnet_type in [18, 34, 50]
assert layer_depth in [1, 2, 3, 4]
if resnet_type == 18:
self.base = resnet18(pretrained=False, num_classes=1000)
# state_dict = torch.load('resnet18.pth')
if layer_depth == 4:
last_fc_in_channel = 512 * 1
elif layer_depth == 3:
last_fc_in_channel = 256 * 1
elif layer_depth == 2:
last_fc_in_channel = 128 * 1
else: # elif layer_depth == 1:
last_fc_in_channel = 64 * 1
elif resnet_type == 34:
self.base = resnet34(pretrained=False, num_classes=1000)
# state_dict = torch.load('resnet34.pth')
if layer_depth == 4:
last_fc_in_channel = 512 * 1
elif layer_depth == 3:
last_fc_in_channel = 256 * 1
elif layer_depth == 2:
last_fc_in_channel = 128 * 1
else: # elif layer_depth == 1:
last_fc_in_channel = 64 * 1
else: # elif resnet_type == 50:
self.base = resnet50(pretrained=False, num_classes=1000)
# state_dict = torch.load('resnet50.pth')
if layer_depth == 4:
last_fc_in_channel = 512 * 4
elif layer_depth == 3:
last_fc_in_channel = 256 * 4
elif layer_depth == 2:
last_fc_in_channel = 128 * 4
else: # elif layer_depth == 1:
last_fc_in_channel = 64 * 4
# self.base.load_state_dict(state_dict)
# def weight_init(m):
# if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
# nn.init.xavier_uniform_(m.weight, gain=nn.init.calculate_gain('relu'))
# if m.bias is not None:
# nn.init.zeros_(m.bias)
#
# self.base.layer3.apply(weight_init)
# self.base.layer4.apply(weight_init)
self.base.fc = nn.Linear(in_features=last_fc_in_channel, out_features=output_class, bias=True)
if drop_block:
self.dropblock = LinearScheduler(
DropBlock2D(drop_prob=drop_prob, block_size=drop_block_size),
start_value=0.,
stop_value=drop_prob,
nr_steps=300
)
else:
self.dropblock = nn.Sequential()
self.drop_pos = drop_pos
self.layer_depth = layer_depth
if model_path is not None:
self.model_path = model_path
assert '.pth' in self.model_path or '.pkl' in self.model_path
self.init_weight()
def init_weight(self):
print('Loading weights into state dict...')
self.load_state_dict(torch.load(self.model_path, map_location=lambda storage, loc: storage))
print('Finished!')
def forward(self, x):
if type(self.dropblock) != nn.Sequential:
self.dropblock.step()
x = self.base.conv1(x)
x = self.base.bn1(x)
x = self.base.relu(x)
x = self.base.maxpool(x)
if self.drop_pos == 1 or self.drop_pos is None:
x = self.dropblock(self.base.layer1(x))
else:
x = self.base.layer1(x)
if self.layer_depth > 1:
if self.drop_pos == 2 or self.drop_pos is None:
x = self.dropblock(self.base.layer2(x))
else:
x = self.base.layer2(x)
if self.layer_depth > 2:
if self.drop_pos == 3 or self.drop_pos is None:
x = self.dropblock(self.base.layer3(x))
else:
x = self.base.layer3(x)
if self.layer_depth > 3:
if self.drop_pos == 4 or self.drop_pos is None:
x = self.dropblock(self.base.layer4(x))
else:
x = self.base.layer4(x)
x = self.base.avgpool(x)
x = torch.flatten(x, 1)
out = self.base.fc(x)
return out
class ResNet(nn.Module):
def __init__(self, cfg, block, layers):
self.inplanes = 64
super(ResNet, self).__init__()
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False) #150*150
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.cfg = cfg
#----------new structure called DropBlock 11sr April-------------------------
if self.cfg.MODEL.BACKBONE.DROP_BLOCK:
drop_prob = 0.5
block_size = 3
self.dropblock = LinearScheduler(DropBlock2D(
drop_prob=drop_prob, block_size=block_size),
start_value=0.,
stop_value=drop_prob,
nr_steps=5)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) #75*75
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2) #38*38
self.layer3 = self._make_layer(block, 256, layers[2], stride=2) #19*19
#self.layer4 = self._make_layer(block, 512,layers[3] , stride=1) #10*10
self.ex_layer0 = self._make_layer(block, 512, 2, stride=2) #10*10
#self.maxpoo2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
#2. extra_layers (ReLU will be used in the foward function) 10thApril,Xiaoyu Zhu
# if cfg.MODEL.BACKBONE.DEPTH>34:
# self.ex_layer1 = nn.Sequential(BasicBlock_modified(2048,512))
# #self.ex_layer1 = self._make_extra_layers(2048,512,3,1) #5*5
# else:
# self.ex_layer1 = nn.Sequential(BasicBlock_modified(512,512))
self.ex_layer1 = self._make_layer(block, 512, 1, stride=2) #5*5
#self.maxpoo3 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.ex_layer2 = self._make_layer(block, 256, 1, stride=2) #3*3
#self.maxpoo4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.ex_layer3 = self._make_layer(block, 128, 1, stride=2)
# if cfg.MODEL.BACKBONE.DEPTH>34:
# self.ex_layer3 = self._make_extra_layers(256*4,128,[2,3],0)
# else:
# #self.ex_layer3 = self._make_extra_layers(256,128,[2,3],0)
# self.ex_layer3 = self._make_layer(block, 128, 1, stride=2)
#BasicBlock_modified(inplanes=256, planes=128, kernel=[2,3],stride=2, padding=0)
# kaiming weight normal after default initialization
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
# construct layer/stage conv2_x,conv3_x,conv4_x,conv5_x
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
# when to need downsample
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes,
planes * block.expansion,
kernel_size=1,
stride=stride,
bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
# inplanes expand for next block
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def _make_extra_layers(self, input_channels, output_channels, k,
p): #10thApril,Xiaoyu Zhu
layers = []
layers.append(
torch.nn.Conv2d(in_channels=input_channels,
out_channels=output_channels,
kernel_size=k,
stride=1,
padding=1))
layers.append(nn.ReLU(inplace=True))
layers.append(nn.BatchNorm2d(output_channels))
#layers.append(nn.Dropout(0.5))
layers.append(
torch.nn.Conv2d(in_channels=output_channels,
out_channels=output_channels,
kernel_size=k,
stride=2,
padding=p))
layers.append(nn.ReLU(inplace=True))
layers.append(nn.BatchNorm2d(output_channels))
return nn.Sequential(*layers)
def forward(self, x):
if self.cfg.MODEL.BACKBONE.DROP_BLOCK:
self.dropblock.step() # increment number of iterations
out_features = []
#print('Input',x.shape) #Original 300*300; For rectange input size :320*240
x = self.conv1(x)
#print('self.conv1',x.shape) #150*150; 160*120
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
#print('self.maxpool',x.shape) #75*75; 80*60
if self.cfg.MODEL.BACKBONE.DROP_BLOCK:
x = self.dropblock(self.layer1(x)) #added 11st April
else:
x = self.layer1(x)
#print('layer1',x.shape) #80*60
#out_features.append(x) #Add new feature map 60*80
if self.cfg.MODEL.BACKBONE.DROP_BLOCK:
x = self.dropblock(self.layer2(x))
else:
x = self.layer2(x)
#print('layer2',x.shape) #38*38 output[0]; 30*40
out_features.append(x)
x = self.layer3(x)
#print('layer3',x.shape) #19*19 output[1]; 15*20
out_features.append(x)
#x = self.layer4(x)
x = self.ex_layer0(x)
#x = self.maxpoo2(x)
#print('layer4',x.shape) #10*10 output[2]; 8*10
out_features.append(x)
#For other output: 10thApril,Xiaoyu Zhu
x = self.ex_layer1(x)
#x = self.maxpoo3(x)
#print('ex_layer1',x.shape) #5*5 output[3] ;4*5
out_features.append(x)
x = self.ex_layer2(x)
#x = self.maxpoo4(x)
#print('ex_layer2',x.shape) #3*2 output[4] ;2*3
out_features.append(x)
x = self.ex_layer3(x)
#print('ex_layer3',x.shape) #1*1 output[5]
out_features.append(x)
#-----------------------------------Old Version-------------------
# #For other outputs:
# for i, f in enumerate(self.extra_layer): #i means the index and f means the function of the layer
# if i== len(self.extra_layer):
# x= f(x)
# else:
# x= torch.nn.functional.relu(f(x),inplace=True)
# if i % 2 == 1:
# out_features.append(x)
#-----------------------------------------------------------------
return tuple(out_features)
class CNN(nn.Module):
def __init__(self, grayscale=False):
super(CNN, self).__init__()
# hyperparameters
# Activation = Mish()
Activation = nn.ReLU()
drop_prob = 0.2
block_size = 5
kernel_sizes = np.array([7, 3, 3, 3, 3])
conv_strides = np.array([3, 1, 1, 1, 1])
pad_sizes = (kernel_sizes - 1) / 2
out_channels = np.array([64, 128, 256, 256, 512])
pool_sizes = np.array([2, 2, 2, 2, 2])
pool_strides = np.array([2, 2, 2, 2, 2])
fcs = [512, 100] #512, 100
if do_dropblock:
self.Dropblock = LinearScheduler(DropBlock2D(
drop_prob=drop_prob, block_size=block_size),
start_value=0.,
stop_value=drop_prob,
nr_steps=5000)
conv_layers = []
assert len(kernel_sizes) == len(out_channels) == len(
pool_sizes), "inconsistent layer length"
layers = range(len(kernel_sizes))
input_dim = 1 if grayscale else 3
out_channels = np.insert(out_channels, 0, input_dim)
for l in layers:
conv_layers.append(
nn.Conv2d(out_channels[l],
out_channels[l + 1],
kernel_sizes[l],
stride=conv_strides[l],
padding=int(pad_sizes[l])))
conv_layers.append(nn.BatchNorm2d(out_channels[l + 1]))
# conv_layers.append(nn.Dropout(0.1))
if do_dropblock:
conv_layers.append(self.Dropblock)
conv_layers.append(Activation)
if pool_sizes[l]:
conv_layers.append(
nn.MaxPool2d(int(pool_sizes[l]),
stride=int(pool_strides[l])))
self.conv = nn.Sequential(*conv_layers)
# compute input shape of FCs
x = torch.zeros([1, input_dim, 224, 224])
x = self.conv(x)
x = x.view(1, -1)
input_FC = x.size(1)
FCs = []
layers = range(len(fcs))
# input length: input_FC
fcs = np.insert(fcs, 0, input_FC)
for l in layers:
FCs.append(nn.Linear(fcs[l], fcs[l + 1]))
FCs.append(Activation)
FCs.append(nn.BatchNorm1d(fcs[l + 1]))
FCs.append(nn.Dropout(0.5))
self.fc = nn.Sequential(*FCs)
self.out = nn.Linear(fcs[-1], 3) # three categories: class A, B, C
def forward(self, x):
if do_dropblock:
self.Dropblock.step() # increment number of iterations
x = self.conv(x)
x = x.view(x.size(0), -1) # flattened
if not do_dropblock:
x = F.dropout(x, p=0.5)
x = self.fc(x)
x = self.out(x)
return x
def get_conv_output(self, x):
if do_dropblock:
self.Dropblock.step() # increment number of iterations
x = self.conv(x)
x = x.view(x.size(0), -1) # flattened
return x.cpu().numpy()
