def convert():
args = parser.parse_args()
switch_conv_bn_impl('DBB')
switch_deploy_flag(False)
train_model = build_model(args.arch)
if 'hdf5' in args.load:
from utils import model_load_hdf5
model_load_hdf5(train_model, args.load)
elif os.path.isfile(args.load):
print("=> loading checkpoint '{}'".format(args.load))
checkpoint = torch.load(args.load)
if 'state_dict' in checkpoint:
checkpoint = checkpoint['state_dict']
checkpoint = copy.deepcopy(checkpoint)
ckpt = {k.replace('module.', ''): v
for k, v in checkpoint.items()} # strip the names
train_model.load_state_dict(ckpt)
else:
print("=> no checkpoint found at '{}'".format(args.load))
for m in train_model.modules():
if hasattr(m, 'switch_to_deploy'):
m.switch_to_deploy()
torch.save(train_model.state_dict(), args.save)
def main(arch, model_path, output_path, input_shape=(224, 224), batch_size=1):
switch_conv_bn_impl('DBB')
switch_deploy_flag(True)
model = build_model(arch)
model.load_state_dict(torch.load(model_path))
dummy_input = torch.autograd.Variable(
torch.randn(batch_size, 3, input_shape[0], input_shape[1]))
torch.onnx.export(model,
dummy_input,
output_path,
verbose=True,
keep_initializers_as_inputs=True,
opset_version=12)
onnx_model = onnx.load(output_path) # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, output_path)
print('finished exporting onnx ')
def test():
args = parser.parse_args()
switch_deploy_flag(args.mode == 'deploy')
switch_conv_bn_impl(args.blocktype)
model = build_model(args.arch)
if not torch.cuda.is_available():
print('using CPU, this will be slow')
use_gpu = False
else:
model = model.cuda()
use_gpu = True
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
if 'hdf5' in args.weights:
from utils import model_load_hdf5
model_load_hdf5(model, args.weights)
elif os.path.isfile(args.weights):
print("=> loading checkpoint '{}'".format(args.weights))
checkpoint = torch.load(args.weights)
if 'state_dict' in checkpoint:
checkpoint = checkpoint['state_dict']
ckpt = {k.replace('module.', ''): v
for k, v in checkpoint.items()} # strip the names
model.load_state_dict(ckpt)
else:
print("=> no checkpoint found at '{}'".format(args.weights))
cudnn.benchmark = True
# Data loading code
valdir = os.path.join(args.data, 'val')
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, val_preprocess(224)),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
validate(val_loader, model, criterion, use_gpu)
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))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size, rank=args.rank)
# =========================== build model
from convnet_utils import switch_deploy_flag, switch_conv_bn_impl, build_model
switch_deploy_flag(False)
switch_conv_bn_impl(args.blocktype)
model = build_model(args.arch)
if gpu == 0:
for name, param in model.named_parameters():
print(name, param.size())
if not torch.cuda.is_available():
print('using CPU, this will be slow')
elif args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif 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
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = sgd_optimizer(model, args.lr, args.momentum, args.weight_decay)
lr_scheduler = CosineAnnealingLR(optimizer=optimizer, T_max=args.epochs * IMAGENET_TRAINSET_SIZE // args.batch_size // ngpus_per_node)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map model to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
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'])
lr_scheduler.load_state_dict(checkpoint['scheduler'])
print("=> 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')
trans = strong_train_preprocess(224) if 'ResNet' in args.arch else standard_train_preprocess(224)
print('aug is ', trans)
train_dataset = datasets.ImageFolder(traindir, trans)
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
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_dataset = datasets.ImageFolder(valdir, val_preprocess(224))
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, args)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
# adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args, lr_scheduler)
# 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)
if not args.multiprocessing_distributed or (args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer' : optimizer.state_dict(),
'scheduler': lr_scheduler.state_dict(),
}, is_best, filename='{}_{}.pth.tar'.format(args.arch, args.blocktype))