def forward(self, embeddings, label):
cos_theta, origin_cos = calc_logits(embeddings, self.kernel)
cos_theta_, _ = calc_logits(embeddings, self.kernel.detach())
mask = torch.zeros_like(cos_theta)
mask.scatter_(1, label.view(-1, 1).long(), 1.0)
sample_num = embeddings.size(0)
tmp_cos_theta = cos_theta - 2 * mask
tmp_cos_theta_ = cos_theta_ - 2 * mask
target_cos_theta = cos_theta[torch.arange(0, sample_num),
label].view(-1, 1)
target_cos_theta_ = cos_theta_[torch.arange(0, sample_num),
label].view(-1, 1)
target_cos_theta_m = target_cos_theta - self.margin
far = 1 / (self.out_features - 1)
# far = 1e-4
topk_mask = torch.greater(tmp_cos_theta, target_cos_theta)
topk_sum = torch.sum(topk_mask.to(torch.int32))
dist.all_reduce(topk_sum)
far_rank = math.ceil(
far * (sample_num *
(self.out_features - 1) * dist.get_world_size() - topk_sum))
cos_theta_neg_topk = torch.topk(
(tmp_cos_theta - 2 * topk_mask.to(torch.float32)).flatten(),
k=far_rank)[0]
cos_theta_neg_topk = all_gather_tensor(cos_theta_neg_topk.contiguous())
cos_theta_neg_th = torch.topk(cos_theta_neg_topk, k=far_rank)[0][-1]
cond = torch.mul(torch.bitwise_not(topk_mask),
torch.greater(tmp_cos_theta, cos_theta_neg_th))
_, cos_theta_neg_topk_index = torch.where(cond)
cos_theta_neg_topk = torch.mul(cond.to(torch.float32), tmp_cos_theta)
cos_theta_neg_topk_ = torch.mul(cond.to(torch.float32), tmp_cos_theta_)
cond = torch.greater(target_cos_theta_m, cos_theta_neg_topk)
cos_theta_neg_topk = torch.where(cond, cos_theta_neg_topk,
cos_theta_neg_topk_)
cos_theta_neg_topk = torch.pow(cos_theta_neg_topk, 2)
times = torch.sum(torch.greater(cos_theta_neg_topk,
0).to(torch.float32),
dim=1,
keepdim=True)
times = torch.where(torch.greater(times, 0), times,
torch.ones_like(times))
cos_theta_neg_topk = torch.sum(cos_theta_neg_topk, dim=1,
keepdim=True) / times
target_cos_theta_m = target_cos_theta_m - (
1 + target_cos_theta_) * cos_theta_neg_topk
cos_theta.scatter_(1, label.view(-1, 1).long(), target_cos_theta_m)
output = cos_theta * self.scale
return output, origin_cos * self.scale
def fn(global_step):
"""Returns a learning rate given the current training iteration."""
float_training_steps = ff(training_steps)
global_step = ff(global_step)
# ensure we don't train longer than training steps
global_step = torch.min(global_step, float_training_steps)
constant_steps = float_training_steps * constant_fraction
x = torch.max(ff(global_step), ff(constant_steps))
min_learning_rate = min_learning_rate_mult * learning_rate
if warmup_fraction:
min_warmup_fraction = torch.max(warmup_fraction, constant_fraction)
warmup_steps = float_training_steps * min_warmup_fraction
is_warmup = ff(torch.greater(ff(warmup_steps), ff(global_step)))
warmup_lr = (global_step / warmup_steps) * learning_rate
else:
warmup_lr = learning_rate
is_warmup = 0.0
step = x - constant_steps
constant_and_decay = (learning_rate - min_learning_rate) * (
torch.cos(step * np.pi /
(float_training_steps - constant_steps)) / 2.0 +
0.5) + min_learning_rate
new_learning_rate = constant_and_decay * (
1.0 - is_warmup) + is_warmup * (warmup_lr)
return new_learning_rate
def forward(self, x):
# check dims
if x.dim() != 4:
raise ValueError('expected 4D input (got {}D input)'.format(
x.dim()))
if self.training:
# batch stats
x_min = torch.amin(x, dim=(0, 1))
x_max = torch.amax(x, dim=(0, 1))
if self.first:
self.max = x_max
self.min = x_min
self.first = False
else:
# update min max with masking correect entries
max_mask = torch.greater(x_max, self.max)
self.max = (max_mask * x_max) + \
(torch.logical_not(max_mask) * self.max)
min_mask = torch.less(x_min, self.min)
self.min = (min_mask * x_min) + \
(torch.logical_not(min_mask) * self.min)
self.max_min = self.max - self.min + 1e-13
# scale batch
x = (x - self.min) / self.max_min
return x
def comparison_ops(self):
a = torch.randn(4)
b = torch.randn(4)
return (
torch.allclose(a, b),
torch.argsort(a),
torch.eq(a, b),
torch.equal(a, b),
torch.ge(a, b),
torch.greater_equal(a, b),
torch.gt(a, b),
torch.greater(a, b),
torch.isclose(a, b),
torch.isfinite(a),
torch.isin(a, b),
torch.isinf(a),
torch.isposinf(a),
torch.isneginf(a),
torch.isnan(a),
torch.isreal(a),
torch.kthvalue(a, 1),
torch.le(a, b),
torch.less_equal(a, b),
torch.lt(a, b),
torch.less(a, b),
torch.maximum(a, b),
torch.minimum(a, b),
torch.fmax(a, b),
torch.fmin(a, b),
torch.ne(a, b),
torch.not_equal(a, b),
torch.sort(a),
torch.topk(a, 1),
torch.msort(a),
)
def kspace_threshhold(image_orig, threshhold=1e-8, name="kspace_threshhold"):
"""Find k-space mask based on threshhold.
Anything less the specified threshhold is set to 0.
Anything above the specified threshhold is set to 1.
"""
mask_x = torch.greater(torch.abs(image_orig), threshhold)
mask_x = torch.cast(mask_x, dtype=torch.float32)
return mask_x
def update(self, output):
y_pred, y = output[0].detach(), output[1].detach()
y_pred = torch.greater(y_pred, self.threshold)
correct = torch.eq(y_pred, y).view(-1)
correct = torch.eq(y_pred, y).view(-1)
self._matches += torch.sum(correct).to(self._device)
self._total += correct.shape[0]
def not_done(self, i):
y = self.score * torch.cast(self.flag, torch.floatx())
y = torch.reduce_min(y, axis=1)
fs = torch.reduce_any(self.flags, axis=1)
old = y + (1.0 - torch.cast(fs, torch.floatx())) * utils.big_neg
n = torch.int_shape(self.tgt)[-1]
new = self.logp[:, 0] / self.penalty(n)
done = torch.reduce_all(torch.greater(old, new))
return torch.logical_and(torch.less(i, n), torch.logical_not(done))
def update_z(a_last, W, b, a, u, v, rho):
z1 = torch.matmul(W, a_last) + b - u / rho
z2 = (z1 + a + v) / 2
z1 = torch.minimum(z1, torch.tensor(0))
z2 = torch.maximum(z2, torch.tensor(0))
value1 = z_obj(a_last, W, b, z1, u, v, a, rho)
value2 = z_obj(a_last, W, b, z2, u, v, a, rho)
imask = torch.greater(value1, value2)
z = torch.where(imask, z2, z1)
return z
def count_incorrects(ramans, predict_confidence_round):
ramans = ramans[0]
target_confidence = ramans[:, 0]
target_position = ramans[:, 1]
target_position = absolute_position(target_position).flatten()
more = torch.greater(predict_confidence_round,
target_confidence).float().sum().detach().item()
less = torch.less(predict_confidence_round,
target_confidence).float().sum().detach().item()
incorrect = more + less
total = target_confidence.float().sum().detach().item()
return int(more), int(less), int(incorrect), int(
total), target_confidence, target_position
def train(epoch):
if epoch > 0 and epoch % decay_every == 0:
lr = lr * decay_rate
for param_group in optimizer.param_groups:
param_group['lr'] = lr
net.train()
optimizer.zero_grad()
jaccard_indices = []
for ligand, protein, site in progressbar.progressbar(
scpdb_dataloader_train):
ligand = ligand.to(device)
protein = protein.to(device)
site = site.to(device)
segmentation = net(protein, ligand)
ground_truth_segmentation = construct_ground_truth_segmentation(site)
ground_truth_segmentation = torch.from_numpy(ground_truth_segmentation)
loss = bce_logit_loss(segmentation, ground_truth_segmentation)
loss.backward()
# track accuracy
predictions = (segmentation > 0.5).int()
true_positives = torch.logical_and(predictions,
ground_truth_segmentation)
false_positives = torch.greater(predictions, ground_truth_segmentation)
false_negatives = torch.less(predictions, ground_truth_segmentation)
num_correct = torch.sum(true_positives)
iou = num_correct / (num_correct + torch.sum(false_positives) +
torch.sum(false_negatives))
jaccard_indices.append(iou)
# optimize
optimizer.step()
optimizer.zero_grad()
mean_iou = sum(jaccard_indices / len(jaccard_indices))
print("Epoch {} - Train: {:06.4f}".format(epoch, mean_iou))
def forward(self):
a = torch.tensor(0)
b = torch.tensor(1)
return len(
torch.allclose(a, b),
torch.argsort(a),
torch.eq(a, b),
torch.eq(a, 1),
torch.equal(a, b),
torch.ge(a, b),
torch.ge(a, 1),
torch.greater_equal(a, b),
torch.greater_equal(a, 1),
torch.gt(a, b),
torch.gt(a, 1),
torch.greater(a, b),
torch.isclose(a, b),
torch.isfinite(a),
torch.isin(a, b),
torch.isinf(a),
torch.isposinf(a),
torch.isneginf(a),
torch.isnan(a),
torch.isreal(a),
torch.kthvalue(a, 1),
torch.le(a, b),
torch.le(a, 1),
torch.less_equal(a, b),
torch.lt(a, b),
torch.lt(a, 1),
torch.less(a, b),
torch.maximum(a, b),
torch.minimum(a, b),
torch.fmax(a, b),
torch.fmin(a, b),
torch.ne(a, b),
torch.ne(a, 1),
torch.not_equal(a, b),
torch.sort(a),
torch.topk(a, 1),
torch.msort(a),
)
def _cos_hinge_loss(self, U_s, q_s, margin=0.1):
"""
:param U_s: mean support embeddings of shape (N, emb_dim)
:param q_s: query embeddings of shape (N, q_n, emb_dim)
:return loss: scalar
:return accuracy: scalar
"""
N = q_s.shape[0]
q_s = q_s.reshape(-1, q_s.shape[-1]) # (N * q_n, emb_dim)
# Similarities of every support mean sentence with every query sentence
similarities = torch.mm(U_s, torch.transpose(q_s, 0,
1)) # (N, N * q_n)
similarities = similarities.view(N, N, -1) # (N, N, q_n)
# Gets the diagonal
mask_pos = torch.eye(N) == 1
positives = similarities[mask_pos, :] # (N, q_n)
positives_ex = positives.view(N, 1, -1) # (N, 1, q_n)
# Gets everything but the diagonal
mask_neg = torch.eye(N) == 0
negatives = similarities[mask_neg, :] # (N * (N - 1), q_n)
negatives = negatives.view(N, N - 1, -1) # (N, N - 1, q_n)
# loss
loss = torch.clamp(margin - positives_ex + negatives, min=0.0)
loss = torch.mean(loss)
# accuracy
max_of_negs = torch.max(negatives, dim=1).values # (N, q_n)
accuracy = torch.greater(positives, max_of_negs).float()
accuracy = torch.mean(accuracy)
return loss, accuracy
def _beam_search_decode(_logprobs, _logprobs_len, _vocab, _beam_size,
_cutoff_top_n, _cutoff_prob, _ext_scoring_func,
_alpha):
assert (_cutoff_top_n is None) or (_cutoff_prob is None)
if (_cutoff_top_n is None) and (_cutoff_prob is None):
_cutoff_prob = 1.0
hypos = set()
indices2tokens = vocab.indices2tokens()
probs_b, probs_nb = defaultdict(float), defaultdict(float)
hypos.add('')
probs_b[''], probs_nb[''] = 1.0, 0.0
probs_b_new, probs_nb_new = dict(), dict()
for t in range(_logprobs_len):
hypos_new = set()
probs_b_new, probs_nb_new = defaultdict(float), defaultdict(float)
probs_t = torch.exp(_logprobs[t]).detach().cpu()
decrease_ids = torch.argsort(-probs_t).detach().cpu().numpy()
if _cutoff_prob is not None:
cond_idxs = torch.where(
torch.greater(torch.cumsum(probs_t[decrease_ids], dim=0),
_cutoff_prob))[0]
if len(cond_idxs) == 0:
_cutoff_top_n = len(decrease_ids)
else:
_cutoff_top_n = cond_idxs[0].item()
if _cutoff_top_n == 0:
_cutoff_top_n = 1
for line in hypos:
for c_idx in decrease_ids[:_cutoff_top_n]:
c, c_prob = indices2tokens[c_idx], probs_t[c_idx]
if c == '<blank>':
probs_b_new[line] += c_prob * (probs_b[line] +
probs_nb[line])
else:
l_end = line[-1] if len(line) > 0 else ''
l_plus = line + c
if c == l_end:
probs_nb_new[line] += c_prob * probs_nb[line]
probs_nb_new[l_plus] += c_prob * probs_b[line]
elif c == ' ':
if _ext_scoring_func is None:
p_w = 1.0
else:
p_w = _ext_scoring_func(line)
probs_b_new[l_plus] += np.power(
p_w, _alpha) * c_prob * (probs_b[line] +
probs_nb[line])
else:
probs_nb_new[l_plus] += c_prob * (probs_b[line] +
probs_nb[line])
hypos_new.add(l_plus)
hypos_new.add(line)
hypos_new = list(
sorted(hypos_new,
key=lambda hypo: probs_b_new[hypo] + probs_nb_new[hypo],
reverse=True))[:min(beam_size, len(hypos_new))]
hypos, probs_b, probs_nb = hypos_new, probs_b_new, probs_nb_new
return sorted([(hypo, probs_b_new[hypo] + probs_nb_new[hypo])
for hypo in hypos],
key=lambda key_value: key_value[1],
reverse=True)
def __gt__(self, other):
x0, x1 = self._to_binary_tensor_args(other)
y = torch.greater(x0._t, x1._t)
s = _ox.greater(*_EagerTensor.ox_args([x0, x1]))
return self.from_torch(y, s)
def global_pointer_f1_score(y_true, y_pred):
y_pred = torch.greater(y_pred, 0)
return torch.sum(y_true * y_pred).item(), torch.sum(y_true + y_pred).item()
def evaluate(inference_fn,
dataset,
height,
width,
progress_bar=False,
plot_dir='',
num_plots=0,
prefix='kitti'):
"""Evaluate an iference function for flow with a kitti eval dataset.
Args:
inference_fn: An inference function that produces a flow_field from two
images, e.g. the infer method of UFlow.
dataset: A dataset produced by the method above with for_eval=True.
height: int, the height to which the images should be resized for inference.
width: int, the width to which the images should be resized for inference.
progress_bar: boolean, flag to indicate whether the function should print a
progress_bar during evaluaton.
plot_dir: string, optional path to a directory in which plots are saved (if
num_plots > 0).
num_plots: int, maximum number of qualitative results to plot for the
evaluation.
prefix: str to prefix evaluation keys with in the returned dictionary.
Returns:
A dictionary of floats that represent different evaluation metrics. The keys
of this dictionary are returned by the method list_eval_keys (see below).
"""
eval_start_in_s = time.time()
it = data.DataLoader(dataset)
# it = tf.compat.v1.data.make_one_shot_iterator(dataset)
epe_occ = [] # End point errors.
errors_occ = []
valid_occ = []
epe_noc = [] # End point errors.
errors_noc = []
valid_noc = []
inference_times = []
all_occlusion_results = defaultdict(lambda: defaultdict(int))
for i, test_batch in enumerate(it):
if progress_bar:
sys.stdout.write(':')
sys.stdout.flush()
(image_batch, flow_uv_occ, flow_uv_noc, flow_valid_occ,
flow_valid_noc) = test_batch
image_batch = torch.squeeze(image_batch,0).to(uflow_gpu_utils.device)
flow_uv_occ = torch.squeeze(flow_uv_occ,0).to(uflow_gpu_utils.device)
flow_uv_noc = torch.squeeze(flow_uv_noc, 0).to(uflow_gpu_utils.device)
flow_valid_occ = torch.squeeze(flow_valid_occ.type(torch.float32),0).to(uflow_gpu_utils.device)
flow_valid_noc = torch.squeeze(flow_valid_noc.type(torch.float32),0).to(uflow_gpu_utils.device)
# pylint:disable=cell-var-from-loop
f = lambda: inference_fn(
image_batch[0], #change back to 0 here and below to 1. Trying to get the backward flow
image_batch[1],
input_height=height,
input_width=width,
infer_occlusion=True)
inference_time_in_ms, (final_flow, soft_occlusion_mask) = uflow_utils.time_it(
f, execute_once_before=i == 0)
inference_times.append(inference_time_in_ms)
occ_mask_gt = flow_valid_occ - flow_valid_noc
f_dict = uflow_utils.compute_f_metrics(soft_occlusion_mask * flow_valid_occ,
occ_mask_gt * flow_valid_occ)
best_thresh = -1.
best_f_score = -1.
for thresh, metrics in f_dict.items():
precision = metrics['tp'] / (metrics['tp'] + metrics['fp'] + 1e-6)
recall = metrics['tp'] / (metrics['tp'] + metrics['fn'] + 1e-6)
f1 = 2 * precision * recall / (precision + recall + 1e-6)
if f1 > best_f_score:
best_thresh = thresh
best_f_score = f1
all_occlusion_results[thresh]['tp'] += metrics['tp']
all_occlusion_results[thresh]['fp'] += metrics['fp']
all_occlusion_results[thresh]['tn'] += metrics['tn']
all_occlusion_results[thresh]['fn'] += metrics['fn']
mask_thresh = torch.greater(soft_occlusion_mask, best_thresh).type(torch.float32)
# Image coordinates are swapped in labels
# final_flow = flow.detach().cpu().numpy()
# final_flow = final_flow[::-1, Ellipsis]
#
# final_flow = torch.from_numpy(final_flow.copy()).to(uflow_gpu_utils.device)
endpoint_error_occ = torch.sum((final_flow - flow_uv_occ)**2, dim= 1, keepdim=True)**0.5
gt_flow_abs = torch.sum(flow_uv_occ**2, dim= 1, keepdim=True)**0.5
outliers_occ = torch.logical_and(endpoint_error_occ > 3.,
endpoint_error_occ > 0.05 * gt_flow_abs).type(torch.float32)
endpoint_error_noc = torch.sum((final_flow - flow_uv_noc)**2, dim=1, keepdim=True)**0.5
gt_flow_abs = torch.sum(flow_uv_noc**2, dim= 1, keepdim=True)**0.5
outliers_noc = torch.logical_and(endpoint_error_noc > 3.,
endpoint_error_noc > 0.05 * gt_flow_abs).type(torch.float32)
epe_occ.append(torch.sum(flow_valid_occ * endpoint_error_occ).item())
errors_occ.append(torch.sum(flow_valid_occ * outliers_occ).item())
valid_occ.append(torch.sum(flow_valid_occ).item())
epe_noc.append(
torch.sum(flow_valid_noc * endpoint_error_noc).item())
errors_noc.append(torch.sum(flow_valid_noc * outliers_noc).item())
valid_noc.append(torch.sum(flow_valid_noc).item())
if plot_dir and i < num_plots:
uflow_plotting.complete_paper_plot(
plot_dir,
i,
image_batch[0].detach().cpu().numpy(),
image_batch[1].detach().cpu().numpy(),
final_flow.detach().cpu().numpy(),
flow_uv_occ.detach().cpu().numpy(),
flow_valid_occ.detach().cpu().numpy(), (1. - mask_thresh).detach().cpu().numpy(),
(1. - occ_mask_gt).detach().cpu().numpy(),
frame_skip=None)
if progress_bar:
sys.stdout.write('\n')
sys.stdout.flush()
fmax, best_thresh = uflow_utils.get_fmax_and_best_thresh(all_occlusion_results)
eval_stop_in_s = time.time()
results = {
prefix + '-occl-f-max':
fmax,
prefix + '-best-occl-thresh':
best_thresh,
prefix + '-EPE(occ)':
np.clip(np.mean(np.array(epe_occ) / np.array(valid_occ)), 0.0, 50.0),
prefix + '-ER(occ)':
np.mean(np.array(errors_occ) / np.array(valid_occ)),
prefix + '-EPE(noc)':
np.clip(np.mean(np.array(epe_noc) / np.array(valid_noc)), 0.0, 50.0),
prefix + '-ER(noc)':
np.mean(np.array(errors_noc) / np.array(valid_noc)),
prefix + '-inf-time(ms)':
np.mean(inference_times),
prefix + '-eval-time(s)':
eval_stop_in_s - eval_start_in_s,
}
writer = SummaryWriter(log_dir=FLAGS.tensorboard_logdir_eval)
for key, val in results.items():
writer.add_scalar(key, val, uflow_flags.step // FLAGS.epoch_length)
return results
def forward(self, feature_pyramid1, feature_pyramid2, training=False):
"""Run the model."""
context = None
flow = None
flow_up = None
context_up = None
flows = []
# Go top down through the levels to the second to last one to estimate flow.
for level, (features1, features2) in reversed(
list(enumerate(zip(feature_pyramid1, feature_pyramid2)))[1:]):
# init flows with zeros for coarsest level if needed
if self._shared_flow_decoder and flow_up is None:
batch_size, _, height, width = list(features1.shape)
flow_up = torch.zeros([batch_size, 2, height, width],
dtype=torch.bfloat16
if self._use_bfloat16 else torch.float32)
if self._num_context_up_channels:
num_channels = int(self._num_context_up_channels *
self._channel_multiplier)
context_up = torch.zeros(
[batch_size, num_channels, height, width],
dtype=torch.bfloat16
if self._use_bfloat16 else torch.float32)
# Warp features2 with upsampled flow from higher level.
if flow_up is None or not self._use_feature_warp:
warped2 = features2
else:
warp_up = uflow_utils.flow_to_warp(flow_up)
warped2 = uflow_utils.resample(features2, warp_up)
# Compute cost volume by comparing features1 and warped features2.
features1_normalized, warped2_normalized = normalize_features(
[features1, warped2],
normalize=self._normalize_before_cost_volume,
center=self._normalize_before_cost_volume,
moments_across_channels=True,
moments_across_images=True)
if self._use_cost_volume:
cost_volume = compute_cost_volume(features1_normalized,
warped2_normalized,
max_displacement=4)
else:
concat_features = torch.cat(
(features1_normalized, warped2_normalized), dim=1)
concat_features = F.pad(concat_features, (2, 1, 2, 1))
cost_volume = self._cost_volume_surrogate_convs[level](
concat_features)
cost_volume = torch.nn.LeakyReLU(
self._leaky_relu_alpha)(cost_volume)
if self._shared_flow_decoder:
# This will ensure to work for arbitrary feature sizes per level.
conv_1x1 = self._1x1_shared_decoder[level - 1]
features1 = conv_1x1(features1)
# Compute context and flow from previous flow, cost volume, and features1.
if flow_up is None:
x_in = torch.cat([cost_volume, features1], dim=1)
else:
if context_up is None:
x_in = torch.cat([flow_up, cost_volume, features1], dim=1)
else:
x_in = torch.cat(
[context_up, flow_up, cost_volume, features1], dim=1)
# Use dense-net connections.
x_out = None
if self._shared_flow_decoder:
# self._flow_layers = self._build_flow_layers(in_channels=x_in.shape[1])
# reuse the same flow decoder on all levels
flow_layers = self._flow_layers
else:
# self._flow_layers = self._build_flow_layers(in_channels=x_in.shape[1])
flow_layers = self._flow_layers[level - 1]
for layer in flow_layers[:-1]:
pad = torch.nn.ConstantPad2d((1, 1, 1, 1), 0)
x_in = pad(x_in)
x_out = layer(x_in)
x_in = x_in[:, :, 1:-1, 1:-1]
x_in = torch.cat([x_in, x_out], dim=1)
context = x_out
pad_context = torch.nn.ConstantPad2d((1, 1, 1, 1), 0)
context_pad = pad_context(context)
flow = flow_layers[-1](context_pad)
# context = context[:,:,1:-1,1:-1]
if (training and self._drop_out_rate):
maybe_dropout = torch.greater(torch.rand([]),
self._drop_out_rate)
maybe_dropout = maybe_dropout.type(
torch.bfloat16 if self._use_bfloat16 else torch.float32)
context *= maybe_dropout
flow *= maybe_dropout
if flow_up is not None and self._accumulate_flow:
flow += flow_up
# Upsample flow for the next lower level.
flow_up = uflow_utils.upsample(flow, is_flow=True)
if self._num_context_up_channels:
# self._context_up_layers = self._build_upsample_layers(in_channels= context.shape[1] , out_channels=int(self._num_context_up_channels * self._channel_multiplier))
context_up = self._context_up_layers[level - 1](context)
# Append results to list.
flows.insert(0, flow)
# Refine flow at level 1.
refine = torch.cat([context, flow], dim=1)
pad = torch.nn.ConstantPad2d((1, 1, 1, 1), 0)
refine = pad(refine)
# self._refine_model = self._build_refinement_model(in_channel = refine.shape[1])
refinement = self._refine_model(refine)
if (training and self._drop_out_rate):
refine_if = torch.greater(torch.rand([]), self._drop_out_rate)
refinement *= refine_if.type(
torch.bfloat16 if self._use_bfloat16 else torch.float32)
refined_flow = flow + refinement
flows[0] = refined_flow
return [flow.type(torch.float32) for flow in flows]
