def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict)
return model
def vgg16_depth(pretrained=False, **kwargs):
"""VGG 16-layer depth model (configuration "D")
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = VGG(make_layers(cfg['D']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
return model
def vgg16(pretrained=True, **kwargs):
"""VGG 16-layer model (configuration "D") - only the conv part
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
conf = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512]
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(conf), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
# Froze bottom layers
for module in model.features[0:10]:
for param in module.parameters():
param.requires_grad = False
return model.features
def __init__(self,
pretrained=True,
model='vgg16',
requires_grad=True,
remove_fc=True,
show_params=False):
VGG.__init__(self, make_layers(cfg[model]))
self.ranges = ranges[model]
if pretrained:
exec(
"self.load_state_dict(models.%s(pretrained=True).state_dict())"
% model)
if not requires_grad:
for param in super().parameters():
param.requires_grad = False
if remove_fc: # delete redundant fully-connected layer params, can save memory
del self.classifier
if show_params:
for name, param in self.named_parameters():
print(name, param.size())
def __init__(self, batch_norm=True, pretrained=True):
super(Extractor, self).__init__()
vgg = VGG(make_layers(cfg['D'], batch_norm))
if pretrained:
if batch_norm:
vgg.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
else:
vgg.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
self.features = vgg.features
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
cfg = [
64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512,
512, 512, 'M'
]
model = VGG(make_layers(cfg, batch_norm=True))
path_pre_model = "/home/jjchu/My_Research/SegGAN/pre_model/vgg16_bn-6c64b313.pth"
model.load_state_dict(torch.load(path_pre_model))
# torch.load(path)
# ('features.41.running_mean',)
# ('classifier.6.bias',)
# model_ft = prevgg16(pretrained=True)
# num_ftrs = model_ft.fc.in_features
# model_ft.fc = nn.Linear(num_ftrs, 2)
if use_gpu:
model_ft = model_ft.cuda()
def main():
global args, best_prec1, train_log, test_log
args = parser.parse_args()
dir_name = args.arch + '_' + datetime.datetime.now().strftime('%m%d_%H%M')
log_dir = os.path.join('logs', os.path.join('prune', dir_name))
checkpoint_dir = os.path.join('checkpoints', os.path.join('prune', dir_name))
os.makedirs(log_dir)
os.makedirs(checkpoint_dir)
train_log = Logger(os.path.join(log_dir, 'train.log'))
test_log = Logger(os.path.join(log_dir, 'test.log'))
config_log = Logger(os.path.join(log_dir, 'config.log'))
for k, v in vars(args).items():
config_log.write(content="{k} : {v}".format(k=k, v=v), wrap=True, flush=True)
config_log.close()
# create model
print("=" * 89)
print("=> creating model '{}'".format(args.arch))
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
vgg_cfg, batch_norm = checkpoint['cfg']
from torchvision.models.vgg import VGG, make_layers
model = VGG(make_layers(cfg=vgg_cfg, batch_norm=batch_norm), init_weights=False)
model.features = torch.nn.DataParallel(model.features)
model.cuda()
model.load_state_dict(checkpoint['state_dict'])
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
return
elif args.pretrained:
if args.pretrained == 'pytorch':
print("=> using pre-trained model from model zoo")
model = models.__dict__[args.arch](pretrained=True)
args.pretrained_parallel = False
else:
model = models.__dict__[args.arch]()
if args.pretrained_parallel:
model.features = torch.nn.DataParallel(model.features)
model.cuda()
if os.path.isfile(args.pretrained):
print("=> using pre-trained model '{}'".format(args.pretrained))
checkpoint = torch.load(args.pretrained)
model.load_state_dict(checkpoint['state_dict'])
if not args.pretrained_parallel:
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
print("=> no checkpoint found at '{}'".format(args.pretrained))
return
else:
model = models.__dict__[args.arch]()
model.features = torch.nn.DataParallel(model.features)
model.cuda()
# define loss function (criterion)
criterion = nn.CrossEntropyLoss().cuda()
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
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(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
if not args.resume:
rcn_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.rcn_batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
prune(train_loader=train_loader, val_loader=val_loader, rcn_loader=rcn_loader,
model=model, criterion=criterion)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(lr_decay_step=args.lr_decay_step,
optimizer=optimizer, epoch=epoch)
# train for one epoch
train(train_loader=train_loader, model=model, criterion=criterion,
optimizer=optimizer, epoch=epoch, log=True)
# evaluate on validation set
prec1 = validate(val_loader=val_loader, model=model,
criterion=criterion, epoch=epoch, log=True)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'cfg': get_vgg_cfg(model),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best=is_best, checkpoint_dir=checkpoint_dir)
def get_vgg(in_channels=3, **kwargs): model = VGG(make_layers(cfg['D'], in_channels), **kwargs) return model