I am trying to convert torchvision vgg layers into a block of layers which can be seen in the following code. I am trying to put them in a defaultdict and use it as a block. When I try to print the model the defaultdict doesnot appear. How do I make defaultdict appear in my model ?
# here feature head is vgg model taken from torchvision.models
class FCN8(_FCN_BASE):
def __init__(self, cfgs, feature_head, num_classes=2):
super(FCN8, self).__init__()
self.cfgs = cfgs
self.feature_head = self.convert_features_to_blocks(feature_head)
self.fc6 = nn.Conv2d(512, 4096, kernel_size=7)
self.relu6 = nn.ReLU(inplace=True)
self.drop6 = nn.Dropout2d()
# fc7
self.fc7 = nn.Conv2d(4096, 4096, 1)
self.relu7 = nn.ReLU(inplace=True)
self.drop7 = nn.Dropout2d()
self.score_fr = nn.Conv2d(4096, num_classes, 1)
self.score_pool3 = nn.Conv2d(256, num_classes, 1)
self.score_pool4 = nn.Conv2d(512, num_classes, 1)
self.upscore2 = nn.ConvTranspose2d(
num_classes, num_classes, 4, stride=2, bias=False)
self.upscore8 = nn.ConvTranspose2d(
num_classes, num_classes, 16, stride=8, bias=False)
self.upscore_pool4 = nn.ConvTranspose2d(
num_classes, num_classes, 4, stride=2, bias=False)
self._initialize_weights()
def convert_features_to_blocks(self, feature_head):
blocks = defaultdict()
if feature_head.__class__.__name__ in ['VGG']:
features = feature_head.features
start = 0
layer = 1
for i, feature in enumerate(features):
if feature.__class__.__name__ == 'MaxPool2d':
block = nn.Sequential(features[start: i + 1])
blocks['l_{}'.format(layer)] = block
start = i + 1
layer += 1
return blocks
def forward(self, x):
x = self.feature_head['l_1'](x)
x = self.feature_head['l_2'](x)
When I print the model. It does not show the defaultdict
FCN8(
(fc6): Conv2d(512, 4096, kernel_size=(7, 7), stride=(1, 1))
(relu6): ReLU(inplace)
(drop6): Dropout2d(p=0.5)
(fc7): Conv2d(4096, 4096, kernel_size=(1, 1), stride=(1, 1))
(relu7): ReLU(inplace)
(drop7): Dropout2d(p=0.5)
(score_fr): Conv2d(4096, 1, kernel_size=(1, 1), stride=(1, 1))
(score_pool3): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))
(score_pool4): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1))
(upscore2): ConvTranspose2d(1, 1, kernel_size=(4, 4), stride=(2, 2), bias=False)
(upscore8): ConvTranspose2d(1, 1, kernel_size=(16, 16), stride=(8, 8), bias=False)
(upscore_pool4): ConvTranspose2d(1, 1, kernel_size=(4, 4), stride=(2, 2), bias=False)
)