def __init__(self, loadweights=True, downsample=4, model_path='pretrained_model/shufflenetv2_x1_69.402_88.374.pth.tar'):
super(shufflenetv2_base, self).__init__()
model = shufflenetv2(width_mult=1.0)
if loadweights:
model.load_state_dict(torch.load(model_path))
self.feature3 = nn.Sequential(model.conv1, model.maxpool, model.features[:4])
self.feature4 = nn.Sequential(model.features[4:12])
self.feature5 = nn.Sequential(model.features[12:])
def main():
global args
args = parser.parse_args()
# create model
print("=> creating model '{}'".format(args.arch))
if args.arch.startswith('shufflenetv2'):
model = shufflenetv2(width_mult=args.width_mult)
elif args.arch.startswith('mobilenetv2'):
model = MobileNetV2()
else:
model = models.__dict__[args.arch]()
print(model)
if args.evaluate:
if os.path.isfile(args.evaluate):
print("=> loading model '{}'".format(args.evaluate))
model.load_state_dict(torch.load(args.evaluate))
else:
print("=> no model found at '{}'".format(args.evaluate))
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
# Data loading code
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
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)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
if args.evaluate:
validate(val_loader, model, criterion)
return
def build_refinedet(phase, size=320, num_classes=13):
if phase != "test" and phase != "train":
print("ERROR: Phase: " + phase + " not recognized")
return
if size != 320 and size != 512 and size != 768 and size != 960:
print("ERROR: You specified size " + repr(size) + ". However, " +
"currently only RefineDet320 and RefineDet512 is supported!")
return
base_ = shufflenetv2(1.)
# extras_ = add_extras(extras[str(size)], size, 1024)
ARM_ = arm_multibox_shuffle()
print('passed # classes = ', num_classes)
ODM_ = odm_multibox_shuffle(num_classes)
TCB_ = add_tcb_shuffle()
print('size = ', size) # debug
return RefineDet(phase, size, base_, ARM_, ODM_, TCB_, num_classes)
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
# create model
print("=> creating model '{}'".format(args.arch))
logging.info("=> creating model '{}'".format(args.arch))
if args.arch.startswith('shufflenetv2'):
model = shufflenetv2(ratio=args.ratio)
elif args.arch.startswith('shufflenetv1'):
model = shufflenetv1(groups=args.groups, ratio=args.ratio)
else:
model = models.__dict__[args.arch]()
# print(model)
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(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_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logging.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
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_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
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 args.evaluate:
if os.path.isfile(args.evaluate):
print("=> loading model '{}'".format(args.evaluate))
logging.info("=> loading model '{}'".format(args.evaluate))
checkpoint = torch.load(args.evaluate)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch']))
validate(val_loader, model, criterion, args)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)