I want to extract features from some pictures, and use res152 with the following code.
When it runs to 17/1000 pics, my ubuntu computer with 8G memory will crash.
import torchvision, torch, os
from torch import nn
from torch.autograd import Variable
from dataset import DATASET
mydata = DATASET(os.getcwd())
model = torchvision.models.resnet152(pretrained=True)
new_classifier = nn.Sequential(*list(model.children())[:-1])
model.classifier = new_classifier
pos1 = model(Variable(mydata[0][0].view(1, 3, 224, 224)))
for index in range(1, 1000):
f = model(Variable(mydata[index][0].view(1, 3, 224, 224)))
pos1 = torch.cat((pos1, f))
with open('pos1_fea.pt', 'wb') as f:
torch.save(pos1, f)