refactor evaluation and training
This commit is contained in:
parent
2c97f20fe9
commit
89d860df2a
@ -129,7 +129,7 @@ by the image name to easily access ground truth files for evaluation and predict
|
||||
|
||||
# Training
|
||||
Provide the json file containing the preprocess joints as argument.
|
||||
As simple as `python3 -m monstereo.run train --joints <json file path>`
|
||||
As simple as `python3 -m monstereo.run train --joints <json file path> `
|
||||
All the hyperparameters options can be checked at `python3 -m monstereo.run train --help`.
|
||||
|
||||
# Evaluation (KITTI Dataset)
|
||||
|
||||
@ -73,3 +73,13 @@ python -m monstereo.run eval --activity \
|
||||
--model <MonoLoco++ model path> --dir_ann <pifpaf annotations directory>
|
||||
```
|
||||
|
||||
## Training
|
||||
We train on KITTI or nuScenes dataset specifying the path of the input joints.
|
||||
|
||||
Our results are obtained with:
|
||||
|
||||
`python -m monstereo.run train --lr 0.001 --joints data/arrays/joints-kitti-201202-1743.json --save --monocular`
|
||||
|
||||
For a more extensive list of available parameters, run:
|
||||
|
||||
`python -m monstereo.run train --help`
|
||||
@ -25,7 +25,7 @@ class EvalKitti:
|
||||
'27', '29', '31', '49')
|
||||
ALP_THRESHOLDS = ('<0.5m', '<1m', '<2m')
|
||||
OUR_METHODS = ['geometric', 'monoloco', 'monoloco_pp', 'pose', 'reid', 'monstereo']
|
||||
METHODS_MONO = ['m3d', 'monopsr']
|
||||
METHODS_MONO = ['m3d', 'monopsr', 'monodis', 'smoke']
|
||||
METHODS_STEREO = ['3dop', 'psf', 'pseudo-lidar', 'e2e', 'oc-stereo']
|
||||
BASELINES = ['task_error', 'pixel_error']
|
||||
HEADERS = ('method', '<0.5', '<1m', '<2m', 'easy', 'moderate', 'hard', 'all')
|
||||
@ -56,6 +56,8 @@ class EvalKitti:
|
||||
self.dic_thresh_conf['monopsr'] += 0.3
|
||||
self.dic_thresh_conf['e2e-pl'] = -100 # They don't have enough detections
|
||||
self.dic_thresh_conf['oc-stereo'] = -100
|
||||
self.dic_thresh_conf['smoke'] = -100
|
||||
self.dic_thresh_conf['monodis'] = -100
|
||||
|
||||
# Extract validation images for evaluation
|
||||
names_gt = tuple(os.listdir(self.dir_gt))
|
||||
|
||||
@ -1,9 +1,10 @@
|
||||
|
||||
#pylint: disable=too-many-branches
|
||||
# pylint: disable=too-many-branches
|
||||
|
||||
"""
|
||||
Run MonoLoco/MonStereo and converts annotations into KITTI format
|
||||
"""
|
||||
from typing import Dict, List
|
||||
|
||||
import os
|
||||
import math
|
||||
@ -22,42 +23,35 @@ from .reid_baseline import get_reid_features, ReID
|
||||
|
||||
class GenerateKitti:
|
||||
|
||||
METHODS = ['monstereo', 'monoloco_pp', 'monoloco', 'geometric']
|
||||
dir_gt = os.path.join('data', 'kitti', 'gt')
|
||||
dir_gt_new = os.path.join('data', 'kitti', 'gt_new')
|
||||
dir_kk = os.path.join('data', 'kitti', 'calib')
|
||||
dir_byc = '/data/lorenzo-data/kitti/object_detection/left'
|
||||
monoloco_checkpoint = 'data/models/monoloco-190717-0952.pkl'
|
||||
baselines = {'mono': [], 'stereo': []}
|
||||
|
||||
def __init__(self, model, dir_ann, p_dropout=0.2, n_dropout=0, hidden_size=1024):
|
||||
def __init__(self, args):
|
||||
|
||||
self.dir_ann = dir_ann
|
||||
assert os.listdir(self.dir_ann), "Annotation directory is empty"
|
||||
# Load Network
|
||||
self.net = args.net
|
||||
assert args.net in ('monstereo', 'monoloco_pp'), "net not recognized"
|
||||
|
||||
# Load monoloco
|
||||
use_cuda = torch.cuda.is_available()
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
|
||||
if 'monstereo' in self.METHODS:
|
||||
self.monstereo = Loco(model=model, net='monstereo', device=device, n_dropout=n_dropout, p_dropout=p_dropout,
|
||||
linear_size=hidden_size)
|
||||
# model_mono_pp = 'data/models/monoloco-191122-1122.pkl' # KITTI_p
|
||||
# model_mono_pp = 'data/models/monoloco-191018-1459.pkl' # nuScenes_p
|
||||
# model_mono_pp = 'data/models/stereoloco-200604-0949.pkl' # KITTI_pp
|
||||
model_mono_pp = 'data/models/monstereo-201202-1745.pkl'
|
||||
# model_mono_pp = 'data/models/stereoloco-200608-1550.pkl' # nuScenes_pp
|
||||
|
||||
if 'monoloco_pp' in self.METHODS:
|
||||
self.monoloco_pp = Loco(model=model_mono_pp, net='monoloco_pp', device=device, n_dropout=n_dropout,
|
||||
p_dropout=p_dropout)
|
||||
|
||||
if 'monoloco' in self.METHODS:
|
||||
model_mono = 'data/models/monoloco-190717-0952.pkl' # KITTI
|
||||
# model_mono = 'data/models/monoloco-190719-0923.pkl' # NuScenes
|
||||
self.monoloco = Loco(model=model_mono, net='monoloco', device=device, n_dropout=n_dropout,
|
||||
p_dropout=p_dropout, linear_size=256)
|
||||
self.model = Loco(
|
||||
model=args.model,
|
||||
net=args.net,
|
||||
device=device,
|
||||
n_dropout=args.n_dropout,
|
||||
p_dropout=args.dropout,
|
||||
linear_size=args.hidden_size
|
||||
)
|
||||
|
||||
# Extract list of pifpaf files in validation images
|
||||
self.dir_gt = os.path.join('data', 'kitti', 'gt')
|
||||
self.dir_gt_new = os.path.join('data', 'kitti', 'gt_new')
|
||||
self.set_basename = factory_basename(dir_ann, self.dir_gt)
|
||||
self.dir_kk = os.path.join('data', 'kitti', 'calib')
|
||||
self.dir_byc = '/data/lorenzo-data/kitti/object_detection/left'
|
||||
self.dir_ann = args.dir_ann
|
||||
self.generate_official = args.generate_official
|
||||
assert os.listdir(self.dir_ann), "Annotation directory is empty"
|
||||
self.set_basename = factory_basename(args.dir_ann, self.dir_gt)
|
||||
|
||||
# For quick testing
|
||||
# ------------------------------------------------------------------------------------------------------------
|
||||
@ -65,33 +59,48 @@ class GenerateKitti:
|
||||
# self.set_basename = ('002282',)
|
||||
# ------------------------------------------------------------------------------------------------------------
|
||||
|
||||
# Calculate stereo baselines
|
||||
# self.baselines = ['pose', 'reid']
|
||||
self.baselines = []
|
||||
self.cnt_disparity = defaultdict(int)
|
||||
self.cnt_no_stereo = 0
|
||||
self.dir_images = os.path.join('data', 'kitti', 'images')
|
||||
self.dir_images_r = os.path.join('data', 'kitti', 'images_r')
|
||||
# ReID Baseline
|
||||
if 'reid' in self.baselines:
|
||||
weights_path = 'data/models/reid_model_market.pkl'
|
||||
self.reid_net = ReID(weights_path=weights_path, device=device, num_classes=751, height=256, width=128)
|
||||
# Add monocular and stereo baselines (they require monoloco as backbone)
|
||||
if args.baselines:
|
||||
|
||||
# Load MonoLoco
|
||||
self.baselines['mono'] = ['monoloco', 'geometric']
|
||||
self.monoloco = Loco(
|
||||
model=self.monoloco_checkpoint,
|
||||
net='monoloco',
|
||||
device=device,
|
||||
n_dropout=args.n_dropout,
|
||||
p_dropout=args.dropout,
|
||||
linear_size=256
|
||||
)
|
||||
# Stereo baselines
|
||||
if args.net == 'monstereo':
|
||||
self.baselines['stereo'] = ['pose', 'reid']
|
||||
self.cnt_disparity = defaultdict(int)
|
||||
self.cnt_no_stereo = 0
|
||||
self.dir_images = os.path.join('data', 'kitti', 'images')
|
||||
self.dir_images_r = os.path.join('data', 'kitti', 'images_r')
|
||||
|
||||
# ReID Baseline
|
||||
weights_path = 'data/models/reid_model_market.pkl'
|
||||
self.reid_net = ReID(weights_path=weights_path, device=device, num_classes=751, height=256, width=128)
|
||||
|
||||
def run(self):
|
||||
"""Run Monoloco and save txt files for KITTI evaluation"""
|
||||
|
||||
cnt_ann = cnt_file = cnt_no_file = 0
|
||||
dir_out = {key: os.path.join('data', 'kitti', key) for key in self.METHODS}
|
||||
print("\n")
|
||||
for key in self.METHODS:
|
||||
make_new_directory(dir_out[key])
|
||||
|
||||
for key in self.baselines:
|
||||
dir_out[key] = os.path.join('data', 'kitti', key)
|
||||
make_new_directory(dir_out[key])
|
||||
print("Created empty output directory for {}".format(key))
|
||||
# Prepare empty folder
|
||||
di = os.path.join('data', 'kitti', self.net)
|
||||
make_new_directory(di)
|
||||
dir_out = {self.net: di}
|
||||
|
||||
# Run monoloco over the list of images
|
||||
for mode, names in self.baselines.items():
|
||||
for name in names:
|
||||
di = os.path.join('data', 'kitti', name)
|
||||
make_new_directory(di)
|
||||
dir_out[name] = di
|
||||
|
||||
# Run the model
|
||||
for basename in self.set_basename:
|
||||
path_calib = os.path.join(self.dir_kk, basename + '.txt')
|
||||
annotations, kk, tt = factory_file(path_calib, self.dir_ann, basename)
|
||||
@ -101,58 +110,58 @@ class GenerateKitti:
|
||||
annotations_r, _, _ = factory_file(path_calib, self.dir_ann, basename, mode='right')
|
||||
_, keypoints_r = preprocess_pifpaf(annotations_r, im_size=(1242, 374))
|
||||
|
||||
if self.net == 'monstereo':
|
||||
dic_out = self.model.forward(keypoints, kk, keypoints_r=keypoints_r)
|
||||
elif self.net == 'monoloco_pp':
|
||||
dic_out = self.model.forward(keypoints, kk)
|
||||
|
||||
all_outputs = {self.net: [dic_out['xyzd'], dic_out['bi'], dic_out['epi'],
|
||||
dic_out['yaw'], dic_out['h'], dic_out['w'], dic_out['l']]}
|
||||
zzs = [float(el[2]) for el in dic_out['xyzd']]
|
||||
|
||||
# Save txt files
|
||||
params = [kk, tt]
|
||||
path_txt = os.path.join(dir_out[self.net], basename + '.txt')
|
||||
save_txts(path_txt, boxes, all_outputs[self.net], params, mode=self.net, cat=cat)
|
||||
cnt_ann += len(boxes)
|
||||
cnt_file += 1
|
||||
all_inputs, all_outputs = {}, {}
|
||||
|
||||
# STEREOLOCO
|
||||
dic_out = self.monstereo.forward(keypoints, kk, keypoints_r=keypoints_r)
|
||||
all_outputs['monstereo'] = [dic_out['xyzd'], dic_out['bi'], dic_out['epi'],
|
||||
dic_out['yaw'], dic_out['h'], dic_out['w'], dic_out['l']]
|
||||
|
||||
# MONOLOCO++
|
||||
if 'monoloco_pp' in self.METHODS:
|
||||
dic_out = self.monoloco_pp.forward(keypoints, kk)
|
||||
all_outputs['monoloco_pp'] = [dic_out['xyzd'], dic_out['bi'], dic_out['epi'],
|
||||
dic_out['yaw'], dic_out['h'], dic_out['w'], dic_out['l']]
|
||||
zzs = [float(el[2]) for el in dic_out['xyzd']]
|
||||
|
||||
# MONOLOCO
|
||||
if 'monoloco' in self.METHODS:
|
||||
# MONO (+ STEREO BASELINES)
|
||||
if self.baselines['mono']:
|
||||
# MONOLOCO
|
||||
dic_out = self.monoloco.forward(keypoints, kk)
|
||||
zzs_geom, xy_centers = geometric_coordinates(keypoints, kk, average_y=0.48)
|
||||
all_outputs['monoloco'] = [dic_out['d'], dic_out['bi'], dic_out['epi']] + [zzs_geom, xy_centers]
|
||||
all_outputs['geometric'] = all_outputs['monoloco']
|
||||
|
||||
params = [kk, tt]
|
||||
# monocular baselines
|
||||
for key in self.baselines['mono']:
|
||||
path_txt = {key: os.path.join(dir_out[key], basename + '.txt')}
|
||||
save_txts(path_txt[key], boxes, all_outputs[key], params, mode=key, cat=cat)
|
||||
|
||||
for key in self.METHODS:
|
||||
path_txt = {key: os.path.join(dir_out[key], basename + '.txt')}
|
||||
save_txts(path_txt[key], boxes, all_outputs[key], params, mode=key, cat=cat)
|
||||
|
||||
# STEREO BASELINES
|
||||
if self.baselines:
|
||||
dic_xyz = self._run_stereo_baselines(basename, boxes, keypoints, zzs, path_calib)
|
||||
|
||||
for key in dic_xyz:
|
||||
all_outputs[key] = all_outputs['monoloco'].copy()
|
||||
all_outputs[key][0] = dic_xyz[key]
|
||||
all_inputs[key] = boxes
|
||||
# stereo baselines
|
||||
if self.baselines['stereo']:
|
||||
all_inputs = {}
|
||||
dic_xyz = self._run_stereo_baselines(basename, boxes, keypoints, zzs, path_calib)
|
||||
for key in dic_xyz:
|
||||
all_outputs[key] = all_outputs['monoloco'].copy()
|
||||
all_outputs[key][0] = dic_xyz[key]
|
||||
all_inputs[key] = boxes
|
||||
|
||||
path_txt[key] = os.path.join(dir_out[key], basename + '.txt')
|
||||
save_txts(path_txt[key], all_inputs[key], all_outputs[key], params, mode='baseline', cat=cat)
|
||||
save_txts(path_txt[key], all_inputs[key], all_outputs[key], params, mode='baseline', cat='cat')
|
||||
|
||||
print("\nSaved in {} txt {} annotations. Not found {} images".format(cnt_file, cnt_ann, cnt_no_file))
|
||||
|
||||
if 'monstereo' in self.METHODS:
|
||||
if self.net == 'monstereo':
|
||||
print("STEREO:")
|
||||
for key in self.baselines:
|
||||
for key in self.baselines['stereo']:
|
||||
print("Annotations corrected using {} baseline: {:.1f}%".format(
|
||||
key, self.cnt_disparity[key] / cnt_ann * 100))
|
||||
print("Maximum possible stereo associations: {:.1f}%".format(self.cnt_disparity['max'] / cnt_ann * 100))
|
||||
print("Not found {}/{} stereo files".format(self.cnt_no_stereo, cnt_file))
|
||||
|
||||
create_empty_files(dir_out) # Create empty files for official evaluation
|
||||
if self.generate_official:
|
||||
create_empty_files(dir_out, self.net) # Create empty files for official evaluation
|
||||
|
||||
def _run_stereo_baselines(self, basename, boxes, keypoints, zzs, path_calib):
|
||||
|
||||
@ -247,11 +256,10 @@ def save_txts(path_txt, all_inputs, all_outputs, all_params, mode='monoloco', ca
|
||||
ff.write("\n")
|
||||
|
||||
|
||||
def create_empty_files(dir_out):
|
||||
def create_empty_files(dir_out, net):
|
||||
"""Create empty txt files to run official kitti metrics on MonStereo and all other methods"""
|
||||
|
||||
methods = ['pseudo-lidar', 'monopsr', '3dop', 'm3d', 'oc-stereo', 'e2e']
|
||||
methods = []
|
||||
methods = ['pseudo-lidar', 'monopsr', '3dop', 'm3d', 'oc-stereo', 'e2e', 'monodis', 'smoke']
|
||||
dirs = [os.path.join('data', 'kitti', method) for method in methods]
|
||||
dirs_orig = [os.path.join('data', 'kitti', method + '-orig') for method in methods]
|
||||
|
||||
@ -266,8 +274,7 @@ def create_empty_files(dir_out):
|
||||
# If the file exits, rewrite in new folder, otherwise create empty file
|
||||
read_and_rewrite(path_orig, path)
|
||||
|
||||
for method in ('monoloco_pp', 'monstereo'):
|
||||
for i in range(7481):
|
||||
name = "0" * (6 - len(str(i))) + str(i) + '.txt'
|
||||
ff = open(os.path.join(dir_out[method], name), "a+")
|
||||
ff.close()
|
||||
for i in range(7481):
|
||||
name = "0" * (6 - len(str(i))) + str(i) + '.txt'
|
||||
ff = open(os.path.join(dir_out[net], name), "a+")
|
||||
ff.close()
|
||||
|
||||
@ -70,7 +70,7 @@ def cli():
|
||||
# Training
|
||||
training_parser.add_argument('--joints', help='Json file with input joints',
|
||||
default='data/arrays/joints-nuscenes_teaser-190513-1846.json')
|
||||
training_parser.add_argument('--save', help='whether to not save model and log file', action='store_true')
|
||||
training_parser.add_argument('--no_save', help='to not save model and log file', action='store_true')
|
||||
training_parser.add_argument('-e', '--epochs', type=int, help='number of epochs to train for', default=500)
|
||||
training_parser.add_argument('--bs', type=int, default=512, help='input batch size')
|
||||
training_parser.add_argument('--monocular', help='whether to train monoloco', action='store_true')
|
||||
@ -83,7 +83,9 @@ def cli():
|
||||
training_parser.add_argument('--hyp', help='run hyperparameters tuning', action='store_true')
|
||||
training_parser.add_argument('--multiplier', type=int, help='Size of the grid of hyp search', default=1)
|
||||
training_parser.add_argument('--r_seed', type=int, help='specify the seed for training and hyp tuning', default=1)
|
||||
training_parser.add_argument('--activity', help='new', action='store_true')
|
||||
training_parser.add_argument('--print_loss', help='print training and validation losses', action='store_true')
|
||||
training_parser.add_argument('--auto_tune_mtl', help='whether to use uncertainty to autotune losses',
|
||||
action='store_true')
|
||||
|
||||
# Evaluation
|
||||
eval_parser.add_argument('--dataset', help='datasets to evaluate, kitti or nuscenes', default='kitti')
|
||||
@ -104,6 +106,9 @@ def cli():
|
||||
eval_parser.add_argument('--variance', help='evaluate keypoints variance', action='store_true')
|
||||
eval_parser.add_argument('--activity', help='evaluate activities', action='store_true')
|
||||
eval_parser.add_argument('--net', help='Choose network: monoloco, monoloco_p, monoloco_pp, monstereo')
|
||||
eval_parser.add_argument('--baselines', help='whether to evaluate stereo baselines', action='store_true')
|
||||
eval_parser.add_argument('--generate_official', help='whether to add empty txt files for official evaluation',
|
||||
action='store_true')
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@ -141,10 +146,7 @@ def main():
|
||||
else:
|
||||
|
||||
from .train import Trainer
|
||||
training = Trainer(joints=args.joints, epochs=args.epochs, bs=args.bs,
|
||||
monocular=args.monocular, dropout=args.dropout, lr=args.lr, sched_step=args.sched_step,
|
||||
n_stage=args.n_stage, sched_gamma=args.sched_gamma, hidden_size=args.hidden_size,
|
||||
r_seed=args.r_seed, save=args.save)
|
||||
training = Trainer(args)
|
||||
|
||||
_ = training.train()
|
||||
_ = training.evaluate()
|
||||
@ -171,8 +173,7 @@ def main():
|
||||
else:
|
||||
if args.generate:
|
||||
from .eval.generate_kitti import GenerateKitti
|
||||
kitti_txt = GenerateKitti(args.model, args.dir_ann, p_dropout=args.dropout, n_dropout=args.n_dropout,
|
||||
hidden_size=args.hidden_size)
|
||||
kitti_txt = GenerateKitti(args)
|
||||
kitti_txt.run()
|
||||
|
||||
if args.dataset == 'kitti':
|
||||
@ -183,7 +184,7 @@ def main():
|
||||
|
||||
elif 'nuscenes' in args.dataset:
|
||||
from .train import Trainer
|
||||
training = Trainer(joints=args.joints, hidden_size=args.hidden_size)
|
||||
training = Trainer(args)
|
||||
_ = training.evaluate(load=True, model=args.model, debug=False)
|
||||
|
||||
else:
|
||||
|
||||
@ -34,10 +34,9 @@ class Trainer:
|
||||
tasks = ('d', 'x', 'y', 'h', 'w', 'l', 'ori', 'aux')
|
||||
val_task = 'd'
|
||||
lambdas = (1, 1, 1, 1, 1, 1, 1, 1)
|
||||
clusters = ['10', '20', '30', '40']
|
||||
|
||||
def __init__(self, joints, epochs=100, bs=256, dropout=0.2, lr=0.002,
|
||||
sched_step=20, sched_gamma=1, hidden_size=256, n_stage=3, r_seed=0, n_samples=100,
|
||||
monocular=False, save=False, print_loss=True):
|
||||
def __init__(self, args):
|
||||
"""
|
||||
Initialize directories, load the data and parameters for the training
|
||||
"""
|
||||
@ -49,31 +48,29 @@ class Trainer:
|
||||
dir_logs = os.path.join('data', 'logs')
|
||||
if not os.path.exists(dir_logs):
|
||||
warnings.warn("Warning: default logs directory not found")
|
||||
assert os.path.exists(joints), "Input file not found"
|
||||
assert os.path.exists(args.joints), "Input file not found"
|
||||
|
||||
self.joints = joints
|
||||
self.num_epochs = epochs
|
||||
self.save = save
|
||||
self.print_loss = print_loss
|
||||
self.monocular = monocular
|
||||
self.lr = lr
|
||||
self.sched_step = sched_step
|
||||
self.sched_gamma = sched_gamma
|
||||
self.clusters = ['10', '20', '30', '40']
|
||||
self.hidden_size = hidden_size
|
||||
self.n_stage = n_stage
|
||||
self.joints = args.joints
|
||||
self.num_epochs = args.epochs
|
||||
self.no_save = args.no_save
|
||||
self.print_loss = args.print_loss
|
||||
self.monocular = args.monocular
|
||||
self.lr = args.lr
|
||||
self.sched_step = args.sched_step
|
||||
self.sched_gamma = args.sched_gamma
|
||||
self.hidden_size = args.hidden_size
|
||||
self.n_stage = args.n_stage
|
||||
self.dir_out = dir_out
|
||||
self.n_samples = n_samples
|
||||
self.r_seed = r_seed
|
||||
self.auto_tune_mtl = False
|
||||
self.r_seed = args.r_seed
|
||||
self.auto_tune_mtl = args.auto_tune_mtl
|
||||
|
||||
# Select the device
|
||||
use_cuda = torch.cuda.is_available()
|
||||
self.device = torch.device("cuda" if use_cuda else "cpu")
|
||||
print('Device: ', self.device)
|
||||
torch.manual_seed(r_seed)
|
||||
torch.manual_seed(self.r_seed)
|
||||
if use_cuda:
|
||||
torch.cuda.manual_seed(r_seed)
|
||||
torch.cuda.manual_seed(self.r_seed)
|
||||
|
||||
# Remove auxiliary task if monocular
|
||||
if self.monocular and self.tasks[-1] == 'aux':
|
||||
@ -95,25 +92,28 @@ class Trainer:
|
||||
input_size = 34
|
||||
output_size = 9
|
||||
|
||||
name = 'monoloco_pp' if self.monocular else 'monstereo'
|
||||
now = datetime.datetime.now()
|
||||
now_time = now.strftime("%Y%m%d-%H%M")[2:]
|
||||
name_out = 'monstereo-' + now_time
|
||||
if self.save:
|
||||
name_out = name + '-' + now_time
|
||||
if not self.no_save:
|
||||
self.path_model = os.path.join(dir_out, name_out + '.pkl')
|
||||
self.logger = set_logger(os.path.join(dir_logs, name_out))
|
||||
self.logger.info("Training arguments: \nepochs: {} \nbatch_size: {} \ndropout: {}"
|
||||
"\nmonocular: {} \nlearning rate: {} \nscheduler step: {} \nscheduler gamma: {} "
|
||||
"\ninput_size: {} \noutput_size: {}\nhidden_size: {} \nn_stages: {} "
|
||||
"\nr_seed: {} \nlambdas: {} \ninput_file: {}"
|
||||
.format(epochs, bs, dropout, self.monocular, lr, sched_step, sched_gamma, input_size,
|
||||
output_size, hidden_size, n_stage, r_seed, self.lambdas, self.joints))
|
||||
.format(args.epochs, args.bs, args.dropout, self.monocular,
|
||||
args.lr, args.sched_step, args.sched_gamma, input_size,
|
||||
output_size, args.hidden_size, args.n_stage, args.r_seed,
|
||||
self.lambdas, self.joints))
|
||||
else:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
# Dataloader
|
||||
self.dataloaders = {phase: DataLoader(KeypointsDataset(self.joints, phase=phase),
|
||||
batch_size=bs, shuffle=True) for phase in ['train', 'val']}
|
||||
batch_size=args.bs, shuffle=True) for phase in ['train', 'val']}
|
||||
|
||||
self.dataset_sizes = {phase: len(KeypointsDataset(self.joints, phase=phase))
|
||||
for phase in ['train', 'val']}
|
||||
@ -122,15 +122,15 @@ class Trainer:
|
||||
self.logger.info('Sizes of the dataset: {}'.format(self.dataset_sizes))
|
||||
print(">>> creating model")
|
||||
|
||||
self.model = MonStereoModel(input_size=input_size, output_size=output_size, linear_size=hidden_size,
|
||||
p_dropout=dropout, num_stage=self.n_stage, device=self.device)
|
||||
self.model = MonStereoModel(input_size=input_size, output_size=output_size, linear_size=args.hidden_size,
|
||||
p_dropout=args.dropout, num_stage=self.n_stage, device=self.device)
|
||||
self.model.to(self.device)
|
||||
print(">>> model params: {:.3f}M".format(sum(p.numel() for p in self.model.parameters()) / 1000000.0))
|
||||
print(">>> loss params: {}".format(sum(p.numel() for p in self.mt_loss.parameters())))
|
||||
|
||||
# Optimizer and scheduler
|
||||
all_params = chain(self.model.parameters(), self.mt_loss.parameters())
|
||||
self.optimizer = torch.optim.Adam(params=all_params, lr=lr)
|
||||
self.optimizer = torch.optim.Adam(params=all_params, lr=args.lr)
|
||||
self.scheduler = lr_scheduler.ReduceLROnPlateau(self.optimizer, 'min')
|
||||
self.scheduler = lr_scheduler.StepLR(self.optimizer, step_size=self.sched_step, gamma=self.sched_gamma)
|
||||
|
||||
@ -243,7 +243,7 @@ class Trainer:
|
||||
self.cout_stats(dic_err['val'], size_eval, clst=clst)
|
||||
|
||||
# Save the model and the results
|
||||
if self.save and not load:
|
||||
if not (self.no_save or load):
|
||||
torch.save(self.model.state_dict(), self.path_model)
|
||||
print('-' * 120)
|
||||
self.logger.info("\nmodel saved: {} \n".format(self.path_model))
|
||||
@ -265,7 +265,6 @@ class Trainer:
|
||||
|
||||
# Distance
|
||||
errs = torch.abs(extract_outputs(outputs)['d'] - extract_labels(labels)['d'])
|
||||
|
||||
assert rel_frac > 0.99, "Variance of errors not supported with partial evaluation"
|
||||
|
||||
# Uncertainty
|
||||
|
||||
@ -58,7 +58,7 @@ def make_new_directory(dir_out):
|
||||
if os.path.exists(dir_out):
|
||||
shutil.rmtree(dir_out)
|
||||
os.makedirs(dir_out)
|
||||
print("Created empty output directory for {} txt files".format(dir_out))
|
||||
print("Created empty output directory {} ".format(dir_out))
|
||||
|
||||
|
||||
def normalize_hwl(lab):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user