def fit_truncated(data, discrete=True, calc_p=False, manual_xmin=None):
results = powerlaw.Fit(data, discrete=discrete)
if manual_xmin:
xmin = manual_xmin
else:
xmin = results.xmin
dist = powerlaw.Truncated_Power_Law(xmin=xmin, discrete=discrete)
dist.fit(data)
empirical_ks = dist.KS(data)
alpha = dist.alpha
Lambda = dist.Lambda
p_val = None
if calc_p:
num_synthetic = 1000
n = len(data)
n_tail = sum(data >= xmin)
p_tail = float(n_tail) / n
p_count = 0.0
for _ in range(num_synthetic):
test_set = []
for jjj in range(n):
if np.random.rand() > p_tail:
x = np.random.choice(np.arange(0, xmin))
test_set.append(x)
else:
while True:
r = np.random.rand()
x = np.floor((xmin - 0.5) *
(1 - r)**(-1 / (alpha - 1)) + 0.5)
p = (x / xmin)**(-alpha)
if np.random.rand() <= p:
test_set.append(x)
break
test_set = np.asarray(test_set)
test_dist = powerlaw.Truncated_Power_Law(xmin=xmin,
discrete=discrete)
test_dist.fit(test_set)
if test_dist.KS(test_set) > empirical_ks:
p_count += 1
p_val = p_count / 1000
return dist.alpha, dist.Lambda, results.xmin, p_val
def _time_generator(self):
return powerlaw.Truncated_Power_Law(
xmin=self.min_wait_time,
parameters=[1. + self._beta, 1.0 / self._tau]).generate_random()[0]
def train_network():
network = Stage2CountingNet()
model_save_dir = './models_stage_2'
model_save_path = os.path.join(model_save_dir, 'train2')
if not os.path.exists(model_save_path):
os.makedirs(model_save_path)
os.makedirs(os.path.join(model_save_path, 'snapshots'))
os.makedirs(os.path.join(model_save_dir, 'dump'))
os.makedirs(os.path.join(model_save_dir, 'dump_test'))
global f
snapshot_path = os.path.join(model_save_path, 'snapshots')
f = open(os.path.join(model_save_path, 'train0.log'), 'w')
# -- Logging Parameters
log(f, 'args: ' + str(args))
log(f, 'model: ' + str(network), False)
log(f, 'Stage2...')
log(f, 'LR: %.12f.' % (args.lr))
start_epoch = 0
num_epochs = args.epochs
valid_losses = {}
train_losses = {}
for metric in ['loss1', 'new_mae']:
valid_losses[metric] = []
for metric in ['loss1']:
train_losses[metric] = []
batch_size = args.batch_size
num_train_images = len(dataset.data_files['train'])
num_patches_per_image = args.patches
assert (batch_size < (num_patches_per_image * num_train_images))
num_batches_per_epoch = num_patches_per_image * num_train_images // batch_size
assert (num_batches_per_epoch >= 1)
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
network.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
network = load_rot_model_blocks(
network,
snapshot_path='models_stage_1/train2/snapshots/',
excluded_layers=excluded_layers)
shift_thresh = get_shift_thresh()
Lambda = get_lambda()
log(f, "Shift Thresh: {}, Lambda: {}".format(shift_thresh, Lambda))
# -- Main Training Loop
min_valid_loss = 100.
min_valid_epoch = -1
before_BN_weights_sum = check_BN_no_gradient_change(
network, exclude_list=excluded_layers)
before_conv_weights_sum = check_conv_no_gradient_change(
network, exclude_list=excluded_layers)
stop_training = False
global sampled_GT
for e_i, epoch in enumerate(range(start_epoch, num_epochs)):
avg_loss = []
# b_i - batch index
for b_i in range(num_batches_per_epoch):
# Generate next training sample
Xs, _ = dataset.train_get_data(batch_size=args.batch_size)
after_conv_weights_sum = check_conv_no_gradient_change(
network, exclude_list=excluded_layers)
assert (np.all(before_conv_weights_sum == after_conv_weights_sum))
sampled_GT = None
sampled_GT_shape = args.sbs * 7 * 7 * \
(8 // args.kernel_size) * (8 // args.kernel_size)
sampling_parameters = [args.alpha, Lambda]
sampled_GT = powerlaw.Truncated_Power_Law(
parameters=sampling_parameters).generate_random(
sampled_GT_shape)
for s_i, s_val in enumerate(sampled_GT):
if s_val < shift_thresh:
sampled_GT[s_i] = np.random.uniform(0, shift_thresh)
assert (sampled_GT.shape[0] == (sampled_GT_shape)
and sampled_GT.ndim == 1)
train_loss = train_function(Xs, sampled_GT, network, optimizer)
avg_loss.append(train_loss)
# Logging losses after each iteration.
if b_i % 1 == 0:
log(
f, 'Epoch %d [%d]: %s loss: %s.' %
(epoch, b_i, [network.name], train_loss))
after_BN_weights_sum = check_BN_no_gradient_change(
network, exclude_list=excluded_layers)
after_conv_weights_sum = check_conv_no_gradient_change(
network, exclude_list=excluded_layers)
assert (np.all(before_BN_weights_sum == after_BN_weights_sum))
assert (np.all(before_conv_weights_sum == after_conv_weights_sum))
# -- Stats update
avg_loss = np.mean(np.array(avg_loss))
train_losses['loss1'].append(avg_loss)
log(
f, 'TRAIN epoch: ' + str(epoch) + ' train mean loss1:' +
str(avg_loss))
torch.cuda.empty_cache()
log(f, 'Validating...')
epoch_val_losses, valid_mae = test_network(dataset, 'test_valid',
network, True)
log(
f, 'TEST valid epoch: ' + str(epoch) + ' test valid loss1, mae' +
str(epoch_val_losses))
for metric in ['loss1', 'new_mae']:
valid_losses[metric].append(epoch_val_losses[metric])
if e_i > args.ma_window:
valid_losses_smooth = np.mean(
valid_losses['loss1'][-args.ma_window:])
if valid_losses_smooth < min_valid_loss:
min_valid_loss = valid_losses_smooth
min_valid_epoch = e_i
count = 0
else:
count = count + 1
if count > args.patience:
stop_training = True
log(
f, 'Best valid so far epoch: {}, valid_loss: {}'.format(
min_valid_epoch, valid_losses['loss1'][min_valid_epoch]))
# Save networks
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': network.state_dict(),
'optimizer': optimizer.state_dict(),
}, snapshot_path, get_filename(network.name, epoch + 1))
print('saving graphs...')
with open(os.path.join(snapshot_path, 'losses.pkl'), 'wb') as lossfile:
pickle.dump((train_losses, valid_losses), lossfile, protocol=2)
for metric in train_losses.keys():
if "maxima_split" not in metric:
if isinstance(train_losses[metric][0], list):
for i in range(len(train_losses[metric][0])):
plt.plot([a[i] for a in train_losses[metric]])
plt.savefig(
os.path.join(snapshot_path,
'train_%s_%d.png' % (metric, i)))
plt.clf()
plt.close()
plt.plot(train_losses[metric])
plt.savefig(
os.path.join(snapshot_path, 'train_%s.png' % metric))
plt.clf()
plt.close()
for metric in valid_losses.keys():
if isinstance(valid_losses[metric][0], list):
for i in range(len(valid_losses[metric][0])):
plt.plot([a[i] for a in valid_losses[metric]])
plt.savefig(
os.path.join(snapshot_path,
'valid_%s_%d.png' % (metric, i)))
plt.clf()
plt.close()
plt.plot(valid_losses[metric])
plt.savefig(os.path.join(snapshot_path, 'valid_%s.png' % metric))
plt.clf()
plt.close()
if stop_training:
break
network = load_net(network, snapshot_path,
get_filename(network.name, min_valid_epoch + 1))
log(f, 'Testing on best model {}'.format(min_valid_epoch))
epoch_test_losses, mae = test_network(dataset,
'test',
network,
print_output=os.path.join(
model_save_dir, 'dump_test'))
log(
f, 'TEST epoch: ' + str(epoch) + ' test loss1, mae:' +
str(epoch_test_losses) + ", " + str(mae))
log(f, 'Exiting train...')
f.close()
return