def get_model(paths, feature_type):
if feature_type == 'vgg':
feature_network = feature_model.Vgg16(
num_classes=len(metadata.action_classes))
elif feature_type == 'resnet':
feature_network = feature_model.Resnet152(
num_classes=len(metadata.action_classes))
elif feature_type == 'densenet':
feature_network = feature_model.Densenet(
num_classes=len(metadata.action_classes))
else:
raise ValueError('feature type not recognized')
if feature_type.startswith('alexnet') or feature_type.startswith('vgg'):
feature_network.features = torch.nn.DataParallel(
feature_network.features)
feature_network.cuda()
else:
feature_network = torch.nn.DataParallel(feature_network).cuda()
checkpoint_dir = os.path.join(paths.tmp_root, 'checkpoints', 'vcoco',
'finetune_{}'.format(feature_type))
best_model_file = os.path.join(checkpoint_dir, 'model_best.pth')
checkpoint = torch.load(best_model_file)
feature_network.load_state_dict(checkpoint['state_dict'])
return feature_network
def get_model(paths):
# vgg16 = Vgg16(last_layer=1).cuda()
if feature_mode == 'None':
return
feature_network = feature_model.Resnet152(num_classes=len(metadata.action_classes))
feature_network.cuda()
checkpoint_dir = os.path.join(os.path.dirname(__file__), '../../../../data/model_resnet_noisy/finetune_{}_noisy'.format(feature_mode))
best_model_file = os.path.join(checkpoint_dir, 'model_best.pth')
checkpoint = torch.load(best_model_file)
for k in list(checkpoint['state_dict'].keys()):
if k[:7] == 'module.':
checkpoint['state_dict'][k[7:]] = checkpoint['state_dict'][k]
del checkpoint['state_dict'][k]
feature_network.load_state_dict(checkpoint['state_dict'])
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
return feature_network, transform
def main(args):
best_prec1 = 0.0
args.distributed = args.world_size > 1
if args.distributed:
torch.distributed.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.feature_type == 'vgg':
model = feature_model.Vgg16(num_classes=len(metadata.action_classes))
elif args.feature_type == 'resnet':
model = feature_model.Resnet152(
num_classes=len(metadata.action_classes))
elif args.feature_type == 'densenet':
model = feature_model.Densenet(
num_classes=len(metadata.action_classes))
input_imsize = (244, 244)
if not args.distributed:
if args.feature_type.startswith(
'alexnet') or args.feature_type.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
else:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss(
weight=torch.FloatTensor(metadata.action_class_weight)).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(os.path.join(args.resume, 'model_best.pth')):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(os.path.join(args.resume,
'model_best.pth'))
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(
os.path.join(args.resume, 'model_best.pth')))
torch.backends.cudnn.benchmark = True
# Data loading code
normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transform = torchvision.transforms.Compose([
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.RandomVerticalFlip(),
torchvision.transforms.ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.3,
hue=0.15),
torchvision.transforms.RandomAffine(degrees=15,
shear=15,
resample=PIL.Image.BILINEAR),
torchvision.transforms.ToTensor(), normalize
])
test_transform = torchvision.transforms.Compose(
[torchvision.transforms.ToTensor(), normalize])
train_dataset = feature_model.VCOCO(args.data, input_imsize,
train_transform, 'train')
val_dataset = feature_model.VCOCO(args.data, input_imsize, train_transform,
'val')
test_dataset = feature_model.VCOCO(args.data, input_imsize, test_transform,
'test')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(val_loader, model, criterion, optimizer, epoch)
train(train_loader, model, criterion, optimizer, epoch)
if epoch == 0 or epoch >= 3:
# evaluate on validation set
prec1 = validate(test_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
print('Best precision: {:.03f}'.format(best_prec1))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.feature_type,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
test_prec = validate(test_loader, model, criterion, test=True)
print('Testing precision: {:.04f}'.format(test_prec))
elif len(img.shape) == 5:
img = img.transpose([0, 1, 4, 2, 3])
img = torch.autograd.Variable(torch.Tensor(img)).cuda()
return img
meta_dir = os.path.join(os.path.dirname(__file__),
'../../../data/vcoco_features')
img_dir = '/home/tengyu/dataset/mscoco/images'
checkpoint_dir = '/home/tengyu/github/Part-GPNN/data/model_resnet_noisy/finetune_resnet'
vcoco_root = '/home/tengyu/dataset/v-coco/data'
save_data_path = '/home/tengyu/github/Part-GPNN/data/feature_resnet_tengyu2'
os.makedirs(save_data_path, exist_ok=True)
feature_network = feature_model.Resnet152(
num_classes=len(metadata.action_classes))
feature_network.cuda()
best_model_file = os.path.join(checkpoint_dir, 'model_best.pth')
checkpoint = torch.load(best_model_file)
for k in list(checkpoint['state_dict'].keys()):
if k[:7] == 'module.':
checkpoint['state_dict'][k[7:]] = checkpoint['state_dict'][k]
del checkpoint['state_dict'][k]
feature_network.load_state_dict(checkpoint['state_dict'])
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])