def main():
config_path = Path(args.config_path)
config = yaml.load(open(config_path))
net_config = config['Net']
data_config = config['Data']
evaluate_config = config['Evaluate']
val_dir = evaluate_config['eval_dir']
val_name = evaluate_config['eval_name']
val_type = evaluate_config['eval_type']
target_size = data_config["target_size"]
use_bined = False
num_workers = data_config['num_worker']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# pretrained_paths = ["/media/2tb/output_models/headpose_resnet/model_epoch_79_4.179329837522199.pth"] # 4.6
# pretrained_paths = ["/media/2tb/output_models/headpose_resnet/model_epoch_78_3.985460854345752.pth"] # 4.3
pretrained_paths = [
"/home/linhnv/projects/fsanet_pytorch/model_train/FSA_MSE/model_epoch_79_5.914248631846521.pth"
]
# pretrained_paths = ["/media/2tb/output_models/headpose_fsanet/model_epoch_228_5.439456623458148.pth"] # 7.8
# models_path = glob("/home/linhnv/projects/RankPose/model/headpose_resnet/*")
# models_path = [x for x in models_path if x.startswith("/home/linhnv/projects/RankPose/model/headpose_resnet/model_epoch")]
# print(models_path)
model = load_model(**net_config)
# To device
model = model.to(device)
modelname = config_path.stem
output_dir = Path('../model') / modelname
output_dir.mkdir(parents=True, exist_ok=True)
log_dir = Path('../logs') / modelname
log_dir.mkdir(parents=True, exist_ok=True)
# logger = debug_logger(log_dir)
# logger.debug(config)
# logger.info(f'Device: {device}')
valid_dataset = load_dataset(data_type=val_type,
base_dir=val_dir,
filename=val_name,
n_class=net_config['n_class'],
target_size=target_size,
debug=False)
# top_10 = len(train_dataset) // 10
# top_30 = len(train_dataset) // 3.33
# train_weights = [ 3 if idx<top_10 else 2 if idx<top_30 else 1 for idx in train_dataset.labels_sort_idx]
# train_sample = WeightedRandomSampler(train_weights, num_samples=len(train_dataset), replacement=True)
# train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sample, num_workers=num_workers,
# pin_memory=True, drop_last=True)
valid_loader = DataLoader(valid_dataset,
batch_size=32,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
if torch.cuda.is_available():
model = nn.DataParallel(model)
for pretrained_path in pretrained_paths:
print(f"[INFO] Pretrained path: {pretrained_path}")
# logger.info(f'Load pretrained from {pretrained_path}')
param = torch.load(pretrained_path, map_location='cpu')
model.load_state_dict(param)
model.to(device)
del param
valid_diffs = []
model.eval()
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
images, labels = batched
images, labels = images.to(device), labels.to(device)
preds = model(images)
# loss = loss_fn([preds], [labels])
diff, dif_yaw, dif_pitch, dif_roll = calculate_diff(
preds, labels)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_diffs.append(diff)
valid_diff = np.mean(valid_diffs)
print(f'valid diff: {valid_diff}')
def main():
config_path = Path(args.config_path)
config = yaml.load(open(config_path))
net_config = config['Net']
data_config = config['Data']
train_config = config['Train']
# Config for data:
train_dir = data_config["train_dir"]
train_name = data_config["train_name"]
train_type = data_config["train_type"]
val_dir = data_config["val_dir"]
val_name = data_config["val_name"]
val_type = data_config["val_type"]
target_size = data_config["target_size"]
num_workers = data_config["num_worker"]
# Config for train:
num_epoch = train_config["num_epoch"]
batch_size = train_config["batch_size"]
val_every = train_config["val_every"]
resume = train_config["resume"]
pretrained_path = train_config["pretrained_path"]
saved_dir = train_config["saved_dir"]
epoch_start = 0
loss_type = train_config["loss_type"]
optimizer_config = train_config["optimizer"]
del data_config
del train_config
model = load_model(**net_config)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"[INFO] Device: {device}")
# To device
model = model.to(device)
# if torch.cuda.is_available():
# model.cuda()
modelname = config_path.stem
output_dir = Path(saved_dir) / "models" / modelname
output_dir.mkdir(parents=True, exist_ok=True)
log_dir = Path(saved_dir) / "logs" / modelname
log_dir.mkdir(parents=True, exist_ok=True)
# logger = debug_logger(log_dir)
# logger.debug(config)
# logger.info(f'Device: {device}')
# logger.info(f'Max Epoch: {max_epoch}')
loss_fn = Criterion(loss_type=loss_type).to(device)
params = model.parameters()
optimizer, scheduler = create_optimizer(params, **optimizer_config)
# Dataset
affine_augmenter = albu.Compose([
albu.GaussNoise(var_limit=(0, 25), p=.2),
albu.GaussianBlur(3, p=0.2),
albu.JpegCompression(50, 100, p=0.2)
])
image_augmenter = albu.Compose([
albu.OneOf([
albu.RandomBrightnessContrast(0.25, 0.25),
albu.CLAHE(clip_limit=2),
albu.RandomGamma(),
],
p=0.5),
albu.HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=20,
p=0.2),
albu.RGBShift(p=0.2),
albu.RandomSizedCrop(min_max_height=[45, 64],
height=64,
width=64,
p=0.5),
])
train_dataset = load_dataset(data_type=train_type,
base_dir=train_dir,
filename=train_name,
n_class=net_config['n_class'],
target_size=target_size,
affine_augmenter=affine_augmenter,
image_augmenter=image_augmenter,
debug=False)
val_dataset = load_dataset(data_type=val_type,
base_dir=val_dir,
filename=val_name,
n_class=net_config['n_class'],
target_size=target_size,
debug=False)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True,
pin_memory=True,
drop_last=True)
valid_loader = DataLoader(val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
if torch.cuda.is_available():
model = nn.DataParallel(model)
if resume:
checkpoint = torch.load(pretrained_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch_start = checkpoint['epoch'] + 1
loss_history = checkpoint['loss_history']
else:
loss_history = []
model.train()
for i_epoch in range(epoch_start, num_epoch):
print(f"Epoch: {i_epoch}")
print(f'Learning rate: {optimizer.param_groups[0]["lr"]}')
train_losses = []
train_diffs = []
model.train()
with tqdm(train_loader) as _tqdm:
for batched in _tqdm:
optimizer.zero_grad()
if loss_type == "RANK":
img1, img2, lbl1, lbl2, labels = batched
img1, img2, lbl1, lbl2, labels = img1.to(device), img2.to(
device), lbl1.to(device), lbl2.to(device), labels.to(
device)
preds1 = model(img1)
preds2 = model(img2)
preds1 = preds1.to(device)
preds2 = preds2.to(device)
loss = loss_fn([preds1, preds2], [lbl1, lbl2, labels])
diff = calculate_diff(preds1, lbl1)
diff += calculate_diff(preds2, lbl2)
diff /= 2
_tqdm.set_postfix(
OrderedDict(loss=f'{loss.item():.3f}',
mae=f'{diff:.1f}'))
train_losses.append(loss.item())
history_ploter(train_losses, log_dir.joinpath('loss.png'))
train_diffs.append(diff)
loss.backward()
optimizer.step()
elif loss_type == "MSE" or loss_type == "wrapped":
img1, lbl1, _, _, _ = batched
img1, lbl1 = img1.to(device), lbl1.to(device)
if net_config["net_type"] == "Perceiver":
img1 = img1.permute(0, 2, 3, 1)
preds1 = model(img1)
loss = loss_fn([preds1, []], [lbl1, []])
diff = calculate_diff(preds1, lbl1)
_tqdm.set_postfix(
OrderedDict(loss=f'{loss.item():.3f}',
mae=f'{diff:.1f}'))
train_losses.append(loss.item())
history_ploter(train_losses, log_dir.joinpath('loss.png'))
train_diffs.append(diff)
loss.backward()
optimizer.step()
train_loss = np.mean(train_losses)
train_diff = np.nanmean(train_diffs)
print(f'[INFO] train loss: {train_loss}')
print(f'[INFO] train diff: {train_diff}')
scheduler.step()
if (i_epoch + 1) % val_every == 0:
valid_losses = []
valid_diffs = []
model.eval()
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
images, labels, _, _, _ = batched
if net_config["net_type"] == "Perceiver":
images = images.permute(0, 2, 3, 1)
images, labels = images.to(device), labels.to(device)
preds = model(images)
# loss = loss_fn([preds], [labels])
diff = calculate_diff(preds, labels)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_diffs.append(diff)
valid_diff = np.mean(valid_diffs)
loss_history.append([train_diff, valid_diff])
history_ploter(loss_history, log_dir.joinpath('diff.png'))
print(f'[INFO] valid diff: {valid_diff}')
torch.save(
model.state_dict(),
output_dir.joinpath(f'model_epoch_{i_epoch}_{valid_diff}.pth'))
torch.save(
{
'epoch': i_epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss_history': loss_history,
},
output_dir.joinpath(
f'checkpoint_epoch_{i_epoch}_{valid_diff}.pth'))
else:
valid_diff = None
def main():
config_path = Path(args.config_path)
config = yaml.load(open(config_path))
net_config = config['Net']
data_config = config['Data']
train_config = config['Train']
loss_config = config['Loss']
opt_config = config['Optimizer']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
n_class = net_config['n_class']
max_epoch = train_config['max_epoch']
batch_size = train_config['batch_size']
num_workers = train_config['num_workers']
test_every = train_config['test_every']
resume = train_config['resume']
pretrained_path = train_config['pretrained_path']
use_rank = train_config['use_rank']
use_bined = train_config['use_bined']
del train_config['use_rank']
del train_config['use_bined']
train_dir = data_config['train_dir']
val_dir = data_config['val_dir']
train_name = data_config['train_name']
val_name = data_config['val_name']
train_type = data_config['train_type']
val_type = data_config['val_type']
del data_config['train_dir']
del data_config['val_dir']
del data_config['train_name']
del data_config['val_name']
del data_config['train_type']
del data_config['val_type']
model = load_model(**net_config)
# To device
model = model.to(device)
modelname = config_path.stem
output_dir = Path('../model') / modelname
output_dir.mkdir(exist_ok=True)
log_dir = Path('../logs') / modelname
log_dir.mkdir(exist_ok=True)
logger = debug_logger(log_dir)
logger.debug(config)
logger.info(f'Device: {device}')
logger.info(f'Max Epoch: {max_epoch}')
loss_fn = Criterion(**loss_config).to(device)
params = model.parameters()
optimizer, scheduler = create_optimizer(params, **opt_config)
# history
if resume:
with open(log_dir.joinpath('history.pkl'), 'rb') as f:
history_dict = pickle.load(f)
best_metrics = history_dict['best_metrics']
loss_history = history_dict['loss']
diff_history = history_dict['diff']
# start_epoch = len(diff_history)
start_epoch = 47
for _ in range(start_epoch):
scheduler.step()
else:
start_epoch = 0
best_metrics = float('inf')
loss_history = []
diff_history = []
# Dataset
affine_augmenter = albu.Compose([albu.GaussNoise(var_limit=(0,25),p=.2),
albu.GaussianBlur(3, p=0.2),
albu.JpegCompression(50, 100, p=0.2)])
image_augmenter = albu.Compose([
albu.OneOf([
albu.RandomBrightnessContrast(0.25,0.25),
albu.CLAHE(clip_limit=2),
albu.RandomGamma(),
], p=0.5),
albu.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20,p=0.2),
albu.RGBShift(p=0.2),
])
# image_augmenter = None
train_dataset = laod_dataset(data_type=train_type, affine_augmenter=affine_augmenter, image_augmenter=image_augmenter,
base_dir=train_dir, filename=train_name, use_bined=use_bined, n_class=n_class, **data_config)
valid_dataset = laod_dataset(data_type=val_type, split='valid', base_dir=val_dir, filename=val_name,
use_bined=use_bined, n_class=n_class, **data_config)
# top_10 = len(train_dataset) // 10
# top_30 = len(train_dataset) // 3.33
# train_weights = [ 3 if idx<top_10 else 2 if idx<top_30 else 1 for idx in train_dataset.labels_sort_idx]
# train_sample = WeightedRandomSampler(train_weights, num_samples=len(train_dataset), replacement=True)
# train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sample, num_workers=num_workers,
# pin_memory=True, drop_last=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, drop_last=True)
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False, num_workers=num_workers, pin_memory=True)
if torch.cuda.is_available():
model = nn.DataParallel(model)
# Pretrained model
if pretrained_path:
logger.info(f'Load pretrained from {pretrained_path}')
param = torch.load(pretrained_path, map_location='cpu')
if "state_dict" in param:
model.load_state_dict(param['state_dict'], strict=False)
else:
model.load_state_dict(param)
del param
# Restore model
if resume:
print("[INFO] resume training.")
model_path = output_dir.joinpath(f'model_epoch_{start_epoch-1}.pth')
logger.info(f'Resume from {model_path}')
param = torch.load(model_path, map_location='cpu')
model.load_state_dict(param)
del param
opt_path = output_dir.joinpath(f'opt_epoch_{start_epoch-1}.pth')
param = torch.load(opt_path)
optimizer.load_state_dict(param)
del param
file_train_log = open("file_train_log.txt", "a")
file_val_log = open("file_val_log.txt", "a")
# Train
for i_epoch in range(start_epoch, max_epoch):
logger.info(f'Epoch: {i_epoch}')
logger.info(f'Learning rate: {optimizer.param_groups[0]["lr"]}')
train_losses = []
train_diffs = []
model.train()
with tqdm(train_loader) as _tqdm:
for batched in _tqdm:
optimizer.zero_grad()
if use_rank:
if use_bined:
img1, img2, lbl1, lbl2, labels, yaw_lbl1, pitch_lbl1, roll_lbl1, yaw_lbl2, pitch_lbl2, roll_lbl2 = batched
img1, img2, lbl1, lbl2, labels = img1.to(device),img2.to(device),lbl1.to(device),lbl2.to(device),labels.to(device)
yaw_lbl1, pitch_lbl1, roll_lbl1 = yaw_lbl1.to(device), pitch_lbl1.to(device), roll_lbl1.to(device)
yaw_lbl2, pitch_lbl2, roll_lbl2 = yaw_lbl2.to(device), pitch_lbl2.to(device), roll_lbl2.to(device)
preds1, y_pres1, p_pres1, r_pres1 = model(img1, True)
preds2, y_pres2, p_pres2, r_pres2 = model(img2, True)
pre_list = [preds1,preds2,y_pres1,p_pres1,r_pres1,y_pres2,p_pres2,r_pres2]
lbl_list = [lbl1,lbl2,yaw_lbl1,pitch_lbl1,roll_lbl1,yaw_lbl2,pitch_lbl2,roll_lbl2,labels]
loss = loss_fn(pre_list, lbl_list, use_bined=True)
else:
img1, img2, lbl1, lbl2, labels = batched
img1, img2, lbl1, lbl2, labels = img1.to(device),img2.to(device),lbl1.to(device),lbl2.to(device),labels.to(device)
preds1 = model(img1, False)
preds2 = model(img2, False)
loss = loss_fn([preds1,preds2], [lbl1,lbl2,labels], use_bined=False)
# print(f"Preds1: {preds1}")
# print(f"Preds2: {preds2}")
# print(f"lib1: {lbl1}")
# print(f"lib2: {lbl2}")
diff = calculate_diff(preds1, lbl1)
diff += calculate_diff(preds2, lbl2)
diff /= 2
# print(f"Diff: {diff}")
elif use_bined:
images, labels, yaw_labels, pitch_labels, roll_labels = batched
images, labels = images.to(device), labels.to(device)
yaw_labels, pitch_labels, roll_labels = yaw_labels.to(device), pitch_labels.to(device), roll_labels.to(device)
preds, y_pres, p_pres, r_pres = model(images, use_bined)
loss = loss_fn([preds, y_pres, p_pres, r_pres], [labels, yaw_labels, pitch_labels, roll_labels], use_bined)
diff = calculate_diff(preds, labels)
else:
images, labels = batched
images, labels = images.to(device), labels.to(device)
preds = model(images, use_bined)
loss = loss_fn([preds], [labels])
diff = calculate_diff(preds, labels, mean=True)
_tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', mae=f'{diff:.1f}'))
train_losses.append(loss.item())
train_diffs.append(diff)
loss.backward()
optimizer.step()
scheduler.step()
train_loss = np.mean(train_losses)
train_diff = np.nanmean(train_diffs)
logger.info(f'train loss: {train_loss}')
logger.info(f'train diff: {train_diff}')
file_train_log.write(f"{train_loss},{train_diff}")
# torch.save(model.module.state_dict(), output_dir.joinpath(f'model_tmp_epoch_{i_epoch}.pth'))
# torch.save(optimizer.state_dict(), output_dir.joinpath(f'opt_tmp_{i_epoch}.pth'))
if (i_epoch + 1) % test_every == 0:
valid_losses = []
valid_diffs = []
model.eval()
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
if use_bined:
images, labels, yaw_labels, pitch_labels, roll_labels = batched
images, labels = images.to(device), labels.to(device)
# yaw_labels, pitch_labels, roll_labels = yaw_labels.to(device), pitch_labels.to(device), roll_labels.to(device)
preds, y_pres, p_pres, r_pres = model(images, use_bined)
# loss = loss_fn([preds, y_pres, p_pres, r_pres], [labels, yaw_labels, pitch_labels, roll_labels])
diff = calculate_diff(preds, labels)
else:
images, labels = batched
images, labels = images.to(device), labels.to(device)
preds = model(images, use_bined)
# loss = loss_fn([preds], [labels])
diff = calculate_diff(preds, labels)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_losses.append(0)
valid_diffs.append(diff)
valid_loss = np.mean(valid_losses)
valid_diff = np.mean(valid_diffs)
logger.info(f'valid seg loss: {valid_loss}')
logger.info(f'valid diff: {valid_diff}')
file_val_log.write(f"{valid_loss},{valid_diff}")
if best_metrics >= valid_diff:
best_metrics = valid_diff
logger.info('Best Model!\n')
torch.save(model.state_dict(), output_dir.joinpath(f'model_epoch_{i_epoch}_{valid_diff}.pth'))
torch.save(optimizer.state_dict(), output_dir.joinpath(f'opt_epoch_{i_epoch}_{valid_diff}.pth'))
torch.save(model.state_dict(), output_dir.joinpath(f'model_epoch_{i_epoch}_{valid_diff}.pth'))
torch.save(optimizer.state_dict(), output_dir.joinpath(f'opt_epoch_{i_epoch}_{valid_diff}.pth'))
else:
valid_loss = None
valid_diff = None
loss_history.append([train_loss, valid_loss])
diff_history.append([train_diff, valid_diff])
history_ploter(loss_history, log_dir.joinpath('loss.png'))
history_ploter(diff_history, log_dir.joinpath('diff.png'))
history_dict = {'loss': loss_history,
'diff': diff_history,
'best_metrics': best_metrics}
with open(log_dir.joinpath('history.pkl'), 'wb') as f:
pickle.dump(history_dict, f)
file_train_log.close()
file_val_log.close()
def main(pretrained_path, config_path, calculate_diff, val_dir, val_name, val_type):
"""
input:
pretrained_path: path to checkpoints
config_path: path to config
calculate_diff: Metric function
val_dir, val_name, val_type: config for data
output:
None
"""
# config_path = Path(args.config_path)
config = yaml.load(open(config_path))
net_config = config['Net']
data_config = config['Data']
train_config = config['Train']
target_size = data_config["target_size"]
num_workers = data_config["num_worker"]
batch_size = train_config["batch_size"]
del data_config
del train_config
model = load_model(**net_config)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"[INFO] Device: {device}")
# To device
model = model.to(device)
# if torch.cuda.is_available():
# model.cuda()
# logger = debug_logger(log_dir)
# logger.debug(config)
# logger.info(f'Device: {device}')
# logger.info(f'Max Epoch: {max_epoch}')
val_dataset = load_dataset(data_type=val_type, base_dir=val_dir, filename=val_name, n_class=net_config['n_class'], target_size=target_size, debug=False )
valid_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
if torch.cuda.is_available():
model = nn.DataParallel(model)
checkpoint = torch.load(pretrained_path)
model.load_state_dict(checkpoint['model_state_dict'])
valid_diffs = []
model.eval()
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
# images, labels = batched
images, labels, _ , _, _ = batched
images, labels = images.to(device), labels.to(device)
preds = model(images)
# loss = loss_fn([preds], [labels])
dif_yaw, dif_pitch, dif_roll, diff = calculate_diff(preds, labels, True)
# print(diff, dif_yaw, dif_pitch, dif_roll)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_diffs.append([dif_yaw, dif_pitch, dif_roll])
valid_diffs = np.array(valid_diffs)
print(valid_diffs.shape)
valid_diff = np.mean(valid_diffs, 0)
print(f'[INFO] valid diff: {valid_diff}')
print((valid_diff[0] + valid_diff[1] + valid_diff[2]) / 3)
def evaluate_fusion(pretrained_uni_path, config_uni_path, pretrained_var_path, config_var_path, pretrained_wei_path, config_wei_path, calculate_diff, val_dir, val_name, val_type):
"""
input:
pretrained_uni_path, pretrained_var_path, pretrained_wei_path: pretrained paths of 3 model types.
config_uni_path, config_var_path, config_wei_path: config paths of 3 model types.
calculate_diff: Metric function
val_dir, val_name, val_type: config for data
output:
None
"""
# config_path = Path(args.config_path)
uni_model, target_size, num_workers, batch_size, n_class = create_model(config_uni_path)
var_model, target_size, num_workers, batch_size, n_class = create_model(config_var_path)
wei_model, target_size, num_workers, batch_size, n_class = create_model(config_wei_path)
val_dataset = load_dataset(data_type=val_type, base_dir=val_dir, filename=val_name, n_class=n_class, target_size=target_size, debug=False )
valid_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
uni_model = load_pretrain_model(uni_model, pretrained_uni_path)
var_model = load_pretrain_model(var_model, pretrained_var_path)
wei_model = load_pretrain_model(wei_model, pretrained_wei_path)
valid_diffs = []
uni_model.eval()
var_model.eval()
wei_model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
# images, labels = batched
images, labels, _ , _, _ = batched
images, labels = images.to(device), labels.to(device)
uni_preds = uni_model(images)
var_preds = var_model(images)
wei_preds = wei_model(images)
preds = uni_preds.add(var_preds).add(wei_preds)/3
# loss = loss_fn([preds], [labels])
dif_yaw, dif_pitch, dif_roll, diff = calculate_diff(preds, labels, True)
# print(diff, dif_yaw, dif_pitch, dif_roll)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_diffs.append([dif_yaw, dif_pitch, dif_roll, diff])
valid_diffs = np.array(valid_diffs)
print(valid_diffs.shape)
valid_diff = np.mean(valid_diffs, 0)
print(f'[INFO] valid diff: {valid_diff}')
print((valid_diff[0] + valid_diff[1] + valid_diff[2]) / 3)
def main():
config_path = Path(args.config_path)
config = yaml.load(open(config_path))
net_config = config['Net']
# loss_config = config['Loss']
# opt_config = config['Optimizer']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
n_class = net_config['n_class']
val_dir = '../data'
val_name = 'biwi_dataset_list.txt'
val_type = 'BIWI'
use_bined = False
num_workers = 4
pretrained_path = [
"/home/linhnv/projects/RankPose/model/headpose_resnet/model_epoch_77_8.775894704750192.pth"
]
# models_path = glob("/home/linhnv/projects/RankPose/model/headpose_resnet/*")
# models_path = [x for x in models_path if x.startswith("/home/linhnv/projects/RankPose/model/headpose_resnet/model_epoch")]
# print(models_path)
for pretrained_path in pretrained_path:
print(f"[INFO] Pretrained path: {pretrained_path}")
model = load_model(**net_config)
# To device
model = model.to(device)
modelname = config_path.stem
output_dir = Path('../model') / modelname
output_dir.mkdir(exist_ok=True)
log_dir = Path('../logs') / modelname
log_dir.mkdir(exist_ok=True)
# logger = debug_logger(log_dir)
# logger.debug(config)
# logger.info(f'Device: {device}')
params = model.parameters()
valid_dataset = laod_dataset(data_type=val_type,
split='valid',
base_dir=val_dir,
filename=val_name,
use_bined=False,
n_class=n_class)
# top_10 = len(train_dataset) // 10
# top_30 = len(train_dataset) // 3.33
# train_weights = [ 3 if idx<top_10 else 2 if idx<top_30 else 1 for idx in train_dataset.labels_sort_idx]
# train_sample = WeightedRandomSampler(train_weights, num_samples=len(train_dataset), replacement=True)
# train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sample, num_workers=num_workers,
# pin_memory=True, drop_last=True)
valid_loader = DataLoader(valid_dataset,
batch_size=32,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
if torch.cuda.is_available():
model = nn.DataParallel(model)
# logger.info(f'Load pretrained from {pretrained_path}')
param = torch.load(pretrained_path, map_location='cpu')
if "state_dict" in param:
model.load_state_dict(param['state_dict'], strict=False)
else:
model.load_state_dict(param)
del param
valid_losses = []
valid_diffs = []
model.eval()
with torch.no_grad():
with tqdm(valid_loader) as _tqdm:
for batched in _tqdm:
if use_bined:
images, labels, yaw_labels, pitch_labels, roll_labels = batched
images, labels = images.to(device), labels.to(device)
# yaw_labels, pitch_labels, roll_labels = yaw_labels.to(device), pitch_labels.to(device), roll_labels.to(device)
preds, y_pres, p_pres, r_pres = model(
images, use_bined)
# loss = loss_fn([preds, y_pres, p_pres, r_pres], [labels, yaw_labels, pitch_labels, roll_labels])
diff = calculate_diff(preds, labels)
else:
images, labels = batched
images, labels = images.to(device), labels.to(device)
preds = model(images, use_bined)
# loss = loss_fn([preds], [labels])
diff = calculate_diff(preds, labels)
_tqdm.set_postfix(OrderedDict(mae=f'{diff:.2f}'))
# _tqdm.set_postfix(OrderedDict(loss=f'{loss.item():.3f}', d_y=f'{np.mean(diff[:,0]):.1f}', d_p=f'{np.mean(diff[:,1]):.1f}', d_r=f'{np.mean(diff[:,2]):.1f}'))
valid_losses.append(0)
valid_diffs.append(diff)
valid_loss = np.mean(valid_losses)
valid_diff = np.mean(valid_diffs)
print(f'valid diff: {valid_diff}')