I’m trying to use the global-pooled bottleneck outputs of ResNet as inputs to my model. Currently this is what I have:
class MyNet(nn.Module):
def __init__(self, resnet, n_features=12):
super().__init__()
self.initial_layers = nn.Sequential(*list(resnet.children())[:4])
self.bottlenecks = []
self.attribute_weights = nn.Linear(15104, n_features)
# Extract the bottleneck layers
for i, mod in enumerate(list(resnet.children())):
if isinstance(mod, nn.Sequential):
for bn in mod:
self.bottlenecks.append(bn)
# Set the resnet weights to not update
# for param in resnet.parameters():
# param.requires_grad = False
def forward(self, inp):
all_feature_maps = []
output = self.initial_layers(inp)
# Loop to extract the outputs of the bottleneck layers from resnet
for bn in self.bottlenecks:
output = bn(output)
kernel_size = (output.size()[2], output.size()[3])
feature_maps = F.avg_pool2d(output, kernel_size)
all_feature_maps.append(feature_maps)
# Global pool
features = torch.cat(all_feature_maps, dim=1).squeeze()
if len(features.size()) == 1:
features = features.unsqueeze(0)
# Use features to predict scores
attribute_scores = self.attribute_weights(features)
return attribute_scores
This is working and the model can be created/trained via:
from torchvision import models
resnet = models.resnet50(pretrained=True)
net = MyNet(resnet, n_features=12)
optimizer = torch.optim.Adam(params=net.attribute_weights.parameters())
# train code...
Now I want to take this further and not only train the attribute_weights but also fine-tune the weights of the pretrained ResNet. To achieve this I’m rebuilding ResNet and creating “transparent layers” where I global pool the feature maps of the bottleneck outputs.
class MyNet3(nn.Module):
def __init__(self, resnet, n_features=12):
super().__init__()
self.model = nn.Sequential(*list(resnet.children())[:4])
self.all_pooled_features = []
self.attribute_weights = nn.Linear(15104, n_features)
count = 0
for i, mod in enumerate(list(resnet.children())):
# Extract the bottleneck layers
if isinstance(mod, nn.Sequential):
for bn in mod:
self.model.add_module(f"bn_{count}", bn)
# Use a "Transparent layer"
pooled_feature_map = PoolFeatures()
self.model.add_module(f"pooled_feature_{count}", pooled_feature_map)
self.all_pooled_features.append(pooled_feature_map)
count += 1
def forward(self, inp):
_ = self.model(inp)
features = torch.cat([pool_fp.feature_maps for pool_fp in self.all_pooled_features], dim=1).squeeze()
if len(features.size()) == 1:
features = features.unsqueeze(0)
# Use features to predict scores
attribute_scores = self.attribute_weights(features)
return attribute_scores
My Pooling Module is:
class PoolFeatures(nn.Module):
def __init__(self):
super().__init__()
def forward(self, inp):
kernel_size = (inp.size()[2], inp.size()[3])
self.feature_maps = F.avg_pool2d(inp, kernel_size)
return inp
Again this seems (?) to be working and can be created and trained via:
resnet = models.resnet50(pretrained=True)
net = MyNet(resnet, n_features=12)
optimizer = torch.optim.Adam(params=net.attribute_weights.parameters())
# train code...
The problem is that I get a memory error and it seems like after each training batch the memory increases. I can’t help but think that this is as a result of my transparent layer. How would be I solve this problem, or what is the correct way to implement what I want? I’ve looked at the Neural Transfer and it seems like .detach and .clone might be useful, though I’m not that sure. Thanks 