if evaluating:
loader_train = DataLoader(dataset_train,
batch_size=1,
collate_fn=custom_collate)
else:
loader_train = DataLoader(dataset_train,
batch_size=batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
drop_last=True,
collate_fn=custom_collate)
loader_val = DataLoader(dataset_val, batch_size=1, collate_fn=custom_collate)
total_params = get_n_params(model.parameters())
ft_params = get_n_params(model.fine_tune_params())
ran_params = get_n_params(model.random_init_params())
spp_params = get_n_params(model.backbone.spp.parameters())
assert total_params == (ft_params + ran_params)
print(
f'Num params: {total_params:,} = {ran_params:,}(random init) + {ft_params:,}(fine tune)'
)
print(f'SPP params: {spp_params:,}')
if evaluating:
eval_loaders = [(loader_val, 'val')] # , (loader_train, 'train')]
store_dir = f'{dir_path}/out/'
for d in ['', 'val', 'train', 'training']:
os.makedirs(store_dir + d, exist_ok=True)
to_color = ColorizeLabels(color_info)
batch_size = bs = 8
print(f'Batch size: {bs}')
nw = 4
subset_sampler_train = WeightedRandomSampler(weights=weights, num_samples=len(dataset_train))
# subset_sampler_train = None
loader_val = DataLoader(dataset_val, batch_size=1, collate_fn=custom_collate, num_workers=nw)
if evaluating:
loader_train = DataLoader(dataset_train, batch_size=1, collate_fn=custom_collate, num_workers=nw)
else:
loader_train = DataLoader(dataset_train, batch_size=batch_size, num_workers=nw, pin_memory=True,
drop_last=True, collate_fn=custom_collate, sampler=subset_sampler_train)
total_params = get_n_params(model.parameters())
ft_params = get_n_params(model.fine_tune_params())
ran_params = get_n_params(model.random_init_params())
assert total_params == (ft_params + ran_params)
print(f'Num params: {total_params:,} = {ran_params:,}(random init) + {ft_params:,}(fine tune)')
eval_observers = []
if False and evaluating:
eval_loaders = [(loader_val, 'validation'), (loader_train, 'training')]
store_dir = f'{dir_path}/out/'
for d in ['', 'validation', 'training']:
os.makedirs(store_dir + d, exist_ok=True)
to_color = ColorizeLabels(color_info)
to_image = Compose([Numpy(), to_color])
eval_observers += [StorePreds(store_dir, to_image, to_color)]
