default='12',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--ft_model_path',
default='/home/luojh2/.torch/models/vgg16-397923af.pth',
type=str,
help='the path of fine tuned model')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
print(args)
# Phase 1 : Data Upload
print('\n[Phase 1] : Data Preperation')
dset_loaders = {
'train':
torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True),
'val':
torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
}
print('data_loader_success!')
# Phase 2 : Model setup
print('\n[Phase 2] : Model setup')
def main():
global args, best_prec1
args = parser.parse_args()
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
# model = models.vgg16(pretrained=True)
model = vgg16_test('checkpoint/model.pth')
print(model)
model = torch.nn.DataParallel(model.cuda(),
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
cudnn.benchmark = True
# Data loading code from lmdb
if args.load_from_lmdb:
train_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
print('train_loader_success!')
val_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
else:
traindir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'train')
valdir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.Resize(256),
transforms.RandomCrop((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])),
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)
# evaluate and train
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, int(args.epochs / 3.0))
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
if is_best:
folder_path = 'checkpoint/fine_tune'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
torch.save(model.state_dict(), folder_path + '/model.pth')
print('best acc is %.3f' % best_prec1)
parser = argparse.ArgumentParser(description='PyTorch Digital Mammography Training')
parser.add_argument('--batch_size', default=100, type=int, help='batch size')
parser.add_argument('--data_base', default='/data/zhangcl/ImageNet', type=str, help='the path of dataset')
parser.add_argument('--gpu_id', default='1', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--ft_model_path', default='/home/luojh2/.torch/models/resnet50-19c8e357.pth',
type=str, help='the path of fine tuned model')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
print(args)
# Phase 1 : Data Upload
print('\n[Phase 1] : Data Preperation')
dset_loaders = {
'train': torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True),
'val': torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
}
print('data_loader_success!')
# Phase 2 : Model setup
print('\n[Phase 2] : Model setup')
def evaluate():
# Phase 1: load model
model = NetworkNew_test('../checkpoint/model.pth')
# Phase 2 : Model setup
model = model.cuda()
model = torch.nn.DataParallel(model.cuda(), device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
# Phase 2 : Data Upload
print('\n[Phase 2] : Data Preperation')
dset_loaders = {
'train': torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True),
'val': torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
}
print('data_loader_success!')
# Phase 3: Validation
print("\n[Phase 3 : Inference on val")
criterion = torch.nn.CrossEntropyLoss().cuda()
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for batch_idx, (input, target) in enumerate(dset_loaders['val']): # dset_loaders['val']):
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input, volatile=True)
target_var = torch.autograd.Variable(target, volatile=True)
# compute output
output = model(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.data[0], input.size(0))
top1.update(prec1[0], input.size(0))
top5.update(prec5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
batch_idx, len(dset_loaders['val']), batch_time=batch_time, loss=losses,
top1=top1, top5=top5))
sys.stdout.flush()
print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
.format(top1=top1, top5=top5))
def main():
global args, best_prec1
# Phase 1 : Data Upload
print('\n[Phase 1] : Data Preperation')
train_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=16,
pin_memory=True)
print('data_loader_success!')
# Phase 2 : Model setup
print('\n[Phase 2] : Model setup')
if args.layer_id == 0:
model_ft = Network_FT.Vgg16(args.ft_model_path).cuda()
model_ft = torch.nn.DataParallel(model_ft)
model_param = model_ft.state_dict()
torch.save(model_param, 'checkpoint/model.pth')
model_ft = Network_FT.NetworkNew(args.layer_id).cuda()
weight = torch.load('checkpoint/layer_7/model.pth')
model_ft = torch.nn.DataParallel(model_ft)
model_ft.load_state_dict(weight)
cudnn.benchmark = True
print("model setup success!")
# Phase 3: fine_tune model
print('\n[Phase 3] : Model fine tune')
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model_ft.parameters()),
args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
scale_factor = 9.0
for epoch in range(args.start_epoch, args.num_epochs):
adjust_learning_rate(optimizer, epoch, 1)
# train for one epoch
channel_index, scale_factor = train(train_loader, model_ft, criterion,
optimizer, epoch, scale_factor)
# evaluate on validation set
prec1 = validate(val_loader, model_ft, criterion, channel_index)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
if is_best:
best_prec1 = prec1
folder_path = 'checkpoint/layer_' + str(args.layer_id)
if not os.path.exists(folder_path):
os.makedirs(folder_path)
torch.save(model_ft.state_dict(), folder_path + '/model.pth')
torch.save(channel_index, folder_path + '/channel_index.pth')
def main():
global args, best_prec1
args = parser.parse_args()
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size)
# create model
model = models.vgg16(pretrained=True)
model = torch.nn.DataParallel(model.cuda(), device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
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']
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(args.resume))
cudnn.benchmark = True
# Data loading code from lmdb
if args.load_from_lmdb:
train_loader = torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True
)
print('train_loader_success!')
val_loader = torch.utils.data.DataLoader(
lmdbDataset(os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True
)
else:
traindir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'train')
valdir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(traindir, transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])),
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)
# evaluate and train
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# 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(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
torch.save(model.state_dict(), 'checkpoint/model.pth')
def main():
global args, best_prec1, scale_factor_list, resnet_channel_number
# Phase 1 : Data Upload
print('\n[Phase 1] : Data Preperation')
train_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
print('data_loader_success!')
# Phase 2 : Model setup
print('\n[Phase 2] : Model setup')
if args.group_id == 0:
model_ft = models.resnet50(True).cuda()
model_ft = torch.nn.DataParallel(model_ft)
model_param = model_ft.state_dict()
torch.save(model_param, 'checkpoint/model.pth')
model_ft = Network_FT.NetworkNew(args.group_id).cuda()
model_ft = torch.nn.DataParallel(model_ft)
cudnn.benchmark = True
print("model setup success!")
# Phase 3: fine_tune model
print('\n[Phase 3] : Model fine tune')
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model_ft.parameters()),
args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
tmp = np.linspace(
1, 100, int(args.num_epochs * len(train_loader) / args.alpha_range))
scale_factor_list = np.ones(
[resnet_channel_number[args.group_id],
len(tmp)])
for tmp_i in range(resnet_channel_number[args.group_id]):
scale_factor_list[tmp_i, :] = tmp.copy()
reg_lambda = 10.0 * np.ones(resnet_channel_number[args.group_id])
for epoch in range(args.start_epoch, args.num_epochs):
adjust_learning_rate(optimizer, epoch, int(args.num_epochs / 2.0))
# train for one epoch
channel_index, reg_lambda = train(train_loader, model_ft, criterion,
optimizer, epoch, reg_lambda)
# evaluate on validation set
prec1 = validate(val_loader, model_ft, criterion, channel_index)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
if is_best:
best_prec1 = prec1
folder_path = 'checkpoint/group_' + str(args.group_id)
if not os.path.exists(folder_path):
os.makedirs(folder_path)
torch.save(model_ft.state_dict(), folder_path + '/model.pth')
if args.group_id == 3:
tmp = channel_index[0].copy()
tmp[:] = 1.0
channel_index.append(tmp.copy())
channel_index.append(tmp.copy())
torch.save(channel_index, folder_path + '/channel_index.pth')