def train(epoch): train_loader = torch.utils.data.DataLoader( lmdb_utils.lmdbDataset(traindb, shape=(model.module.width, model.module.height), shuffle=True, train=True, seen=model.module.seen), batch_size=batch_size, shuffle=False, **kwargs) logging('epoch %d : processed %d samples' % (epoch, epoch * len(train_loader.dataset))) model.train() adjust_learning_rate(optimizer, epoch) for batch_idx, (data, target) in enumerate(train_loader): if (batch_idx+1) % 70 == 0: sys.stdout.write('.') if use_cuda: data = data.cuda() #target= target.cuda() data, target = Variable(data), Variable(target) optimizer.zero_grad() output = model(data) loss = region_loss(output, target) loss.backward() optimizer.step() print('') logging('save weights to %s/%06d.weights' % (backupdir, epoch+1)) model.module.seen = (epoch + 1) * len(train_loader.dataset) model.module.save_weights('%s/%06d.weights' % (backupdir, epoch+1))
############### torch.manual_seed(seed) if use_cuda: torch.cuda.manual_seed(seed) model = Darknet(cfgfile) region_loss = model.loss model.load_weights(weightfile) model.print_network() init_epoch = model.seen / nsamples kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {} test_loader = torch.utils.data.DataLoader( lmdb_utils.lmdbDataset(testdb, shape=(160, 160), shuffle=False, transform=None, train=False), batch_size=batch_size, shuffle=False, **kwargs) if use_cuda: model = torch.nn.DataParallel(model).cuda() optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum) def adjust_learning_rate(optimizer, epoch): """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" lr = learning_rate * (0.1 ** (epoch // 50)) for param_group in optimizer.param_groups: param_group['lr'] = lr logging('set lr=%f' % (lr))
eps = 1e-5 ############### torch.manual_seed(seed) if use_cuda: torch.cuda.manual_seed(seed) model = Darknet(cfgfile) region_loss = model.loss model.load_weights(weightfile) model.print_network() init_epoch = model.seen / nsamples kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {} test_loader = torch.utils.data.DataLoader(lmdb_utils.lmdbDataset( testdb, shape=(160, 160), shuffle=False, transform=None, train=False), batch_size=batch_size, shuffle=False, **kwargs) if use_cuda: model = torch.nn.DataParallel(model).cuda() optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum) def adjust_learning_rate(optimizer, epoch): """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" lr = learning_rate * (0.1**(epoch // 50)) for param_group in optimizer.param_groups: param_group['lr'] = lr