def main(cfg):
orig_cwd = hydra.utils.get_original_cwd()
print(f" Evaluating the checkpoint "
f"'{cfg.test.trained_ckpt_for_inference}'")
loaded_struct = torch.load(
pjn(orig_cwd, cfg.test.trained_ckpt_for_inference))
cfg.train.train_val_split = 1.0
# ^ associate all of the train data with the train_loader below
# (do not split the train data into train + validation)
train_loader, _, test_loader = train.get_dataloaders(cfg, (1, 0, 1))
interval_bounds, label2count_list = make_label2count_list(cfg)
model = SDCNet(label2count_list,
cfg.model.supervised,
load_pretr_weights_vgg=False)
model.load_state_dict(loaded_struct['model_state_dict'], strict=True)
additional_cfg = {'device': None}
if not cfg.resources.disable_cuda and torch.cuda.is_available():
additional_cfg['device'] = torch.device('cuda')
model = model.cuda()
else:
additional_cfg['device'] = torch.device('cpu')
optimizer = None
trainer = train.TrainManager(
model,
optimizer,
cfg,
additional_cfg,
train_loader=None,
val_loader=None,
)
print()
datadir = pjn(cfg.dataset.dataset_rootdir, f"part_{cfg.dataset.part}")
print(f" Evaluating on the (whole) train data and on the test data "
f"(in '{datadir}')")
mae_train, mse_train = trainer.validate(train_loader)
print(f" Metrics on the (whole) train data: "
f"MAE: {mae_train:.2f}, MSE: {mse_train:.2f}")
mae_test, mse_test = trainer.validate(test_loader)
print(f" Metrics on the test data: "
f"MAE: {mae_test:.2f}, MSE: {mse_test:.2f}")
if cfg.test.visualize:
vis_dir_name = f"visualized_part_{cfg.dataset.part}_test_set_predictions"
vis_dir_print = pjn(os.path.relpath(os.getcwd(), orig_cwd),
vis_dir_name)
print(
f" Visualized predictions are being saved to '{vis_dir_print}':")
visualize_predictions(cfg, model, test_loader, vis_dir_name)
def main(cfg):
orig_cwd = hydra.utils.get_original_cwd()
print(f" Exporting the checkpoint "
f"'{cfg.test.trained_ckpt_for_inference}'")
ckpt_path = pjn(orig_cwd, cfg.test.trained_ckpt_for_inference)
loaded_struct = torch.load(ckpt_path)
interval_bounds, label2count_list = make_label2count_list(cfg)
model = SDCNet(label2count_list,
cfg.model.supervised,
load_pretr_weights_vgg=False)
model.load_state_dict(loaded_struct['model_state_dict'], strict=True)
batch_size = 1
x = torch.randn(batch_size, 3, 64 * 1, 64 * 1, requires_grad=False)
if not cfg.resources.disable_cuda and torch.cuda.is_available():
model = model.cuda()
x = x.cuda()
p1, ext = os.path.splitext(ckpt_path)
dir_, bname = os.path.split(p1)
try:
torch.onnx.export(model, x, p1 + ".onnx", opset_version=13)
except:
print(" (!) Failed to export the checkpoint to ONNX format")
else:
print(" Successfully exported the checkpoint to ONNX format")
try:
traced_script_module = torch.jit.trace(model, x)
traced_script_module.save(bname + "_jit_trace.pt")
except:
print(
" (!) Failed to export the checkpoint to 'torch jit trace' format"
)
else:
print(
" Successfully exported the checkpoint to 'torch jit trace' format"
)
try:
script_module = torch.jit.script(model)
script_module.save(bname + "_jit_script.pt")
except:
print(
" (!) Failed to export the checkpoint to 'torch jit script' format"
)
else:
print(
" Successfully exported the checkpoint to 'torch jit script' format"
)
def main(cfg):
"""
Create data loaders, the model instance, optimizer and TrainManager()
object. Run the training process.
"""
orig_cwd = hydra.utils.get_original_cwd()
train_loader, val_loader, test_loader = get_dataloaders(cfg, (1, 1, 0))
interval_bounds, label2count_list = make_label2count_list(cfg)
model = SDCNet(
label2count_list,
cfg.model.supervised,
load_pretr_weights_vgg=True)
if cfg.train.pretrained_ckpt:
print(f" Using pretrained model and its checkpoint "
f"'{cfg.train.pretrained_ckpt}'")
loaded_struct = torch.load(pjn(orig_cwd, cfg.train.pretrained_ckpt))
model.load_state_dict(loaded_struct['model_state_dict'], strict=True)
additional_cfg = {'device': None}
if not cfg.resources.disable_cuda and torch.cuda.is_available():
additional_cfg['device'] = torch.device('cuda')
model = model.cuda()
else:
additional_cfg['device'] = torch.device('cpu')
optimizer = torch.optim.SGD(
model.parameters(),
lr=cfg.train.lr_schedule.lr_init,
momentum=cfg.train.optimizer.momentum,
weight_decay=cfg.train.optimizer.weight_decay)
additional_cfg['tbx_wrtr_dir'] = os.getcwd()
with SummaryWriter(additional_cfg['tbx_wrtr_dir']) as tbx_wrtr:
additional_cfg['tbx_wrtr'] = tbx_wrtr
trainer = TrainManager(
model,
optimizer,
cfg,
additional_cfg,
train_loader=train_loader,
val_loader=val_loader)
trainer.train()
def main(cfg):
orig_cwd = hydra.utils.get_original_cwd()
print(f" Running inference using checkpoint "
f"'{cfg.test.trained_ckpt_for_inference}'")
ckpt_path = pjn(orig_cwd, cfg.test.trained_ckpt_for_inference)
loaded_struct = torch.load(ckpt_path)
interval_bounds, label2count_list = make_label2count_list(cfg)
model = SDCNet(label2count_list,
cfg.model.supervised,
load_pretr_weights_vgg=False)
model.load_state_dict(loaded_struct['model_state_dict'], strict=True)
if not cfg.resources.disable_cuda and torch.cuda.is_available():
model = model.cuda()
if cfg.test.visualize:
vis_dir_print = pjn(os.path.relpath(os.getcwd(), orig_cwd),
'visualized_predictions')
print(f" Visualized predictions are being saved to '{vis_dir_print}'")
get_predictions(cfg, model)
class Object(object):
pass
class Cfg:
# imitates config from hydra (only 3 parameters from the full config
# are required here to instantiate SDCNet)
def __init__(self, num_intervals, interval_step, partition_method):
self.dataset = Object()
self.dataset.num_intervals = num_intervals
self.model = Object()
self.model.interval_step = interval_step
self.model.partition_method = partition_method
cfg = Cfg(22, 0.5, 2) # for part_A
#cfg = Cfg(7, 0.5, 2) # for part_B
interval_bounds, label2count_list = make_label2count_list(cfg)
# ================================================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sdcnet_instance = SDCNet(label2count_list,
supervised=True,
load_pretr_weights_vgg=True).to(device)
batch_size = 1
h = 64 * 3
w = 64 * 2
summary(sdcnet_instance, (3, h, w))
x = torch.randn(batch_size, 3, h, w, requires_grad=False).to(device)
out_list = sdcnet_instance(x)
def get_dataloaders(cfg, train_val_test_mask):
"""
Constructs Dataset objects and corresponding DataLoader objects.
Args:
cfg: the global configuration (hydra).
train_val_test_mask: a list or tuple containing three values
that will be interpreted as booleans, for example (1,1,0).
The corresponding DataLoaders for the (train, val, test)
filtered list of Datasets will be created and returned.
Returns:
The tuple (train_loader, val_loader, test_loader). The loaders
corresponding to the `False` values in `train_val_test_mask`
are set to `None`.
"""
interval_bounds, label2count_list = make_label2count_list(cfg)
assert len(label2count_list) == interval_bounds.shape[0] + 1, \
"Number of class labels must be equal to (1 + number of interval " \
"boundaries), but the equality does not hold"
rgb_mean_train = dtst.calc_rgb_mean_train(cfg)
composed_transf_train = transforms.Compose([
dtst.Normalize(rgb_mean_train),
dtst.ToCHW(),
dtst.RandomHorizontalFlip(p=0.5),
dtst.QuasiRandomCrop(),
dtst.PadToMultipleOf64(),
dtst.AddGtCountsLabels(interval_bounds),
])
composed_transf_test = transforms.Compose([
dtst.Normalize(rgb_mean_train),
dtst.ToCHW(),
dtst.PadToMultipleOf64(),
])
if train_val_test_mask[0]:
train_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='train_data',
shuffle_seed=42,
rel_inds=(0.0, cfg.train.train_val_split),
transform=composed_transf_train)
train_loader = DataLoader(
train_dataset,
batch_size=cfg.train.batch_size,
shuffle=True,
num_workers=4)
else:
train_loader = None
if train_val_test_mask[1]:
val_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='train_data',
shuffle_seed=42,
rel_inds=(cfg.train.train_val_split, 1.0),
transform=composed_transf_test)
val_loader = DataLoader(
val_dataset,
batch_size=cfg.validation.batch_size,
shuffle=False,
num_workers=4)
else:
val_loader = None
if train_val_test_mask[2]:
test_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='test_data',
shuffle_seed=None,
rel_inds=(0.0, 1.0),
transform=composed_transf_test)
test_loader = DataLoader(
test_dataset,
batch_size=cfg.test.batch_size,
shuffle=False,
num_workers=4)
else:
test_loader = None
return train_loader, val_loader, test_loader
def get_dataloaders(cfg, train_val_test_mask):
interval_bounds, label2count_list = make_label2count_list(cfg)
assert len(label2count_list) == interval_bounds.shape[0] + 1, \
"Number of class labels must be equal to (1 + number of interval " \
"boundaries), but the equality does not hold"
rgb_mean_train = dtst.calc_rgb_mean_train(cfg)
composed_transf_train = transforms.Compose([
dtst.Normalize(rgb_mean_train),
dtst.ToCHW(),
dtst.RandomHorizontalFlip(p=0.5),
dtst.QuasiRandomCrop(),
dtst.PadToMultipleOf64(),
dtst.AddGtCountsLabels(interval_bounds),
])
composed_transf_test = transforms.Compose([
dtst.Normalize(rgb_mean_train),
dtst.ToCHW(),
dtst.PadToMultipleOf64(),
])
if train_val_test_mask[0]:
train_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='train_data',
shuffle_seed=42,
rel_inds=(0.0, cfg.train.train_val_split),
transform=composed_transf_train)
train_loader = DataLoader(
train_dataset,
batch_size=cfg.train.batch_size,
shuffle=True,
num_workers=4)
else:
train_loader = None
if train_val_test_mask[1]:
val_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='train_data',
shuffle_seed=42,
rel_inds=(cfg.train.train_val_split, 1.0),
transform=composed_transf_test)
val_loader = DataLoader(
train_dataset,
batch_size=cfg.validation.batch_size,
shuffle=False,
num_workers=4)
else:
val_loader = None
if train_val_test_mask[2]:
test_dataset = dtst.ShanghaiTechDataset(
cfg,
subdir='test_data',
shuffle_seed=None,
rel_inds=(0.0, 1.0),
transform=composed_transf_test)
test_loader = DataLoader(
test_dataset,
batch_size=cfg.test.batch_size,
shuffle=False,
num_workers=4)
else:
test_loader = None
return train_loader, val_loader, test_loader