File size: 9,602 Bytes
08ec965 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | import argparse
import os
import sys
import torch
import shutil
import utils
import numpy as np
import metrics
import torch.nn as nn
import torch.backends.cudnn as cudnn
import vision_transformer as vits
from tqdm import tqdm
from datasets import build_dataset
from meters import AverageEpochMeter
from object_discovery import ncut, detect_box, get_feats
from visualization import visualize_fms, visualize_predictions_gt, visualize_img, visualize_eigvec
from collections import OrderedDict
#from visualizations import visualize_img, visualize_eigvec, visualize_gt, visualize_fms, visualize_predictions, visualize_predictions_gt
def evaluate():
if args.device != 'cuda':
args.distributed = False
else:
utils.init_distributed_mode(args)
print(args)
device = torch.device(args.device)
cudnn.benchmark = True
# =============build network ==============
if args.arch in vits.__dict__.keys():
model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=args.num_labels)
embed_dim = model.embed_dim * (args.n_last_blocks + int(args.avgpool_patchtokens))
print(f'embed_dim: {embed_dim}')
else:
print(f"Unknow architecture: {args.arch}")
sys.exit(1)
model.to(device)
model.eval()
linear_classifier = LinearClassifier(embed_dim, num_labels=args.num_labels)
linear_classifier = linear_classifier.to(device)
# ============ Load checkpoint =============
checkpoint_model = torch.load(args.pretrained_weights, map_location='cpu')
if args.classifier_weights is None:
linear_classifier.load_state_dict(checkpoint_model['classifier'])
model.load_state_dict(checkpoint_model['model'], strict=False)
else:
model.load_state_dict(checkpoint_model)
checkpoint_classifier = torch.load(args.classifier_weights, map_location='cpu')
new_state_dict = OrderedDict()
new_state_dict = {k[7:]: v for k, v in checkpoint_classifier['state_dict'].items()}
# load params
linear_classifier.load_state_dict(new_state_dict)
print('Load from checkpoint done.')
dataset_val, _ = build_dataset(is_train=False, args=args)
if args.distributed:
sampler_val = torch.utils.data.distributed.DistributedSampler(dataset_val, shuffle=False)
else:
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
val_loader = torch.utils.data.DataLoader(
dataset_val,
sampler=sampler_val,
batch_size=args.batch_size_per_gpu,
num_workers=args.num_workers,
pin_memory=True,
)
num_batches = len(val_loader)
scales = [args.patch_size, args.patch_size]
top1_cls_meter = AverageEpochMeter('Top-1 Cls')
top5_cls_meter = AverageEpochMeter('Top-5 Cls')
gt_loc_meter = AverageEpochMeter('GT-Known Loc')
top1_loc_meter = AverageEpochMeter('Top-1 Loc')
feat_out = {}
bbox_error = 0
cls_error = 0
skip = 0
def hook_fn_forward_qkv(module, input, output):
feat_out["qkv"] = output
model._modules["blocks"][-1]._modules["attn"]._modules["qkv"].register_forward_hook(hook_fn_forward_qkv)
val_loader = tqdm(val_loader)
for i, (images, labels, gt_boxes) in enumerate(val_loader):
init_image_size = images[0].shape
images = utils.padding_img(images[0], args)
images = images.unsqueeze(0)
if images.shape[-1] > 1000 and images.shape[-2] > 1000:
skip = skip + 1
continue
with torch.no_grad():
images = images.to(device)
labels = labels.to(device)
intermediate_output, shape= model.get_intermediate_layers(images, args.n_last_blocks)
output = torch.cat([x[:, 0] for x in intermediate_output], dim=-1)
if args.avgpool_patchtokens:
output = torch.cat((output.unsqueeze(-1), torch.mean(intermediate_output[-1][:, 1:], dim=1).unsqueeze(-1)), dim=-1)
output = output.reshape(output.shape[0], -1)
output = linear_classifier(output)
batch_size = images.size(0)
top1_cls, top5_cls = utils.accuracy(output, labels, topk=(1, 5))
top1_cls_meter.update(top1_cls, batch_size)
top5_cls_meter.update(top5_cls, batch_size)
# ===== Generate Bbox=========
dims = (images[0].shape[-2] // args.patch_size, images[0].shape[-1] // args.patch_size)
k = get_feats(feat_out, shape)
bboxes, mask, seed, eigvec= ncut(k, dims, scales, init_image_size=init_image_size[1:], eps=args.eps, tau=args.tau)
bboxes_copy = np.copy(bboxes)
bboxes = torch.FloatTensor(bboxes).unsqueeze(0)
gt_loc, top1_loc = metrics.loc_accuracy(output, labels, gt_boxes, bboxes)
if args.visual:
visualize_predictions_gt(images[0], bboxes_copy , gt_boxes, i, seed, dims, scales, './output/all')
visualize_eigvec(eigvec, './output/all', i, dims, scales)
#visualize_fms(images[0], mask, seed, i, dims, scales)
if top1_loc == 0 and gt_loc == 0:
bbox_error = bbox_error + 1
if args.visual:
# print(f'i: {i}, gt_boxes {gt_boxes}, prediction: {bboxes_copy}, image size: {images.shape}')
visualize_predictions_gt(images[0], bboxes_copy , gt_boxes, i, seed, dims, scales, './output/bbox_error')
elif top1_loc == 0 and gt_loc == 1:
cls_error = cls_error +1
if args.visual:
visualize_predictions_gt(images[0], bboxes_copy , gt_boxes, i, seed, dims, scales, './output/classification_error')
gt_loc_meter.update(gt_loc, batch_size)
top1_loc_meter.update(top1_loc, batch_size)
top1_cls = top1_cls_meter.compute()
top5_cls = top5_cls_meter.compute()
gt_loc = gt_loc_meter.compute()
top1_loc = top1_loc_meter.compute()
val_loader.set_description(f'Top1_cls: {top1_cls:.4f}, top5_cls{top5_cls:.4f}, gt_loc: {gt_loc:.4f}, top1_loc:{top1_loc:.4f}')
print(f'Top1_cls: {top1_cls}, top5_cls{top5_cls}, gt_loc: {gt_loc}, top1_loc:{top1_loc}')
print(f'Bbox error: {bbox_error}, cls error: {cls_error}')
print(f'Skip large image: {skip}')
class LinearClassifier(nn.Module):
"""Linear layer to train on top of frozen features"""
def __init__(self, dim, num_labels=1000):
super(LinearClassifier, self).__init__()
self.num_labels = num_labels
self.linear = nn.Linear(dim, num_labels)
self.linear.weight.data.normal_(mean=0.0, std=0.01)
self.linear.bias.data.zero_()
def forward(self, x):
# flatten
x = x.view(x.size(0), -1)
# linear layer
return self.linear(x)
if __name__=='__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--n_last_blocks', default=4, type=int, help="""Concatenate [CLS] tokens
for the `n` last blocks. We use `n=4` when evaluating ViT-Small and `n=1` with ViT-Base.""")
parser.add_argument('--avgpool_patchtokens', default=False, type=utils.bool_flag,
help="""Whether ot not to concatenate the global average pooled features to the [CLS] token.
We typically set this to False for ViT-Small and to True with ViT-Base.""")
parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
parser.add_argument('--dataset', default='cub', type=str, choices=['cub', 'imagenet'], help='Architecture')
parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
parser.add_argument('--input_size', default=224, type=int, help='Input image size, default(224).')
parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
parser.add_argument('--classifier_weights', default=None, type=str, help="Path to linear classifier pretrained weights to evaluate.")
parser.add_argument('--batch_size_per_gpu', default=256, type=int, help='Per-GPU batch-size')
parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
distributed training; see https://pytorch.org/docs/stable/distributed.html""")
parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
parser.add_argument('--data_path', default='/path/to/imagenet/', type=str)
parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
parser.add_argument('--val_freq', default=1, type=int, help="Epoch frequency for validation.")
parser.add_argument('--num_labels', default=200, type=int, help='Number of labels for linear classifier')
parser.add_argument('--device', default='cuda', help='device to use for training / testing')
parser.add_argument('--distributed', default=False, action='store_true', help='device to use for training / testing')
parser.add_argument('--ori_size', default=False, action='store_true', help='Evaluate on image raw size')
parser.add_argument('--visual', default=False, action='store_true', help='Visualize error examples on ./test')
parser.add_argument('--eps', default=1e-5, type=float, help='hyperparameter for tokencut')
parser.add_argument('--tau', default=0.05, type=float, help='hyperparamter for tokencut')
parser.add_argument('--no_center_crop', default=False, action='store_true', help='Center crop input image')
args = parser.parse_args()
evaluate()
|