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")
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)
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')
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)