def __init__(self, model_format, color_space, model_type, model_path):
self.model_format = model_format
self.color_space = color_space
self.model_type = model_type
if self.model_format == "torch":
if self.model_type == 'deepEyeNet':
self.modelSession = DeepEyeModel().to(device='cpu')
# saved = torch.load('utility_functions/demo_models/demo_DEN_0_99/best_checkpoint.pth.tar', map_location='cpu')
saved = torch.load(model_path, map_location='cpu')
else:
self.modelSession = ITrackerModel(
self.color_space, self.model_type).to(device='cpu')
# saved = torch.load('utility_functions/demo_models/demo_resNet_0_78/best_checkpoint.pth.tar', map_location='cpu')
saved = torch.load(model_path, map_location='cpu')
self.modelSession.load_state_dict(saved['state_dict'])
self.modelSession.eval()
elif self.model_format == "onnx":
self.modelSession = onnxruntime.InferenceSession(model_path)
class InferenceEngine:
def __init__(self, model_format, color_space, model_type, model_path):
self.model_format = model_format
self.color_space = color_space
self.model_type = model_type
if self.model_format == "torch":
if self.model_type == 'deepEyeNet':
self.modelSession = DeepEyeModel().to(device='cpu')
# saved = torch.load('utility_functions/demo_models/demo_DEN_0_99/best_checkpoint.pth.tar', map_location='cpu')
saved = torch.load(model_path, map_location='cpu')
else:
self.modelSession = ITrackerModel(
self.color_space, self.model_type).to(device='cpu')
# saved = torch.load('utility_functions/demo_models/demo_resNet_0_78/best_checkpoint.pth.tar', map_location='cpu')
saved = torch.load(model_path, map_location='cpu')
self.modelSession.load_state_dict(saved['state_dict'])
self.modelSession.eval()
elif self.model_format == "onnx":
self.modelSession = onnxruntime.InferenceSession(model_path)
def run_inference(self, normalize_image, image_face, image_eye_left,
image_eye_right, image_face_grid):
# compute output
if self.model_format == "torch":
with torch.no_grad():
output = self.modelSession(image_face, image_eye_left,
image_eye_right, image_face_grid)
gaze_prediction_np = output.numpy()[0]
elif self.model_format == "onnx":
# compute output
output = self.modelSession.run(
None, {
"face": image_face.numpy(),
"eyesLeft": image_eye_left.numpy(),
"eyesRight": image_eye_right.numpy(),
"imFaceGrid": image_face_grid.numpy()
})
gaze_prediction_np = (output[0])[0]
else:
print('Error: Invalid mode: ', self.mode)
return gaze_prediction_np
def __init__(self, mode, color_space, model_type):
self.mode = mode
self.color_space = color_space
self.model_type = model_type
if self.mode == "torch":
if self.model_type == 'deepEyeNet':
self.modelSession = DeepEyeModel().to(device='cpu')
saved = torch.load(
'utility_functions/demo_models/demo_0.85/best_checkpoint.pth.tar',
map_location='cpu')
else:
self.modelSession = ITrackerModel(
self.color_space, self.model_type).to(device='cpu')
saved = torch.load(
'utility_functions/demo_models/MSR_0_9171/best_checkpoint.pth.tar',
map_location='cpu')
self.modelSession.load_state_dict(saved['state_dict'])
self.modelSession.eval()
elif self.mode == "onnx":
self.modelSession = onnxruntime.InferenceSession('itracker.onnx')
def main():
global args, best_prec1, weight_decay, momentum
inception_exists = os.path.isfile(get_checkpoint_path(INCEPTION_FILENAME))
model = ITrackerModel(use_pretrained=not inception_exists)
if use_cuda:
model = torch.nn.DataParallel(model)
model = try_cuda(model)
imSize=(224, 224)
cudnn.benchmark = True
epoch = 0
minibatch = 0
if doLoad:
saved = load_checkpoint()
if saved:
print('Loading checkpoint for epoch %05d with error %.5f...' % (saved['epoch'], saved['best_prec1']))
state_dict = model.state_dict()
state = saved['state_dict']
if isinstance(model, torch.nn.DataParallel):
temp_state = OrderedDict()
for k, v in state.items():
if not k.startswith('module.'):
k = 'module.' + k
temp_state[k] = v
state = temp_state
else:
temp_state = OrderedDict()
for k, v in state.items():
if k.startswith('module.'):
k = k[len('module.'):]
temp_state[k] = v
state = temp_state
if not os.path.isfile(get_checkpoint_path(INCEPTION_FILENAME)):
# Delete the connected layers if not the Inception file because
# the modified network does not have these.
if 'eyesFC.0.weight' in state:
del state['eyesFC.0.weight']
if 'module.eyesFC.0.weight' in state:
del state['module.eyesFC.0.weight']
if 'eyesFC.0.bias' in state:
del state['eyesFC.0.bias']
if 'module.eyesFC.0.bias' in state:
del state['module.eyesFC.0.bias']
state_dict.update(state)
try:
model.module.load_state_dict(state_dict)
except:
model.load_state_dict(state_dict)
epoch = saved['epoch']
best_prec1 = saved['best_prec1']
if 'minibatch' in saved:
minibatch = saved['minibatch']
else:
print('Warning: Could not read checkpoint!')
dataTrain = ITrackerData(split='train', imSize = imSize)
dataVal = ITrackerData(split='val', imSize = imSize)
# Gotta mess with the train.
train_loader = torch.utils.data.DataLoader(
dataTrain,
batch_size=batch_size, shuffle=True,
num_workers=workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
dataVal,
batch_size=batch_size, shuffle=False,
num_workers=workers, pin_memory=True)
criterion = try_cuda(nn.MSELoss())
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr,
momentum=momentum,
weight_decay=weight_decay)
# Quick test
if doTest:
validate(val_loader, model, criterion, epoch)
return
# There's a bug here which causes epoch = epoch - 1. However, the checkpoint is already
# saved with the higher number epoch, so we'll just always add 1 when saving the epoch instead.
for epoch in range(0, epoch):
adjust_learning_rate(optimizer, epoch)
for epoch in range(epoch, epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, minibatch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch)
# remember best prec@1, set minibatch to zero, and save checkpoint
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
minibatch = 0
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
def main():
global args, best_prec1, weight_decay, momentum
model = ITrackerModel()
model = torch.nn.DataParallel(model)
model.cuda()
imSize = (224, 224)
cudnn.benchmark = True
epoch = 0
if doLoad:
saved = load_checkpoint()
if saved:
print(
'Loading checkpoint for epoch %05d with loss %.5f (which is L2 = mean of squares)...'
% (saved['epoch'], saved['best_prec1']))
state = saved['state_dict']
try:
model.module.load_state_dict(state)
except:
model.load_state_dict(state)
epoch = saved['epoch']
best_prec1 = saved['best_prec1']
else:
print('Warning: Could not read checkpoint!')
dataTrain = ITrackerData(split='train', imSize=imSize)
dataVal = ITrackerData(split='val', imSize=imSize)
dataTest = ITrackerData(split='test', imSize=imSize)
train_loader = torch.utils.data.DataLoader(dataTrain,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(dataVal,
batch_size=batch_size,
shuffle=False,
num_workers=workers,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(dataTest,
batch_size=batch_size,
shuffle=False,
num_workers=workers,
pin_memory=True)
criterion = nn.MSELoss().cuda() # mean square error
optimizer = torch.optim.SGD(model.parameters(),
lr,
momentum=momentum,
weight_decay=weight_decay)
# Quick test
if doTest:
validate(test_loader, model, criterion, epoch)
return
for epoch in range(0, epoch):
adjust_learning_rate(optimizer, epoch)
for epoch in range(epoch, epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
def initialize_torch(path, model_type, device, color_space):
model = ITrackerModel(color_space, model_type).to(device=device)
saved = torch.load(path, map_location=device)
model.load_state_dict(saved['state_dict'])
model.eval()
return model
def initialize_model(args):
if not args.info:
if args.verbose:
print('')
if args.using_cuda:
print('Using cuda devices:', args.local_rank)
# print('CUDA DEVICE_COUNT {0}'.format(torch.cuda.device_count()))
print('')
# Retrieve model
if args.model_type == "deepEyeNet":
model = DeepEyeModel().to(device=args.device)
else:
model = ITrackerModel(args.color_space,
args.model_type).to(device=args.device)
# print(model)
# GPU optimizations and modes
# Enable Heuristics: cuDNN will apply heuristics before training to figure
# out the most performant algorithm for the model architecture and input.
# This is especially helpful, if input shapes don't change during training.
cudnn.benchmark = True
if args.using_cuda:
if args.mode == 'dp':
print('Using DataParallel Backend')
if not args.disable_sync:
from utility_functions.sync_batchnorm import convert_model
# Convert batchNorm layers into synchronized batchNorm
model = convert_model(model)
model = torch.nn.DataParallel(
model, device_ids=args.local_rank).to(device=args.device)
elif args.mode == 'ddp1':
# Single-Process Multiple-GPU: You'll observe all gpus running a single process (processes with same PID)
print(
'Using DistributedDataParallel Backend - Single-Process Multi-GPU'
)
torch.distributed.init_process_group(backend="nccl")
model = torch.nn.parallel.DistributedDataParallel(
model, find_unused_parameters=True)
elif args.mode == 'ddp2':
# Multi-Process Single-GPU : You'll observe multiple gpus running different processes (different PIDs)
# OMP_NUM_THREADS = nb_cpu_threads / nproc_per_node
torch.distributed.init_process_group(backend='nccl')
if not args.disable_sync:
# Convert batchNorm layers into synchronized batchNorm
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
model = torch.nn.parallel.DistributedDataParallel(
model,
find_unused_parameters=True,
device_ids=args.local_rank,
output_device=args.local_rank[0])
###### code after this place runs in their own process #####
set_print_policy(args.master, torch.distributed.get_rank())
print(
'Using DistributedDataParallel Backend - Multi-Process Single-GPU'
)
else:
print("No Parallelization")
else:
print("Cuda disabled")
else:
model = None
# use new model or load existing
val_RMSErrors, test_RMSErrors, train_RMSErrors, best_RMSErrors, best_RMSError, epoch, learning_rates = checkpoint_manager.extract_checkpoint_data(
args, model)
return val_RMSErrors, test_RMSErrors, train_RMSErrors, best_RMSErrors, best_RMSError, epoch, learning_rates, model
def main():
global args, best_prec1, weight_decay, momentum
model = ITrackerModel()
print("Total number of parameters: ",
sum(p.numel() for p in model.parameters()))
#print("Number of trainable parameters: ", sum(p.numel() for p in model.parameters() if p.requires_grad))
model = nn.DataParallel(model)
model.cuda()
imSize = (224, 224)
cudnn.benchmark = True
epoch = 0
if doLoad:
saved = load_checkpoint()
if saved:
print(
'Loading checkpoint for epoch %05d with loss %.5f (which is the mean squared error not the actual linear error)...'
% (saved['epoch'], saved['best_prec1']))
state = saved['state_dict']
try:
model.module.load_state_dict(state)
except:
model.load_state_dict(state)
epoch = saved['epoch']
best_prec1 = saved['best_prec1']
else:
print('Warning: Could not read checkpoint!')
dataTrain = ITrackerData(dataPath=args.data_path,
split='train',
imSize=imSize)
dataVal = ITrackerData(dataPath=args.data_path,
split='test',
imSize=imSize)
# data = dataVal.__getitem__(0)
# for idx in range(len(data)):
# np.squeeze(data[idx])
# model_stats = get_model_stats(model, data)
train_loader = torch.utils.data.DataLoader(dataTrain,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(dataVal,
batch_size=batch_size,
shuffle=False,
num_workers=workers,
pin_memory=True)
criterion = nn.MSELoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
lr,
momentum=momentum,
weight_decay=weight_decay)
# Quick test
if doTest:
val_mean = validate(val_loader, model, criterion, epoch)
print(val_mean)
return
for epoch in range(0, epoch):
adjust_learning_rate(optimizer, epoch)
for epoch in range(epoch, epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)