def gluon2torch(name,
gluon_path,
torch_path,
base=False,
reorder=False,
force_pair=None):
name = name.lower()
if base:
torch_model = model_zoo.get_model(name,
pretrained=False,
pretrained_base=False,
root=torch_path)
gluon_model = get_model(name,
pretrained=True,
pretrained_base=False,
root=gluon_path)
else:
torch_model = model_zoo.get_model(name,
pretrained=False,
root=torch_path)
gluon_model = get_model(name, pretrained=True, root=gluon_path)
torch_keys = [
k for k in torch_model.state_dict().keys()
if not k.endswith('num_batches_tracked')
]
gluon_keys = gluon_model.collect_params().keys()
if reorder:
key_words = ('running_mean', 'running_var', 'moving_mean',
'moving_var')
if force_pair is not None:
torch_keys = [k for k in torch_keys if not k.endswith(key_words) and not k.startswith(force_pair[1])] + \
[k for k in torch_keys if not k.endswith(key_words) and k.startswith(force_pair[1])] + \
[k for k in torch_keys if k.endswith(key_words) and not k.startswith(force_pair[1])] + \
[k for k in torch_keys if k.endswith(key_words) and k.startswith(force_pair[1])]
gluon_keys = [k for k in gluon_keys if not k.endswith(key_words) and not k.startswith(force_pair[0])] + \
[k for k in gluon_keys if not k.endswith(key_words) and k.startswith(force_pair[0])] + \
[k for k in gluon_keys if k.endswith(key_words) and not k.startswith(force_pair[0])] + \
[k for k in gluon_keys if k.endswith(key_words) and k.startswith(force_pair[0])]
else:
torch_keys = [k for k in torch_keys if not k.endswith(key_words)] + \
[k for k in torch_keys if k.endswith(key_words)]
gluon_keys = [k for k in gluon_keys if not k.endswith(key_words)] + \
[k for k in gluon_keys if k.endswith(key_words)]
assert len(torch_keys) == len(gluon_keys)
map = dict(zip(gluon_keys, torch_keys))
pytorch_model_params = {}
print('Convert Gluon Model to PyTorch Model ...')
for key, value in gluon_model.collect_params().items():
tensor = torch.from_numpy(value.data().asnumpy())
tensor.require_grad = True
pytorch_model_params[map[key]] = tensor
torch.save(pytorch_model_params, os.path.join(torch_path, name + '.pth'))
print('Finish')
def parse_network(network, outputs, pretrained, **kwargs):
"""Parse network with specified outputs and other arguments.
Parameters
----------
network : str or nn.Module
Logic chain: load from gluoncv.model_zoo if network is string.
Convert to Symbol if network is HybridBlock
outputs : str or iterable of str
The name of layers to be extracted as features.
pretrained : bool
Use pretrained parameters as in model_zoo
Returns
-------
results: list of nn.Module (the same size as len(outputs))
"""
l, n = len(outputs), len(outputs[0])
results = [[] for _ in range(l)]
if isinstance(network, str):
from model.model_zoo import get_model
network = get_model(network, pretrained=pretrained, **kwargs).features
# helper func
def recursive(pos, block, arr, j):
if j == n:
results[pos].append([block])
return
child = list(block.children())
results[pos].append(child[:arr[j]])
if pos + 1 < l: results[pos + 1].append(child[arr[j] + 1:])
recursive(pos, child[arr[j]], arr, j + 1)
block = list(network.children())
for i in range(l):
pos = outputs[i][0]
if i == 0:
results[i].append(block[:pos])
elif i < l:
results[i].append(block[outputs[i - 1][0] + 1:pos])
recursive(i, block[pos], outputs[i], 1)
for i in range(l):
results[i] = nn.Sequential(
*[item for sub in results[i] for item in sub if sub])
return results
def __init__(self,
base_name='resnet50_v1b',
pretrained_base=False,
num_joints=17,
num_deconv_layers=3,
num_deconv_filters=(2048, 256, 256, 256),
num_deconv_kernels=(4, 4, 4),
final_conv_kernel=1,
deconv_with_bias=False,
**kwargs):
super(SimplePoseResNet, self).__init__(**kwargs)
from model.model_zoo import get_model
base_network = get_model(base_name, pretrained=pretrained_base)
self.resnet = list()
for layer in ['features']:
self.resnet.append(getattr(base_network, layer))
# if base_name.endswith('v1'):
# for layer in ['features']:
# self.resnet.append(getattr(base_network, layer))
# else:
# for layer in ['conv1', 'bn1', 'relu', 'maxpool',
# 'layer1', 'layer2', 'layer3', 'layer4']:
# self.resnet.append(getattr(base_network, layer))
self.resnet = nn.Sequential(*self.resnet)
self.deconv_with_bias = deconv_with_bias
# used for deconv layers
self.deconv_layers = self._make_deconv_layer(
num_deconv_layers,
num_deconv_filters,
num_deconv_kernels,
)
self.final_layer = nn.Conv2d(
num_deconv_filters[-1],
num_joints,
kernel_size=final_conv_kernel,
stride=1,
padding=1 if final_conv_kernel == 3 else 0,
)
def train(cfg, local_rank, distributed):
pretrained_base = os.path.join(cfg.TRAIN.model_root,
cfg.TRAIN.backbone + '.pth')
model = get_model(cfg.MODEL.model, pretrained_base=pretrained_base)
device = torch.device(cfg.MODEL.device)
model.to(device)
if distributed:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
optimizer = make_optimizer(cfg, model)
scheduler = make_lr_scheduler(cfg, optimizer)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
)
arguments = {}
arguments["iteration"] = 0
output_dir = cfg.CONFIG.output_dir
save_to_disk = ptutil.get_rank() == 0
checkpointer = ptutil.CheckPointer(model, optimizer, scheduler, output_dir,
save_to_disk)
extra_checkpoint_data = checkpointer.load(cfg.TRAIN.weight)
arguments.update(extra_checkpoint_data)
data_loader = build_dataloader(
cfg,
train=True,
distributed=distributed,
start_iter=arguments["iteration"],
)
checkpoint_period = cfg.TRAIN.checkpoint_period
training(model, data_loader, optimizer, scheduler, checkpointer, device,
checkpoint_period, arguments)
return model
def gluon2torch(name, gluon_path, torch_path, num):
name = name.lower()
torch_model = model_zoo.get_model(name, pretrained=False, root=torch_path)
gluon_model = get_model(name, pretrained=True, root=gluon_path)
torch_keys = [k for k in torch_model.state_dict().keys() if not k.endswith('num_batches_tracked')]
gluon_keys = gluon_model.collect_params().keys()
assert len(torch_keys) == len(gluon_keys)
map = dict(zip(gluon_keys, torch_keys))
pytorch_model_params = {}
print('Convert Gluon Model to PyTorch Model ...')
for i, ((key, value), (key2, value2)) in enumerate(
zip(gluon_model.collect_params().items(), torch_model.state_dict().items())):
if i < num:
tensor = torch.from_numpy(value.data().asnumpy())
tensor.require_grad = True
pytorch_model_params[map[key]] = tensor
else:
pytorch_model_params[map[key]] = value2
torch.save(pytorch_model_params, os.path.join(torch_path, name + '.pth'))
print('Finish')
default=os.path.expanduser('~/.torch/models'),
help='Default pre-trained mdoel root.')
opt = parser.parse_args()
return opt
if __name__ == '__main__':
opt = parse_args()
device = torch.device('cpu')
if opt.cuda:
device = torch.device('cuda:0')
# Load Model
model_name = opt.model
pretrained = True if opt.saved_params == '' else False
model = get_model(model_name,
pretrained=pretrained,
pretrained_base=False,
root=opt.root).to(device)
model.eval()
# Load Images
if opt.input_pic is None:
img_map = {
'voc': 'voc_example.jpg',
'ade': 'ade_example.jpg',
'coco': 'voc_example.jpg',
'citys': 'city_example.jpg'
}
opt.input_pic = os.path.join(
cur_path, '../png/' + img_map[model_name.split('_')[-1]])
img = Image.open(opt.input_pic)
torch.backends.cudnn.benchmark = False if args.mode == 'testval' else True
device = torch.device('cuda')
else:
distributed = False
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method=args.init_method)
# Load Model
model = model_zoo.get_model(args.model_name,
pretrained=True,
pretrained_base=False,
base_size=args.base_size,
crop_size=args.crop_size,
root=args.root,
aux=args.aux,
dilated=args.dilated,
jpu=args.jpu)
model.keep_shape = True if args.mode == 'testval' else False
model.to(device)
# testing data
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
data_kwargs = {
'base_size': args.base_size,
device = torch.device('cuda')
else:
distributed = False
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method=args.init_method)
# Load Model
model_name = args.network
if args.pretrained.lower() in ['true', '1', 'yes', 't']:
pretrained = True
else:
pretrained = False
kwargs = {'classes': 10, 'pretrained': pretrained, 'root': args.root}
net = model_zoo.get_model(model_name, **kwargs)
net.to(device)
# testing data
val_metric = Accuracy()
val_data = get_dataloader(args.batch_size, args.num_workers,
args.data_root, distributed)
# testing
metric = validate(net, val_data, device, val_metric)
synchronize()
name, value = accumulate_metric(metric)
if is_main_process():
print(name, value)
distributed = num_gpus > 1
if args.cuda and torch.cuda.is_available():
cudnn.benchmark = True
device = torch.device('cuda')
else:
distributed = False
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl", init_method=args.init_method)
# network
net_name = '_'.join((args.algorithm, args.network, args.dataset))
args.save_prefix += net_name
if args.pretrained.lower() in ['true', '1', 'yes', 't']:
net = model_zoo.get_model(net_name, pretrained=True)
else:
net = model_zoo.get_model(net_name, pretrained=False)
net.load_parameters(args.pretrained.strip())
net.to(device)
net.set_nms(nms_thresh=0.45, nms_topk=400)
# testing data
val_dataset, val_metric = get_dataset(args.dataset, args.data_shape)
val_data = get_dataloader(val_dataset, args.batch_size, args.num_workers, distributed, args.dataset == 'coco')
classes = val_dataset.classes # class names
# testing
val_metric = validate(net, val_data, device, val_metric, args.dataset == 'coco')
synchronize()
help='demo with GPU')
parser.add_argument('--root', type=str, default=os.path.expanduser('~/.torch/models'),
help='Default pre-trained mdoel root.')
parser.add_argument('--pretrained', type=str, default='True',
help='Load weights from previously saved parameters.')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
device = torch.device('cpu')
if args.cuda:
device = torch.device('cuda')
image = args.images
net = get_model(args.network, pretrained=True)
net.to(device)
net.set_nms(0.3, 200)
net.eval()
ax = None
x, img = load_test(image, short=net.short, max_size=net.max_size)
x = x.to(device)
with torch.no_grad():
ids, scores, bboxes, masks = [xx.cpu().numpy() for xx in net(x)]
masks = expand_mask(masks, bboxes, (img.shape[1], img.shape[0]), scores)
img = plot_mask(img, masks)
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(1, 1, 1)
ax = plot_bbox(img, bboxes, scores, ids,
class_names=net.classes, ax=ax)
import time
import numpy as np
import mxnet as mx
from gluoncv.utils import viz
from mxnet import gpu
from mxnet.ndarray import concat
from model import model_zoo
from os.path import basename
import gluoncv as gcv
classes = ['text']
imagePath = ''
netName = 'textboxes_512_mobilenet1.0_custom'
path_to_model = 'textBoxes_512_mobilenet1.0_text_ICDAR_new.params'
gpu_ind = 0
output_path = ''
net, _ = model_zoo.get_model(netName, classes=classes, pretrained_base=False)
net.reset_class(classes)
net.load_parameters(path_to_model)
start = time.time()
x, image = gcv.data.transforms.presets.ssd.load_test(imagePath,
scale,
max_size=2048)
cid, score, bbox = net(y)
ax = viz.plot_bbox(image, bbox[0], score[0], cid[0], class_names=classes)
base_image_name = basename(imagePath)
plt.savefig(output_path + base_image_name)
print(time.time() - start)
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl", init_method=args.init_method)
# network
kwargs = {}
module_list = []
if args.use_fpn:
module_list.append('fpn')
if args.norm_layer is not None:
module_list.append(args.norm_layer)
net_name = '_'.join(('faster_rcnn', *module_list, args.network, args.dataset))
args.save_prefix += net_name
if args.pretrained.lower() in ['true', '1', 'yes', 't']:
net = model_zoo.get_model(net_name, pretrained=True, root=args.root, **kwargs)
else:
net = model_zoo.get_model(net_name, pretrained=False, **kwargs)
net.load_state_dict(args.pretrained.strip())
net.to(device)
# testing data
val_dataset, val_metric = get_dataset(net.short, net.max_size, args.dataset)
val_data = get_dataloader(val_dataset, args.batch_size, args.num_workers,
distributed, args.dataset == 'coco')
classes = val_dataset.classes # class names
# testing
val_metric = validate(net, val_data, device, val_metric, args.dataset == 'coco')
synchronize()
if args.cuda and torch.cuda.is_available():
cudnn.benchmark = True
device = torch.device('cuda')
else:
distributed = False
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method=args.init_method)
input_size = [int(i) for i in args.input_size.split(',')]
val_list = get_dataloader(args.data_dir, args.batch_size, args.num_workers,
input_size, args.mean, args.std, distributed)
val_metric = COCOKeyPointsMetric(val_list[0],
'coco_keypoints',
data_shape=tuple(input_size),
in_vis_thresh=args.score_threshold,
cleanup=True)
use_pretrained = True if not args.params_file else False
net = model_zoo.get_model(args.model,
num_joints=args.num_joints,
pretrained=use_pretrained).to(device)
net.eval()
val_metric = validate(val_list, net, val_metric, device, args.flip_test)
synchronize()
name, value = accumulate_metric(val_metric)
if is_main_process():
print(name, value)
def __init__(self, args):
self.device = torch.device(args.device)
self.save_prefix = '_'.join((args.model, args.backbone, args.dataset))
# image transform
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
# dataset and dataloader
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_segmentation_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
args.per_iter = len(trainset) // (args.num_gpus * args.batch_size)
args.max_iter = args.epochs * args.per_iter
if args.distributed:
sampler = data.DistributedSampler(trainset)
else:
sampler = data.RandomSampler(trainset)
train_sampler = data.sampler.BatchSampler(sampler, args.batch_size,
True)
train_sampler = IterationBasedBatchSampler(
train_sampler, num_iterations=args.max_iter)
self.train_loader = data.DataLoader(trainset,
batch_sampler=train_sampler,
pin_memory=True,
num_workers=args.workers)
if not args.skip_eval or 0 < args.eval_epochs < args.epochs:
valset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
val_sampler = make_data_sampler(valset, False, args.distributed)
val_batch_sampler = data.sampler.BatchSampler(
val_sampler, args.test_batch_size, False)
self.valid_loader = data.DataLoader(
valset,
batch_sampler=val_batch_sampler,
num_workers=args.workers,
pin_memory=True)
# create network
if args.model_zoo is not None:
self.net = get_model(args.model_zoo, pretrained=True)
else:
kwargs = {'oc': args.oc} if args.model == 'ocnet' else {}
self.net = get_segmentation_model(model=args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
dilated=args.dilated,
jpu=args.jpu,
crop_size=args.crop_size,
**kwargs)
if args.distributed:
self.net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.net)
self.net.to(self.device)
# resume checkpoint if needed
if args.resume is not None:
if os.path.isfile(args.resume):
self.net.load_state_dict(torch.load(args.resume))
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
# create criterion
if args.ohem:
min_kept = args.batch_size * args.crop_size**2 // 16
self.criterion = OHEMSoftmaxCrossEntropyLoss(thresh=0.7,
min_kept=min_kept,
use_weight=False)
else:
self.criterion = MixSoftmaxCrossEntropyLoss(
args.aux, aux_weight=args.aux_weight)
# optimizer and lr scheduling
params_list = [{
'params': self.net.base1.parameters(),
'lr': args.lr
}, {
'params': self.net.base2.parameters(),
'lr': args.lr
}, {
'params': self.net.base3.parameters(),
'lr': args.lr
}]
if hasattr(self.net, 'others'):
for name in self.net.others:
params_list.append({
'params':
getattr(self.net, name).parameters(),
'lr':
args.lr * 10
})
if hasattr(self.net, 'JPU'):
params_list.append({
'params': self.net.JPU.parameters(),
'lr': args.lr * 10
})
self.optimizer = optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.scheduler = WarmupPolyLR(self.optimizer,
T_max=args.max_iter,
warmup_factor=args.warmup_factor,
warmup_iters=args.warmup_iters,
power=0.9)
if args.distributed:
self.net = torch.nn.parallel.DistributedDataParallel(
self.net,
device_ids=[args.local_rank],
output_device=args.local_rank)
# evaluation metrics
self.metric = SegmentationMetric(trainset.num_class)
self.args = args
help='demo with GPU')
parser.add_argument('--input-pic', type=str, default=os.path.join(cur_path, '../png/mt_baker.jpg'),
help='path to the input picture')
opt = parser.parse_args()
return opt
if __name__ == '__main__':
opt = parse_args()
device = torch.device('cpu')
if opt.cuda:
device = torch.device('cuda:0')
# Load Model
model_name = opt.model
pretrained = True if opt.saved_params == '' else False
net = get_model(model_name, pretrained=pretrained, root=opt.root).to(device)
net.eval()
# Load Images
img = Image.open(opt.input_pic)
# Transform
transform_fn = transforms.Compose([
transforms.Resize(256 if model_name.lower() != 'inceptionv3' else 299),
transforms.CenterCrop(224 if model_name.lower() != 'inceptionv3' else 299),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = transform_fn(img).to(device)
with torch.no_grad():
default='../configs/retina_resnet101_v1b_coco.yaml')
parser.add_argument('--images', type=str, default=os.path.join(cur_path, '../png/biking.jpg'),
help='Test images.')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
cfg.merge_from_file(args.config_file)
cfg.freeze()
device = torch.device(cfg.MODEL.device)
device_cpu = torch.device('cpu')
image = args.images
net = get_model(cfg.CONFIG.model, pretrained=cfg.TEST.pretrained)
net.to(device)
net.eval()
ax = None
x, img = load_test(image, min_image_size=800)
x = to_image_list(x, 32)
x = x.to(device)
with torch.no_grad():
predictions = net(x)
predictions = [o.to(device_cpu) for o in predictions]
prediction = predictions[0]
height, width = img.shape[:-1]
prediction = prediction.resize((width, height))
top_predictions = select_top_predictions(prediction, conf_thresh=cfg.TEST.vis_thresh)
opt = parser.parse_args()
return opt
if __name__ == '__main__':
opt = parse_args()
device = torch.device('cpu')
if opt.cuda:
device = torch.device('cuda')
# Load Model
model_name = opt.model
pretrained = True if opt.saved_params == '' else False
model = get_model(model_name,
pretrained=pretrained,
pretrained_base=False,
aux=opt.aux,
dilated=opt.dilated,
jpu=opt.jpu,
root=opt.root).to(device)
model.eval()
# Load Images
if opt.input_pic is None:
img_map = {'voc': 'voc_example.jpg', 'citys': 'city_example.jpg'}
opt.input_pic = os.path.join(
cur_path, 'png/' + img_map[model_name.split('_')[-1]])
img = Image.open(opt.input_pic)
# Transform
transform_fn = transforms.Compose([
transforms.ToTensor(),
help='demo with GPU')
parser.add_argument('--root', type=str, default=os.path.expanduser('~/.torch/models'),
help='Default pre-trained model root')
parser.add_argument('--input-pic', type=str, default=os.path.join(cur_path, '../png/soccer.png'),
help='path to the input picture')
opt = parser.parse_args()
return opt
if __name__ == '__main__':
opt = parse_args()
device = torch.device('cpu')
if opt.cuda:
device = torch.device('cuda')
detector = get_model(opt.detector, pretrained=True).to(device)
detector.reset_class(["person"], reuse_weights=['person'])
pose_net = get_model(opt.pose_model, pretrained=True).to(device)
detector.eval()
pose_net.eval()
x, img = load_test(opt.input_pic, short=512)
x = x.to(device)
with torch.no_grad():
class_IDs, scores, bounding_boxs = detector(x)
pose_input, upscale_bbox = detector_to_simple_pose(img, class_IDs, scores, bounding_boxs, device=device)
predicted_heatmap = pose_net(pose_input)
pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox)
plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores,
box_thresh=0.5, keypoint_thresh=0.2)
if __name__ == '__main__':
# config
classes = 10
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck']
args = parse_args()
device = torch.device('cpu')
if args.cuda:
device = torch.device('cuda:0')
# Load Model
model_name = args.model
pretrained = True if args.saved_params == '' else False
kwargs = {'classes': classes, 'pretrained': pretrained, 'root': args.root, }
net = get_model(model_name, **kwargs).to(device)
net.eval()
# Load Images
img = Image.open(args.input_pic)
# Transform
transform_fn = transforms.Compose([
transforms.Resize(32),
transforms.CenterCrop(32),
transforms.ToTensor(),
transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
])
img = transform_fn(img).to(device)
with torch.no_grad():
parser.add_argument(
'--thresh',
type=float,
default=0.5,
help='Threshold of object score when visualize the bboxes.')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
device = torch.device('cpu')
if args.cuda:
device = torch.device('cuda')
image = args.images
net = get_model(args.network, pretrained=True, root=args.root)
net.to(device)
net.set_nms(0.45, 200)
net.eval()
ax = None
x, img = load_test(image, short=512)
x = x.to(device)
with torch.no_grad():
ids, scores, bboxes = [xx[0].cpu().numpy() for xx in net(x)]
ax = plot_bbox(img,
bboxes,
scores,
ids,
thresh=args.thresh,
class_names=net.classes,
dataset = gcv.data.LstDetection(LSTpath, root=images_root)
print(dataset)
image = dataset[0][0]
label = dataset[0][1]
print('label:', label)
# display image and label
ax = viz.plot_bbox(image,
bboxes=label[:, :4],
labels=label[:, 4:5],
class_names=classes)
plt.savefig('labeled_image.jpg')
#initalize model
net, input_size = model_zoo.get_model(netName,
pretrained=False,
classes=classes)
if finetune_model == '':
net.initialize()
net.reset_class(classes)
else:
net.load_parameters(path_to_model)
net.reset_class(classes)
print(net)
train_data = get_dataloader(net, dataset, input_size, batch_size, 0)
#############################################################################################
# Try use GPU for training
try:
def __init__(self, args):
self.device = torch.device(args.device)
# network
net_name = '_'.join(('yolo3', args.network, args.dataset))
self.save_prefix = net_name
self.net = get_model(net_name, pretrained_base=True)
if args.distributed:
self.net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.net)
if args.resume.strip():
logger.info("Resume from the model {}".format(args.resume))
self.net.load_state_dict(torch.load(args.resume.strip()))
else:
logger.info("Init from base net {}".format(args.network))
classes, anchors = self.net.num_class, self.net.anchors
self.net.set_nms(nms_thresh=0.45, nms_topk=400)
if args.label_smooth:
self.net._target_generator._label_smooth = True
self.net.to(self.device)
if args.distributed:
self.net = torch.nn.parallel.DistributedDataParallel(
self.net, device_ids=[args.local_rank], output_device=args.local_rank)
# dataset and dataloader
train_dataset = get_train_data(args.dataset, args.mixup)
width, height = args.data_shape, args.data_shape
batchify_fn = Tuple(*([Stack() for _ in range(6)] + [Pad(axis=0, pad_val=-1) for _ in range(1)]))
train_dataset = train_dataset.transform(
YOLO3DefaultTrainTransform(width, height, classes, anchors, mixup=args.mixup))
args.per_iter = len(train_dataset) // (args.num_gpus * args.batch_size)
args.max_iter = args.epochs * args.per_iter
if args.distributed:
sampler = data.DistributedSampler(train_dataset)
else:
sampler = data.RandomSampler(train_dataset)
train_sampler = data.sampler.BatchSampler(sampler=sampler, batch_size=args.batch_size,
drop_last=False)
train_sampler = IterationBasedBatchSampler(train_sampler, num_iterations=args.max_iter)
if args.no_random_shape:
self.train_loader = data.DataLoader(train_dataset, batch_sampler=train_sampler, pin_memory=True,
collate_fn=batchify_fn, num_workers=args.num_workers)
else:
transform_fns = [YOLO3DefaultTrainTransform(x * 32, x * 32, classes, anchors, mixup=args.mixup)
for x in range(10, 20)]
self.train_loader = RandomTransformDataLoader(transform_fns, train_dataset, batch_sampler=train_sampler,
collate_fn=batchify_fn, num_workers=args.num_workers)
if args.eval_epoch > 0:
# TODO: rewrite it
val_dataset, self.metric = get_test_data(args.dataset)
val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
val_dataset = val_dataset.transform(YOLO3DefaultValTransform(width, height))
val_sampler = make_data_sampler(val_dataset, False, args.distributed)
val_batch_sampler = data.BatchSampler(val_sampler, args.test_batch_size, False)
self.val_loader = data.DataLoader(val_dataset, batch_sampler=val_batch_sampler,
collate_fn=val_batchify_fn, num_workers=args.num_workers)
# optimizer and lr scheduling
self.optimizer = optim.SGD(self.net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.wd)
if args.lr_mode == 'cos':
self.scheduler = WarmupCosineLR(optimizer=self.optimizer, T_max=args.max_iter,
warmup_factor=args.warmup_factor, warmup_iters=args.warmup_iters)
elif args.lr_mode == 'step':
lr_decay = float(args.lr_decay)
milestones = sorted([float(ls) * args.per_iter for ls in args.lr_decay_epoch.split(',') if ls.strip()])
self.scheduler = WarmupMultiStepLR(optimizer=self.optimizer, milestones=milestones, gamma=lr_decay,
warmup_factor=args.warmup_factor, warmup_iters=args.warmup_iters)
else:
raise ValueError('illegal scheduler type')
self.args = args
from torch import nn
def net_xavier_uniform_init(net):
for m in net.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_uniform_(m.weight)
nn.init.zeros_(m.bias)
if __name__ == '__main__':
from model.model_zoo import get_model
net_name = 'ssd_512_mobilenet1.0_coco'
net = get_model(net_name, pretrained_base=True)
net_xavier_uniform_init(net)
