Ejemplo n.º 1
0
def train():

    learning_rate = 0.001
    num_epochs = 300
    batch_size = 32

    net = FaceBox()
    print('load model...')
    net.load_state_dict(torch.load('weight/faceboxes.pt'))
    traced_script_module = torch.jit.trace(net, torch.rand(1, 3, 1024, 1024))
    traced_script_module.save("faceboxes.pt")
Ejemplo n.º 2
0
    f_write.close()


def getFddbList(path, file_name):
    with open(path + file_name) as f:
        file_list = f.readlines()
    f_write = open(path + 'fddblist.txt', 'w')
    for item in file_list:
        if '/' in item:
            f_write.write(item)
    f_write.close()
    print('get fddb list done')


if __name__ == '__main__':
    net = FaceBox()
    net.load_state_dict(
        torch.load('weight/faceboxes.pt',
                   map_location=lambda storage, loc: storage))

    if use_gpu:
        net.cuda()
    net.eval()
    data_encoder = DataEncoder()

    font = cv2.FONT_HERSHEY_SCRIPT_SIMPLEX

    # given video path, predict and show
    path = "/home/lxg/codedata/faceVideo/1208.mp4"
    # testVideo(path)
Ejemplo n.º 3
0
from networks import FaceBox
from multibox_loss import MultiBoxLoss
from dataset import ListDataset

import visdom
import numpy as np

use_gpu = torch.cuda.is_available()
file_root = '/home/lxg/codedata/'

learning_rate = 0.001
num_epochs = 300
batch_size = 64

net = FaceBox()
if use_gpu:
    net.cuda()

print('load model...')
# net.load_state_dict(torch.load('weight/faceboxes.pt'))

criterion = MultiBoxLoss()

# optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0003)
optimizer = torch.optim.Adam(net.parameters(),
                             lr=learning_rate,
                             weight_decay=1e-4)

train_dataset = ListDataset(root=file_root,
                            list_file='label/box_label.txt',
Ejemplo n.º 4
0
def train():
    use_gpu = torch.cuda.is_available()
    file_root = os.path.dirname(os.path.abspath(__file__))

    learning_rate = 0.001
    num_epochs = 300
    batch_size = 32

    net = FaceBox()
    if use_gpu:
        net.cuda()

    print('load model...')
    net.load_state_dict(torch.load('weight/faceboxes_add_norm.pt'))


    criterion = MultiBoxLoss()
    #optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0005)
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[198, 248], gamma=0.1)

    train_dataset = ListDataset(root=file_root,
                                list_file='data/train_rewrite.txt',
                                train=True,
                                transform = [transforms.ToTensor()])
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_dataset = ListDataset(root=file_root,
                                list_file='data/val_rewrite.txt',
                                train=False,
                                transform = [transforms.ToTensor()])
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    print('the dataset has %d images' % (len(train_dataset)))
    print('the batch_size is %d' % (batch_size))

    num_iter = 0
    vis = visdom.Visdom()
    win = vis.line(Y=np.array([0]), X = np.array([0]))

    net.train()
    for epoch in range(num_epochs):
        scheduler.step()

        print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
        print('Learning Rate for this epoch: {}'.format(learning_rate))

        total_loss = 0.
        net.train()
        for i,(images,loc_targets,conf_targets) in enumerate(train_loader):

            if use_gpu:
                images = images.cuda()
                loc_targets = loc_targets.cuda()
                conf_targets = conf_targets.cuda()

            loc_preds, conf_preds = net(images)
            loss = criterion(loc_preds,loc_targets,conf_preds,conf_targets)
            total_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i+1) % 10 == 0:
                print ('Epoch [{}/{}], Iter [{}/{}] Loss: {:.4f}, average_loss: {:.4f}'.format(
                    epoch+1, num_epochs, i+1, len(train_loader), loss.item(), total_loss / (i+1)))

                #train_loss = total_loss /(len(train_dataset) / batch_size)
                vis.line(Y=np.array([total_loss / (i+1)]), X=np.array([num_iter]),
                        win=win,
                        name='train',
                        update='append')
                num_iter += 1
        # val_loss = 0.0
        # net.eval()
        # for idx, (images, loc_targets,conf_targets) in enumerate(val_loader):
        #     with torch.no_grad():
        #         if use_gpu:
        #             images = images.cuda()
        #             loc_targets = loc_targets.cuda()
        #             conf_targets = conf_targets.cuda()
        #
        #         loc_preds, conf_preds = net(images)
        #         loss = criterion(loc_preds, loc_targets, conf_preds, conf_targets)
        #         val_loss += loss.item()
        # val_loss /= len(val_dataset)/batch_size
        # vis.line(Y=np.array([val_loss]), X=np.array([epoch]),
        #          win=win,
        #          name='val',
        #          update='append')
        # print('loss of val is {}'.format(val_loss))

        if not os.path.exists('weight/'):
            os.mkdir('weight')

        print('saving model ...')
        torch.save(net.state_dict(),'weight/faceboxes_add_norm.pt')
Ejemplo n.º 5
0
    parser.add_argument('--nms_threshold', type=float, help='learning rate', default=0.0)
    parser.add_argument('--cuda', type=bool, help='use gpu', default=True)

    input_args = parser.parse_args()
    print(input_args)
    return input_args


if __name__ == '__main__':
    args = parse_argvs()

    # train_path = os.path.expanduser('~/deeplearning/Data/car_rough_detect/car_detect_train/')
    # test_path = os.path.expanduser('~/deeplearning/Data/car_rough_detect/car_detect_test/')

    # common.mkdir_if_not_exist('./weight')
    # output_model_path = 'weight/car_rough_detect.pt'

    # batch_size = 20
    # lr = 0.001
    # img_size = 1024

    model = FaceBox()
    model_train = ModuleTrain(train_path=args.train_path, test_path=args.test_path,
                              label_name=args.label_name, model_file=args.output_model_path,
                              model=model, batch_size=args.batch_size, img_size=args.img_size,
                              lr=args.lr, use_gpu=args.cuda, nms_threshold=args.nms_threshold)

    model_train.train(200, 60)
    # model_train.test(show_info=True)

Ejemplo n.º 6
0
from networks import FaceBox
from multibox_loss import MultiBoxLoss
from dataset import ListDataset

import visdom
import numpy as np

use_gpu = torch.cuda.is_available()
file_root = os.path.expanduser('~/codedata/aflw/')

learning_rate = 0.001
num_epochs = 300
batch_size = 32

net = FaceBox()
if use_gpu:
    net.cuda()

print('load model...')
# net.load_state_dict(torch.load('weight/faceboxes.pt'))

criterion = MultiBoxLoss()

# optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0003)
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=1e-4)

train_dataset = ListDataset(root=file_root,list_file='label/box_label.txt',train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=5)
print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
Ejemplo n.º 7
0
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
from torch.autograd import Variable

from networks import FaceBox
from multibox_loss import MultiBoxLoss
from dataset import ListDataset

use_gpu = torch.cuda.is_available()

file_root = '/home/lxg/codedata/aflw/'
learning_rate = 0.001
num_epochs = 50
batch_size = 10
net = FaceBox()

net.load_state_dict(torch.load('weight/faceboxes.pt'))

print(net)
print('load pre-trined model')
print('cuda', torch.cuda.current_device(), torch.cuda.device_count())

criterion = MultiBoxLoss()
if use_gpu:
    net.cuda()

net.train()
# optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-4)
optimizer = torch.optim.Adam(net.parameters(), weight_decay=1e-4)