Created
May 13, 2020 12:12
-
-
Save gauravbansal98/5c299e44c46aade5e8f822aac597843c to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
class Encoder(nn.Module): | |
""" | |
Encodes the input image to a vector. | |
# """ | |
def __init__(self): | |
super(Encoder, self).__init__() | |
vgg = models.vgg16(pretrained=True) | |
model = torch.nn.Sequential() | |
for name, child in vgg.named_children(): | |
if isinstance(child, torch.nn.Sequential): | |
for cnt, layer in child.named_children(): | |
layer_name = name + str(cnt) | |
model.add_module(layer_name, layer) | |
else: | |
model.add_module(name, child) | |
model.add_module('flatten', nn.Flatten()) | |
# remove last two layers | |
modules = list(model.children())[:-2] | |
self.enc_model = nn.Sequential(*modules) | |
for p in self.enc_model.parameters(): | |
p.requires_grad = False | |
def forward(self, images): | |
encoded_out = self.enc_model(images) | |
return encoded_out | |
# extract features from all images | |
directory = 'Flicker8k_Dataset' | |
features = extract_features(directory) | |
print('Extracted Features: %d' % len(features)) | |
# save to file | |
dump(features, open('features.pkl', 'wb')) | |
print("Features are extracted and saved into the file") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment