net = FPNSSD512(num_classes=9).to(device) net.load_state_dict(torch.load(args.model)) if device == 'cuda': net = torch.nn.DataParallel(net) cudnn.benchmark = True best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load(args.checkpoint) net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] criterion = SSDLoss(num_classes=9) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = inputs.to(device) loc_targets = loc_targets.to(device)
NUM_WORKERS = 2 trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS) testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS) net.cuda() # net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) cudnn.benchmark = True criterion = SSDLoss(num_classes=NUM_CLASSES) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda())
def main(**kwargs): opt._parse(kwargs) vis = Visualizer(env=opt.env) # Model print('==> Building model..') net = DSOD(num_classes=21) start_epoch = 0 # start from epoch 0 or last epoch if opt.load_path is not None: print('==> Resuming from checkpoint..') checkpoint = torch.load(opt.load_path) net.load_state_dict(checkpoint['net']) # Dataset print('==> Preparing dataset..') box_coder = SSDBoxCoder(net) trainset = ListDataset(root=opt.data_root, list_file=[opt.voc07_trainval, opt.voc12_trainval], transform=Transform(box_coder, True)) trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.batch_size, shuffle=True, num_workers=8) net.cuda() net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) cudnn.benchmark = True criterion = SSDLoss(num_classes=21) optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.9, weight_decay=5e-4) best_map_ = 0 for epoch in range(start_epoch, start_epoch + 200): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in tqdm(enumerate(trainloader)): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad() loc_preds, cls_preds = net(inputs) loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss.backward() optimizer.step() train_loss += loss.data[0] if (batch_idx + 1) % opt.plot_every == 0: vis.plot('loss', train_loss / (batch_idx + 1)) img = predict( net, box_coder, os.path.join(opt.data_root, trainset.fnames[batch_idx])) vis.img('predict', np.array(img).transpose(2, 0, 1)) if os.path.exists(opt.debug_file): import ipdb ipdb.set_trace() aps = eval(net.module, test_num=epoch * 100 + 100) map_ = aps['map'] if map_ > best_map_: print('Saving..') state = { 'net': net.state_dict(), 'map': best_map_, 'epoch': epoch, } best_map_ = map_ if not os.path.isdir(os.path.dirname(opt.checkpoint)): os.mkdir(os.path.dirname(opt.checkpoint)) torch.save(state, opt.checkpoint + '/%s.pth' % best_map_)
def main(**kwargs): opt._parse(kwargs) vis = Visualizer(env=opt.env) # Model print('==> Building model..') net = DSOD(num_classes=21) # Dataset print('==> Preparing dataset..') box_coder = SSDBoxCoder(net) trainset = ListDataset(root=opt.train_img_root, list_file=opt.train_img_list, transform=Transform(box_coder, True)) trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_worker, pin_memory=True) net.cuda() net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) cudnn.benchmark = True criterion = SSDLoss(num_classes=21) optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.9, weight_decay=1e-4) best_map_ = 0 best_loss = 1e100 start_epoch = 0 if opt.load_path is not None: print('==> Resuming from checkpoint..') checkpoint = torch.load(opt.load_path) net.load_state_dict(checkpoint['net']) best_loss = checkpoint['map'] start_epoch = checkpoint['epoch'] + 1 print('start_epoch = ', start_epoch, 'best_loss = ', best_loss) for epoch in range(start_epoch, start_epoch + 100): print('\nEpoch: ', epoch) net.train() train_loss = 0 optimizer.zero_grad() ix = 0 for batch_idx, (inputs, loc_targets, cls_targets) in tqdm(enumerate(trainloader)): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) loc_preds, cls_preds = net(inputs) ix += 1 loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss.backward() train_loss += loss.data.item() current_loss = train_loss / (1 + batch_idx) if (batch_idx + 1) % (opt.iter_size) == 0: for name, p in net.named_parameters(): p.grad.data.div_(ix) ix = 0 optimizer.step() optimizer.zero_grad() if (batch_idx + 1) % opt.plot_every == 0: vis.plot('loss', current_loss) # img = predict(net, box_coder, os.path.join(opt.train_img_root, trainset.fnames[batch_idx])) # vis.img('predict', np.array(img).transpose(2, 0, 1)) # if os.path.exists(opt.debug_file): # import ipdb # ipdb.set_trace() print('current_loss: ', current_loss, 'best_loss: ', best_loss) if (epoch + 1) % 20 == 0: for param_group in optimizer.param_groups: param_group['lr'] *= 0.1 if (epoch + 1) % opt.save_state_every == 0: state = { 'net': net.state_dict(), 'map': current_loss, 'epoch': epoch, } torch.save(state, opt.checkpoint + '%s.pth' % epoch) if current_loss < best_loss: best_loss = current_loss print('saving model at epoch: ', epoch) state = { 'net': net.state_dict(), 'map': best_loss, 'epoch': epoch, } torch.save(state, opt.checkpoint + 'dsod.pth')
def main(**kwargs): opt._parse(kwargs) vis = Visualizer(env=opt.env) # Model print('==> Building model..') net = DSOD(num_classes=21) start_epoch = 0 # start from epoch 0 or last epoch # Dataset print('==> Preparing dataset..') box_coder = SSDBoxCoder(net) trainset = ListDataset(root=opt.data_root, list_file=[opt.voc07_trainval, opt.voc12_trainval], transform=Transform(box_coder, True)) trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.batch_size, shuffle=True, num_workers=8, pin_memory=True) net.cuda() net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) cudnn.benchmark = True if opt.load_path is not None: print('==> Resuming from checkpoint..') checkpoint = torch.load(opt.load_path) net.module.load_state_dict(checkpoint['net']) criterion = SSDLoss(num_classes=21) optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.9, weight_decay=5e-4) best_map_ = 0 best_loss = 1e100 for epoch in range(start_epoch, start_epoch + 200): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 optimizer.zero_grad() ix = 0 for batch_idx, (inputs, loc_targets, cls_targets) in tqdm(enumerate(trainloader)): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) loc_preds, cls_preds = net(inputs) ix += 1 loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss.backward() train_loss += loss.data[0] if (batch_idx + 1) % (opt.iter_size) == 0: # if True: for name, p in net.named_parameters(): p.grad.data.div_(ix) ix = 0 optimizer.step() optimizer.zero_grad() if (batch_idx + 1) % opt.plot_every == 0: vis.plot('loss', train_loss / (batch_idx + 1)) img = predict( net, box_coder, os.path.join(opt.data_root, trainset.fnames[batch_idx])) vis.img('predict', np.array(img).transpose(2, 0, 1)) if os.path.exists(opt.debug_file): import ipdb ipdb.set_trace() # if (epoch+1)%10 == 0 : # state = { # 'net': net.module.state_dict(), # # 'map': best_map_, # 'epoch': epoch, # } # torch.save(state, opt.checkpoint + '/%s.pth' % epoch) # if (epoch+1) % 30 == 0: # for param_group in optimizer.param_groups: # param_group['lr'] *= 0.1 current_loss = train_loss / (1 + batch_idx) if current_loss < best_loss: best_loss = current_loss torch.save(net.module.state_dict(), '/tmp/dsod.pth') if (epoch + 1) % opt.eval_every == 0: net.module.load_state_dict(torch.load('/tmp/dsod.pth')) aps = eval(net.module) map_ = aps['map'] if map_ > best_map_: print('Saving..') state = { 'net': net.module.state_dict(), 'map': best_map_, 'epoch': epoch, } best_map_ = map_ if not os.path.isdir(os.path.dirname(opt.checkpoint)): os.mkdir(os.path.dirname(opt.checkpoint)) best_path = opt.checkpoint + '/%s.pth' % best_map_ torch.save(state, best_path) else: net.module.load_state_dict(torch.load(best_path)['net']) for param_group in optimizer.param_groups: param_group['lr'] *= 0.1 vis.log( dict(epoch=(epoch + 1), map=map_, loss=train_loss / (batch_idx + 1)))