Compute feature vectors for an image dataset using a neural network. Parameters: dataset (torch.utils.data.Dataset): dataset model (str or torch.nn.Module, optional): pretrained model. If it is a str, use the last hidden model of that model.
(self, dataset, model="resnet50", batch_size=128)
| 361 | file.close() |
| 362 | |
| 363 | def image_feature_data(self, dataset, model="resnet50", batch_size=128): |
| 364 | """ |
| 365 | Compute feature vectors for an image dataset using a neural network. |
| 366 | |
| 367 | Parameters: |
| 368 | dataset (torch.utils.data.Dataset): dataset |
| 369 | model (str or torch.nn.Module, optional): pretrained model. |
| 370 | If it is a str, use the last hidden model of that model. |
| 371 | batch_size (int, optional): batch size |
| 372 | """ |
| 373 | import torch |
| 374 | import torchvision |
| 375 | from torch import nn |
| 376 | |
| 377 | logger.info("computing %s feature" % model) |
| 378 | if isinstance(model, str): |
| 379 | full_model = getattr(torchvision.models, model)(pretrained=True) |
| 380 | model = nn.Sequential(*list(full_model.children())[:-1]) |
| 381 | num_worker = multiprocessing.cpu_count() |
| 382 | data_loader = torch.utils.data.DataLoader(dataset, |
| 383 | batch_size=batch_size, num_workers=num_worker, shuffle=False) |
| 384 | model = model.cuda() |
| 385 | model.eval() |
| 386 | |
| 387 | features = [] |
| 388 | with torch.no_grad(): |
| 389 | for i, (batch_images, batch_labels) in enumerate(data_loader): |
| 390 | if i % 100 == 0: |
| 391 | logger.info("%g%%" % (100.0 * i * batch_size / len(dataset))) |
| 392 | batch_images = batch_images.cuda() |
| 393 | batch_features = model(batch_images).view(batch_images.size(0), -1).cpu().numpy() |
| 394 | features.append(batch_features) |
| 395 | features = np.concatenate(features) |
| 396 | |
| 397 | return features |
| 398 | |
| 399 | |
| 400 | class BlogCatalog(Dataset): |
no test coverage detected