Esempio n. 1
0
    torch.set_default_tensor_type('torch.FloatTensor')

net = Yolact()
net.train()

if args.resume == 'latest':
    weight = glob.glob('weights/latest*')[0]
    net.load_weights(weight)
    resume_step = int(weight.split('.pth')[0].split('_')[-1])
    print(f'\nResume training with \'{weight}\'.\n')
elif args.resume and 'yolact' in args.resume:
    net.load_weights('weights/' + args.resume)
    resume_step = int(args.resume.split('.pth')[0].split('_')[-1])
    print(f'\nResume training with \'{args.resume}\'.\n')
else:
    net.init_weights(backbone_path='weights/' + cfg.backbone.path)
    print('\nTraining from begining, weights initialized.\n')

optimizer = optim.SGD(net.parameters(),
                      lr=cfg.lr,
                      momentum=cfg.momentum,
                      weight_decay=cfg.decay)
criterion = Multi_Loss(num_classes=cfg.num_classes,
                       pos_thre=cfg.pos_iou_thre,
                       neg_thre=cfg.neg_iou_thre,
                       np_ratio=3)

if cuda:
    cudnn.benchmark = True
    net = nn.DataParallel(net).cuda()
    criterion = nn.DataParallel(criterion).cuda()
Esempio n. 2
0
                      lr=cfg.lr,
                      momentum=0.9,
                      weight_decay=5e-4)
criterion = Multi_Loss(cfg)

if args.resume == 'latest':
    weight = glob.glob('weights/latest*')[0]
    net.load_weights(weight, cuda)
    start_step = int(weight.split('.pth')[0].split('_')[-1])
    print(f'\nResume training with \'{weight}\'.\n')
elif args.resume and 'yolact' in args.resume:
    net.load_weights(cfg.weight, cuda)
    start_step = int(cfg.weight.split('.pth')[0].split('_')[-1])
    print(f'\nResume training with \'{args.resume}\'.\n')
else:
    net.init_weights(cfg.weight)
    print(
        f'\nTraining from begining, weights initialized with {cfg.weight}.\n')
    start_step = 0

dataset = COCODetection(cfg, mode='train')
train_sampler = None
main_gpu = False
if cuda:
    cudnn.benchmark = True
    cudnn.fastest = True
    main_gpu = dist.get_rank() == 0
    num_gpu = dist.get_world_size()

    net_with_loss = NetWithLoss(net, criterion)
    net = DDP(net_with_loss.cuda(), [args.local_rank],