def infogan_comp(): from infogan import InfoGAN STYLE_DIM = 2 LABEL_DIM = 10 RAND_DIM = 88 IMG_SHAPE = (28, 28, 1) FIX_STD = True model = InfoGAN(RAND_DIM, STYLE_DIM, LABEL_DIM, IMG_SHAPE, FIX_STD) model.load_weights("./models/infogan/model.ckpt").expect_partial() img_label = np.arange(0, 10).astype(np.int32).repeat(10, axis=0) noise = tf.repeat(tf.random.normal((1, model.rand_dim)), len(img_label), axis=0) def plot(noise, img_label, img_style, n): img_label = tf.convert_to_tensor(img_label, dtype=tf.int32) img_style = tf.convert_to_tensor(img_style, dtype=tf.float32) imgs = model.g.call([noise, img_label, img_style], training=False).numpy() plt.clf() nc, nr = 10, 10 plt.figure(0, (nc * 2, nr * 2)) for c in range(nc): for r in range(nr): i = r * nc + c plt.subplot(nc, nr, i + 1) plt.imshow(imgs[i], cmap="gray_r") plt.axis("off") plt.text(23, 26, int(r), fontsize=23) plt.tight_layout() model_name = model.__class__.__name__.lower() dir_ = "visual/{}".format(model_name) os.makedirs(dir_, exist_ok=True) path = dir_ + "/style{}.png".format(n) plt.savefig(path) img_style = np.concatenate( [np.linspace(-model.style_scale, model.style_scale, 10)] * 10).reshape( (100, 1)).astype(np.float32) plot(noise, img_label, np.concatenate((img_style, np.zeros_like(img_style)), axis=1), 1) plot(noise, img_label, np.concatenate((np.zeros_like(img_style), img_style), axis=1), 2)
def _gaussian_kernel(self, pos): ''' Generates all gaussian kernels needed for the loss function ''' # Get vector of spatial indices idxs = np.indices((self.dim_len[0], self.dim_len[1])) idxs = idxs.reshape(2, -1) idxs = tf.convert_to_tensor(idxs) # Each value of the kernel is computed as exp(dist^2/sigma) # where dist is the distance between position (i, j) and the center pos_rep = tf.repeat(tf.expand_dims(pos, 1), self.dim_len[0] * self.dim_len[1], axis=1) dists = tf.math.reduce_euclidean_norm(tf.cast(idxs - pos_rep, tf.float64), axis=0) gaussian = tf.math.exp(-tf.math.square(dists) / self.sigma) return gaussian
def _spatial_transform(image, rotation=5.0, shear=2.0, hzoom=8.0, wzoom=8.0, hshift=8.0, wshift=8.0): ydim = tf.gather(tf.shape(image), 0) xdim = tf.gather(tf.shape(image), 1) xxdim = xdim % 2 yxdim = ydim % 2 # random rotation, shear, zoom and shift rotation = rotation * tf.random.normal([1], dtype='float32') shear = shear * tf.random.normal([1], dtype='float32') hzoom = 1.0 + tf.random.normal([1], dtype='float32') / hzoom wzoom = 1.0 + tf.random.normal([1], dtype='float32') / wzoom hshift = hshift * tf.random.normal([1], dtype='float32') wshift = wshift * tf.random.normal([1], dtype='float32') m = _get_transform_matrix(rotation, shear, hzoom, wzoom, hshift, wshift) # origin pixels y = tf.repeat(tf.range(ydim // 2, -ydim // 2, -1), xdim) x = tf.tile(tf.range(-xdim // 2, xdim // 2), [ydim]) z = tf.ones([ydim * xdim], dtype='int32') idx = tf.stack([y, x, z]) # destination pixels idx2 = tf.matmul(m, tf.cast(idx, dtype='float32')) idx2 = tf.cast(idx2, dtype='int32') # clip to origin pixels range idx2y = tf.clip_by_value(idx2[0, ], -ydim // 2 + yxdim + 1, ydim // 2) idx2x = tf.clip_by_value(idx2[1, ], -xdim // 2 + xxdim + 1, xdim // 2) idx2 = tf.stack([idx2y, idx2x, idx2[2, ]]) # apply destinations pixels to image idx3 = tf.stack([ydim // 2 - idx2[0, ], xdim // 2 - 1 + idx2[1, ]]) d = tf.gather_nd(image, tf.transpose(idx3)) image = tf.reshape(d, [ydim, xdim, 3]) return image
def draft_decoder(self, input_ids, enc_output, beam_size, length_penalty, temperature, top_p, top_k, batch_size, training=False): """ Inference call, builds a draft output_sequence auto-regressively """ start_ids = tf.repeat(config.CLS_ID, repeats=batch_size) input_ids = tfa.seq2seq.tile_batch(input_ids, multiplier=beam_size) enc_output = tfa.seq2seq.tile_batch(enc_output, multiplier=beam_size) #start_ids = tf.cast(start_ids, dtype=tf.int64) def perform_beam_search(dec_input): return query_decoder(self, enc_output, input_ids, dec_input, batch_size, temperature, top_p, top_k, training=training) predicted_beam_search_op, _, _, attention_dist = beam_search( perform_beam_search, initial_ids=start_ids, beam_size=beam_size, decode_length=config.target_seq_length, vocab_size=config.target_vocab_size, alpha=length_penalty, stop_early=False, eos_id=config.SEP_ID) predicted_output_sequence = predicted_beam_search_op[:, 0, :] return predicted_output_sequence, attention_dist
def grad(dx): # Gradients with respect to off grid signal fft_dx = scale_and_fft_on_image_volume(dx, scaling_coef, grid_size, im_size, norm, im_rank=im_rank) dx_dy = kbinterp(fft_dx, om, interpob) if grad_traj: # Gradients with respect to trajectory locations r = [ tf.linspace(-im_size[i] / 2, im_size[i] / 2 - 1, im_size[i]) for i in range(im_rank) ] # This wont work for multicoil case as the dimension for dx is `batch_size x coil x Nx x Ny` grid_r = tf.cast(tf.meshgrid(*r, indexing='ij'), dx.dtype)[None, ...] ifft_dxr = scale_and_fft_on_image_volume(tf.math.conj(dx) * grid_r, scaling_coef, grid_size, im_size, norm, im_rank=im_rank, do_ifft=True) # Do this when handling batches ifft_dxr = tf.reshape(ifft_dxr, shape=(-1, 1, *ifft_dxr.shape[2:])) inufft_dxr = kbinterp(ifft_dxr, tf.repeat(om, im_rank, axis=0), interpob, conj=True) # Unbatch back the data inufft_dxr = tf.reshape(inufft_dxr, shape=(-1, im_rank, *inufft_dxr.shape[2:])) dx_dom = tf.cast(1j * y * inufft_dxr, om.dtype) # dx_dom = tf.math.reduce_sum(dx_dom, axis=1)[None, :] else: dx_dom = None return dx_dy, dx_dom
def si_weight(states, actions, n_agents): data = tf.concat([states, actions], axis=-1) all_head_key, all_head_agents, all_head_action = [[] for _ in range(3)] with tf.variable_scope('adv_hyer'): for i in range(num_kernel): all_head_key.append( tf.layers.dense(states, 1, kernel_initializer=w_initializer, bias_initializer=b_initializer, name='all_head_key-{}'.format(i))) all_head_agents.append( tf.layers.dense(states, n_agents, kernel_initializer=w_initializer, bias_initializer=b_initializer, name='all_head_agents-{}'.format(i))) all_head_action.append( tf.layers.dense(data, n_agents, kernel_initializer=w_initializer, bias_initializer=b_initializer, name='all_head_action-{}'.format(i))) head_attend_weights = [] for curr_head_key, curr_head_agents, curr_head_action in zip( all_head_key, all_head_agents, all_head_action): x_key = tf.repeat(tf.abs(curr_head_key), n_agents) + 1e-10 x_key = tf.reshape(x_key, shape=[-1, n_agents]) x_agents = tf.sigmoid(curr_head_agents) x_action = tf.sigmoid(curr_head_action) weights = x_key * x_agents * x_action head_attend_weights.append(weights) head_attend = tf.stack(head_attend_weights, axis=1) head_attend = tf.reshape(head_attend, shape=(-1, num_kernel, n_agents)) head_attend = tf.reduce_sum(head_attend, axis=1, keepdims=False) return head_attend
def update_output(self, input_latent, input_labels, img_width=1600): input_latent = tf.convert_to_tensor(input_latent, dtype=tf.float32) input_latent = tf.repeat(input_latent, self.num_img, axis=0) input_labels = tf.convert_to_tensor(input_labels, dtype=tf.float32) img_per_batch = 12 self.output = [] for i in range(self.num_img // img_per_batch): index_start = i * img_per_batch index_end = min((i + 1) * img_per_batch, self.num_img) self.output.extend( self.model([ input_latent[index_start:index_end], input_labels[index_start:index_end] ])) self.output = (np.asarray(self.output) * 0.5) + 0.5 self.img_size = self.output.shape[1] imgs = [] for i in range(self.output.shape[0]): imgs.append( tf.keras.preprocessing.image.array_to_img(self.output[i])) self.output_img = Image.new( 'L', (self.num_img // self.num_rows_chars * self.img_size, self.img_size * self.num_rows_chars)) for i in range(len(imgs)): self.output_img.paste( imgs[i], (self.img_size * (i % (self.num_img // self.num_rows_chars)), (i // (self.num_img // self.num_rows_chars)) * self.img_size)) img_height = int(self.output_img.size[1] * img_width / self.output_img.size[0]) self.output_img = self.output_img.resize((img_width, img_height), resample=Image.NEAREST) q_pix = QPixmap.fromImage(ImageQt.ImageQt(self.output_img)) self.output_label.setPixmap(q_pix)
def shear_img(image, dim, n_channels, shear_factor = 7.0): ''' Shears an image by `shear_factor` degrees Args: image: array of one squared image of size (dim, dim, n_channels) n_channels: (int) shear_factor: (int or float) the shear factor in degrees Returns: An sheared image ''' xdim = dim%2 shr = shear_factor * tf.ones([1], dtype='float32') shear = 3.141593 * shr / 180. def get_3x3_mat(lst): return tf.reshape(tf.concat([lst],axis=0), [3,3]) # shear matrix c2 = tf.math.cos(shear) s2 = tf.math.sin(shear) one = tf.constant([1],dtype='float32') zero = tf.constant([0],dtype='float32') m = get_3x3_mat([one, s2, zero, zero, c2, zero, zero, zero, one]) # list destination pixel indices x = tf.repeat(tf.range(dim//2, -dim//2,-1), dim) y = tf.tile(tf.range(-dim//2, dim//2), [dim]) z = tf.ones([dim*dim], dtype='int32') idx = tf.stack( [x,y,z] ) # rotate destination pixels onto origin pixels idx2 = K.dot(m, tf.cast(idx, dtype='float32')) idx2 = K.cast(idx2, dtype='int32') idx2 = K.clip(idx2, -dim//2+xdim+1, dim//2) # find origin pixel values idx3 = tf.stack([dim//2-idx2[0,], dim//2-1+idx2[1,]]) d = tf.gather_nd(image, tf.transpose(idx3)) return tf.reshape(d,[dim, dim, n_channels])
def test_get_box_as_dotted_lines(self): corners = tf.repeat(tf.expand_dims(tf.expand_dims(tf.range( 0, 8, dtype=tf.float32), axis=0), axis=2), 3, axis=2) corners = tf.concat([corners, corners + 10.], axis=0) num_of_points_per_line = 100 points = box_utils.get_box_as_dotted_lines( corners, num_of_points_per_line=num_of_points_per_line) expected_points = np.repeat(np.linspace( 0., 1., num_of_points_per_line)[:, np.newaxis], repeats=3, axis=1) self.assertAllEqual(points.shape, [2, 12 * num_of_points_per_line, 3]) self.assertAllClose(points[0, :num_of_points_per_line, :], expected_points) self.assertAllClose(points[1, :num_of_points_per_line, :], expected_points + 10.)
def _replace_tensor_values(_data_to_edit, _new_value, _replace_idx): #function performs _data_to_edit[_replace_idx] = _new_value #acting as a work around for eager tensors if (tf.size(_replace_idx) == 1): _replace_value = _new_value else: _replace_value = tf.repeat(_new_value, tf.size(_replace_idx)) #find the location of index's to keep constant _keep_idx = tf.where( tf.equal( tf.constant([0]), tf.cast(tf.scatter_nd( _replace_idx, tf.ones(tf.size(_replace_idx)), tf.expand_dims(tf.size(_data_to_edit), axis=0)), dtype=tf.int32))) _keep_value = tf.gather_nd(_data_to_edit, _keep_idx) return tf.dynamic_stitch([ tf.cast(tf.squeeze(_replace_idx), dtype=tf.int32), tf.cast(tf.squeeze(_keep_idx), dtype=tf.int32) ], [_replace_value, _keep_value])
def random_walk_layer(B, adj_list, degree_matrix, num_steps, num_sims): """Computes omega using random walks Omega has shape (batch size, k, num sims) """ P = tf.tile(B, (num_sims, )) Ps = [P] for i in range(num_steps): x = tf.random.uniform((num_sims, )) x = tf.repeat(x, B.shape[0]) degrees = tf.gather(degree_matrix, P) neighbor_indices = tf.cast(tf.floor(x * tf.cast(degrees, 'float32')), 'int32') adj_indices = tf.transpose(tf.stack([P, neighbor_indices])) P_ = tf.gather_nd(adj_list, adj_indices) Ps.append(P_) P = P_ return tf.reshape(tf.transpose(tf.stack(Ps)), (B.shape[0], num_steps + 1, num_sims))
def function_with_pending_points(x: TensorType) -> TensorType: # stuff below is quite tricky, and thus deserves an elaborate comment # we receive unknown number N of batches to evaluate # and need to collect batch of B new points # so the shape of `x` is [N, B, D] # we want to add P pending points to each batch # so that acquisition actually receives N batches of shape [P+B, D] each # therefore here we prepend each batch with all pending points # resulting a shape [N, P+B, D] # we do that by repeating pending points N times and concatenating with x # pending points are 2D, we need 3D and repeat along first axis expanded = tf.expand_dims(pending_points, axis=0) pending_points_repeated = tf.repeat(expanded, [tf.shape(x)[0]], axis=0) all_points = tf.concat([pending_points_repeated, x], axis=1) return cast(AcquisitionFunction, self._acquisition_function)(all_points)
def rbf_neighborhood(datacube): datacube_reshape = tf.reshape(datacube, (datacube.shape[0], datacube.shape[1] * datacube.shape[2], datacube.shape[3])) centers = datacube[:, int((datacube.shape[1] - 1) / 2), int((datacube.shape[2] - 1) / 2), :] centers_repmat = tf.repeat(centers[:, tf.newaxis, :], datacube_reshape.shape[1], axis=1) dists = tf.math.pow( tf.math.reduce_euclidean_norm(datacube_reshape - centers_repmat, axis=2), 2) betas = (1 / (tf.math.reduce_mean(dists)))**1 betas = matlib.repmat(betas, dists.shape[0], dists.shape[1]) return np.squeeze(betas)
def test_DotProductAttention(self): """ query and key are vectors of same length Returns: """ query_size = key_size = 2 queries, keys = tf.random.normal(shape=(2, 1, query_size)), tf.ones( (2, 10, key_size)) # The two value matrices in the `values` minibatch are identical values = tf.repeat(tf.reshape(tf.range(40, dtype=tf.float32), shape=(1, 10, 4)), repeats=2, axis=0) valid_lens = tf.constant([2, 6]) attention = DotProductAttention( # num_hiddens=8, dropout=0.1) output = attention(queries, keys, values, valid_lens, training=False) # output.shape self.assertTrue(output.shape == [2, 1, 4])
def costs_decorated_to_logits(cls, symbol_costs_decorated, questions): symbol_costs_valid = tf.ragged.boolean_mask( symbol_costs_decorated['costs'], symbol_costs_decorated['valid']) questions_valid = tf.ragged.boolean_mask( questions, symbol_costs_decorated['valid']) logits_valid = cls.costs_to_logits(symbol_costs_valid, questions_valid) index_mask = tf.repeat(symbol_costs_decorated['valid'], questions.row_lengths()) indices_all = tf.range(questions.row_splits[-1]) indices_valid = indices_all[index_mask] indices_valid = tf.expand_dims(indices_valid, axis=1) # tf.debugging.assert_equal(indices_valid.shape[:-1], logits_valid.flat_values.shape) # We assign the value 0 to the invalid logits. `scatter_nd` defaults the output values to 0. logits_values = tf.scatter_nd( indices_valid, logits_valid.flat_values, tf.reshape(questions.row_splits[-1], (1, ))) logits = tf.RaggedTensor.from_row_splits(logits_values, questions.row_splits) tf.debugging.assert_equal(logits.row_splits, questions.row_splits) return logits
def propagating(self, x, mask=None, pos=None, past=None, training=False, softmax=True): p1, p2 = (tf.unstack(past, axis=1) if past is not None else [None] * self.param['n_layer']), [] t1 = pos if pos is not None else tf.repeat( [past.shape[3]], x.shape[0], 0) if past is not None else None x1 = self.embedding.propagating(x, None, t1, training) for i1 in range(self.param['n_layer']): x1, a1 = self.encoder[i1].propagating(x1, mask, p1[i1], training) p2.append(a1) x1, h1 = self.norm(x1), tf.stack(p2, 1) x2 = tf.matmul(x1, self.embedding.emb, transpose_b=True) return tf.nn.softmax(x2) if softmax else x2, x1, h1
def add_selfloops_edge(edge_index, num_nodes, edge_weight=None, fill_weight=1.0): diagnal_edge_index = tf.reshape( tf.repeat(tf.range(num_nodes, dtype=intx()), 2), [num_nodes, 2]) updated_edge_index = tf.concat([edge_index, diagnal_edge_index], axis=0) if edge_weight: diagnal_edge_weight = tf.cast(tf.fill([num_nodes], fill_weight), dtype=floatx()) updated_edge_weight = tf.concat([edge_weight, diagnal_edge_weight], axis=0) else: updated_edge_weight = None return updated_edge_index, updated_edge_weight
def add_self_loops_indices(indices, n_nodes=None): """ Given the indices of a square SparseTensor, adds the diagonal entries (i, i) and returns the reordered indices. :param indices: Tensor of rank 2, the indices to a SparseTensor. :param n_nodes: the size of the n_nodes x n_nodes SparseTensor indexed by the indices. If `None`, n_nodes is calculated as the maximum entry in the indices plus 1. :return: Tensor of rank 2, the indices to a SparseTensor. """ n_nodes = tf.reduce_max(indices) + 1 if n_nodes is None else n_nodes row, col = indices[..., 0], indices[..., 1] mask = tf.ensure_shape(row != col, row.shape) sl_indices = tf.range(n_nodes, dtype=row.dtype)[:, None] sl_indices = tf.repeat(sl_indices, 2, -1) indices = tf.concat((indices[mask], sl_indices), 0) dummy_values = tf.ones_like(indices[:, 0]) indices, _ = gen_sparse_ops.sparse_reorder(indices, dummy_values, (n_nodes, n_nodes)) return indices
def get_vector_elements_equalities_matrix_op(vector_op): """ Given a vector_op, return a square matrix such that element (i,j) is 1 if vector_op[i] == vector_op[j] and 0 otherwise :param vector_op: 1D tensor of ints :return: 2D matrix of ints """ elements_count_op = tf.shape(vector_op)[0] # Unroll vector so that each element can be matched with each other element vector_repeated_elements_wise = tf.repeat(vector_op, repeats=elements_count_op) vector_repeated_vector_wise = tf.tile(vector_op, multiples=[elements_count_op]) # Compute equalities, cast booleans to floats equalities_vector_op = tf.cast(vector_repeated_elements_wise == vector_repeated_vector_wise, tf.float32) # Reshape vector to square matrix return tf.reshape(equalities_vector_op, shape=(elements_count_op, elements_count_op))
def call(self, x): batch_size = tf.shape(x)[0] # Expand the class token batch number of times. class_token = tf.repeat(self.cls, repeats=batch_size, axis=0) # Concat the input with the trainable class token. x = tf.concat([class_token, x], axis=1) # Apply attention to x. x = self.layer_norm1(x) x, viz_weights = self.attention(query=x[:, 0:1], key=x, value=x, return_attention_scores=True) class_token = class_token + x class_token = self.layer_norm2(class_token) class_token = self.flatten(class_token) class_token = self.layer_norm3(class_token) class_token = class_token + self.mlp(class_token) # Build the logits logits = self.dense(class_token) return logits, tf.squeeze(viz_weights)[..., 1:]
def histmatch(x, y, nbins=256): x_shape = tf.shape(x) y_shape = tf.shape(y) ch = x_shape[-1] size = x_shape[0] * x_shape[1] x_flat = tf.transpose(tf.reshape(x, [-1, ch]), [1, 0]) y_flat = tf.transpose(tf.reshape(y, [-1, ch]), [1, 0]) x_min = tf.reduce_min(x_flat, -1)[..., None] x_max = tf.reduce_max(x_flat, -1)[..., None] y_min = tf.reduce_min(y_flat, -1)[..., None] y_max = tf.reduce_max(y_flat, -1)[..., None] x_flat = (x_flat - x_min) / (x_max - x_min) y_flat = (y_flat - y_min) / (y_max - y_min) x_bins = tf.minimum(tf.cast(x_flat * tf.cast(nbins, tf.float32), tf.int32), nbins - 1) x_hist = tf.math.bincount(x_bins, minlength=nbins, axis=-1) x_cdf = tf.cumsum(x_hist, -1) x_cdf = tf.cast(x_cdf, tf.float32) x_cdf = x_cdf / x_cdf[:, -1, None] y_bins = tf.minimum(tf.cast(y_flat * tf.cast(nbins, tf.float32), tf.int32), nbins - 1) y_hist = tf.math.bincount(y_bins, minlength=nbins, axis=-1) y_cdf = tf.cumsum(y_hist, -1) y_cdf = tf.cast(y_cdf, tf.float32) y_cdf = y_cdf / y_cdf[:, -1, None] x_match = x_cdf[:, :, None] y_match = y_cdf[:, None, :] where = (y_match >= x_match) amax = tf.argmax(where, -1) indices = tf.stack([tf.repeat(tf.range(ch)[..., None], size, -1), x_bins], -1) corr = tf.gather_nd(amax, indices) corr = tf.cast(corr, tf.float32) / tf.cast(nbins, tf.float32) corr = corr * (y_max - y_min) + y_min corr = tf.reshape(tf.transpose(corr, [1, 0]), x_shape) return corr
def get_delta_errors_fixed_size(y_true, all_but_one_auxiliary_outputs, error_with_all_features, loss_function, log_transform, math_ops): n_out = math_ops.int_shape(all_but_one_auxiliary_outputs[0])[1] n_feats = len(all_but_one_auxiliary_outputs) x = tf.stack(all_but_one_auxiliary_outputs, axis=1) #(batch, n_feats, n_out) y_true_tile = tf.reshape(tf.repeat(y_true, n_feats, axis=1), (-1, n_feats, n_out)) # (batch, n_feats, n_out) error_without_one_feature = safe_evaluate_custom_loss_function( loss_function, y_true_tile, x, math_ops) # (batch, n_feats, n_out) error_with_all_features = tf.reshape(error_with_all_features, (-1, n_out)) # (batch, n_out) delta_errors = math_ops.maximum( error_without_one_feature - error_with_all_features, math_ops.epsilon()) # delta_errors = error_without_one_feature - error_with_all_features # delta_errors = delta_errors - math_ops.reduce_min(delta_errors, axis=1) + math_ops.epsilon() if log_transform: delta_errors = math_ops.log(1 + delta_errors) # --- # delta_errors = [] # for all_but_one_auxiliary_output in all_but_one_auxiliary_outputs: # error_without_one_feature = safe_evaluate_custom_loss_function( # loss_function, y_true, all_but_one_auxiliary_output, math_ops # ) # # The error without the feature is an indicator as to how potent the left-out feature is as a predictor. # delta_error = math_ops.maximum(error_without_one_feature - error_with_all_features, math_ops.epsilon()) # if log_transform: # delta_error = math_ops.log(1 + delta_error) # delta_errors.append(delta_error) # delta_errors = math_ops.stack(delta_errors, axis=-1) return delta_errors
def transform(image, inv_mat, image_shape): '''Used in random_rotate Adapted from https://www.kaggle.com/xiejialun/gridmask-data-augmentation-with-tensorflow/notebook''' h, w, c = image_shape cx, cy = w // 2, h // 2 new_xs = tf.repeat(tf.range(-cx, cx, 1), h) new_ys = tf.tile(tf.range(-cy, cy, 1), [w]) new_zs = tf.ones([h * w], dtype=tf.int32) old_coords = tf.matmul( inv_mat, tf.cast(tf.stack([new_xs, new_ys, new_zs]), tf.float32)) old_coords_x, old_coords_y = tf.round(old_coords[0, :] + w // 2), tf.round(old_coords[1, :] + h // 2) clip_mask_x = tf.logical_or(old_coords_x < 0, old_coords_x > w - 1) clip_mask_y = tf.logical_or(old_coords_y < 0, old_coords_y > h - 1) clip_mask = tf.logical_or(clip_mask_x, clip_mask_y) old_coords_x = tf.boolean_mask(old_coords_x, tf.logical_not(clip_mask)) old_coords_y = tf.boolean_mask(old_coords_y, tf.logical_not(clip_mask)) new_coords_x = tf.boolean_mask(new_xs + cx, tf.logical_not(clip_mask)) new_coords_y = tf.boolean_mask(new_ys + cy, tf.logical_not(clip_mask)) old_coords = tf.cast(tf.stack([old_coords_y, old_coords_x]), tf.int32) new_coords = tf.cast(tf.stack([new_coords_y, new_coords_x]), tf.int64) rotated_image_values = tf.gather_nd(image, tf.transpose(old_coords)) rotated_image_channel = list() for i in range(c): vals = rotated_image_values[:, i] sparse_channel = tf.SparseTensor(tf.transpose(new_coords), vals, [h, w]) rotated_image_channel.append( tf.sparse.to_dense(sparse_channel, default_value=0, validate_indices=False)) return tf.transpose(tf.stack(rotated_image_channel), [1, 2, 0])
def _train_step(self, times, sequences, optimizer): with tf.GradientTape() as tape: params = self.encoder((times, sequences), training=True) times_list = tf.expand_dims(times.values, axis=-1) params_list = tf.repeat(params, times.row_lengths(), axis=0) latents = self.trajectory((times_list, params_list)) reconstructed_sequences_seq = self.decoder(latents, training=True) reconstructed_sequences = tf.RaggedTensor.from_row_splits( reconstructed_sequences_seq, row_splits=times.row_splits) loss, reconstruction_error, regularisation = self._loss( sequences, reconstructed_sequences, params) heart_rates = tf.exp(params[:, 0]) * 60 # update model weights variables = self.trainable_variables grads = tape.gradient(loss, variables) optimizer.apply_gradients(zip(grads, variables)) # logging self._train_metrics['loss'](loss) self._train_metrics['reconstruction_error'](reconstruction_error) self._train_metrics['regularisation'](regularisation) self._train_metrics['heart_rate'](heart_rates) self._train_metrics['regularisation_strength'](self._sigma_squared) # reconstruction_stddev = tf.reduce_mean(tf.math.reduce_std(reconstructed_sequences, axis=1, keepdims=False).values) means = tf.expand_dims(tf.reduce_mean(reconstructed_sequences, axis=1), axis=1) reconstruction_var = tf.reduce_mean( tf.reduce_mean((reconstructed_sequences - means)**2, axis=1)) reconstruction_stddev = tf.sqrt(reconstruction_var) self._train_metrics['reconstruction_stddev'](reconstruction_stddev) # update regularisation weight update_rate = 0.99 self._sigma_squared.assign(update_rate * self._sigma_squared + (1 - update_rate) * tf.reduce_mean(reconstruction_error))
def rotate_img(image, dim, n_channels, degrees): ''' Rotates an image by `degrees` degrees Args: image: array of one squared image of size (dim, dim, n_channels) n_channels: (int) degrees: (int or float) the angle of the rotation in degrees Returns: An image rotated by an angle of `degrees` ''' xdim = dim%2 degrees = degrees * tf.ones([1], dtype='float32') rotation = 3.141593 * degrees / 180. def get_3x3_mat(lst): return tf.reshape(tf.concat([lst],axis=0), [3,3]) c1 = tf.math.cos(rotation) s1 = tf.math.sin(rotation) one = tf.constant([1],dtype='float32') zero = tf.constant([0],dtype='float32') m = get_3x3_mat([c1, s1, zero, -s1, c1, zero, zero, zero, one]) # list destination pixel indices x = tf.repeat(tf.range(dim//2, -dim//2,-1), dim) y = tf.tile(tf.range(-dim//2, dim//2), [dim]) z = tf.ones([dim*dim], dtype='int32') idx = tf.stack( [x,y,z] ) # rotate destination pixels onto origin pixels idx2 = K.dot(m, tf.cast(idx, dtype='float32')) idx2 = K.cast(idx2, dtype='int32') idx2 = K.clip(idx2, -dim//2+xdim+1, dim//2) # find origin pixel values idx3 = tf.stack([dim//2-idx2[0,], dim//2-1+idx2[1,]]) d = tf.gather_nd(image, tf.transpose(idx3)) return tf.reshape(d,[dim, dim, n_channels])
def cem_planner(state, num_actions, horizon, proposals, topk, iterations, imagine, objective): dtype = prec.global_policy().compute_dtype B, P = list(state.values())[0].shape[0], proposals H, A = horizon, num_actions flat_state = {k: tf.repeat(v, P, 0) for k, v in state.items()} mean = tf.zeros((B, H, A), dtype) std = tf.ones((B, H, A), dtype) for _ in range(iterations): proposals = tf.random.normal((B, P, H, A), dtype=dtype) proposals = proposals * std[:, None] + mean[:, None] proposals = tf.clip_by_value(proposals, -1, 1) flat_proposals = tf.reshape(proposals, (B * P, H, A)) states = imagine(flat_proposals, flat_state) scores = objective(states) scores = tf.reshape(tf.reduce_sum(scores, -1), (B, P)) _, indices = tf.math.top_k(scores, topk, sorted=False) best = tf.gather(proposals, indices, axis=1, batch_dims=1) mean, var = tf.nn.moments(best, 1) std = tf.sqrt(var + 1e-6) return mean[:, 0, :]
def add_noise(audio, noises=None, scale=0.5): if noises is not None: # Create a random tensor of the same size as audio ranging from # 0 to the number of noise stream samples that we have. tf_rnd = tf.random.uniform((tf.shape(audio)[0], ), 0, noises.shape[0], dtype=tf.int32) noise = tf.gather(noises, tf_rnd, axis=0) # Get the amplitude proportion between the audio and the noise prop = tf.math.reduce_max(audio, axis=1) / tf.math.reduce_max( noise, axis=1) prop = tf.repeat(tf.expand_dims(prop, axis=1), tf.shape(audio)[1], axis=1) # Adding the rescaled noise to audio audio = audio + noise * prop * scale return audio
def body(i, batch_size, outputs, encoder_masks, encoder_hidden_states, durations_gt, max_durations): repeats = durations_gt[i] real_length = tf.reduce_sum(repeats) pad_size = max_durations - real_length masks = tf.sequence_mask([real_length], max_durations, dtype=tf.int32) repeat_encoder_hidden_states = tf.repeat(encoder_hidden_states[i], repeats=repeats, axis=0) repeat_encoder_hidden_states = tf.expand_dims( tf.pad(repeat_encoder_hidden_states, [[0, pad_size], [0, 0]]), 0) # [1, max_durations, hidden_size] outputs = tf.concat([outputs, repeat_encoder_hidden_states], axis=0) encoder_masks = tf.concat([encoder_masks, masks], axis=0) return [ i + 1, batch_size, outputs, encoder_masks, encoder_hidden_states, durations_gt, max_durations ]
def broadcast_wff_and_mask(wff: Formula, mask: Formula) -> Formula: """Broadcast the wff to include all vars in mask; put the vars of the mask in the first axes""" wff = wff._copy() # 1. broadcast wff with vars that are in the mask but not yet in the formula mask_vars_not_in_wff = [ var for var in mask.free_vars if var not in wff.free_vars ] for var in mask_vars_not_in_wff: new_idx = len(wff.free_vars) wff.tensor = tf.expand_dims(wff.tensor, axis=new_idx) wff.tensor = tf.repeat(wff.tensor, mask._get_dim_of_free_var(var), axis=new_idx) wff.free_vars.append(var) # 2. transpose wff so that the masked vars on the first axes vars_not_in_mask = [ var for var in wff.free_vars if var not in mask.free_vars ] wff = transpose_free_vars(wff, new_var_order=mask.free_vars + vars_not_in_mask) return wff
def init_qstar_mean(self, m, batch_index, batch_num): """Initialize the q0 with mean.""" # use q0=avg(m) (or q0=0) # batch_shape = ksb.shape(batch_num) q = tf.math.segment_mean(m, batch_index) # (batch,feat) # r0 qt = tf.repeat(q, batch_num, axis=0) # (batch*num,feat) et = self.f_et(m, qt) # (batch*num,) # get at = exp(et)/sum(et) with sum(et)=norm at = ksb.exp(et - self.get_scale_per_sample(et, batch_index, batch_num)) # (batch*num,) norm = self.get_norm(at, batch_index, batch_num) # (batch*num,) at = norm * at # (batch*num,) x (batch*num,) # calculate rt at = ksb.expand_dims(at, axis=1) # (batch*num,1) rt = m * at # (batch*num,feat) x (batch*num,1) rt = tf.math.segment_sum(rt, batch_index) # (batch,feat) # [q0,r0] qstar = ksb.concatenate([q, rt], axis=1) # (batch,2*feat) qstar = ksb.expand_dims(qstar, axis=1) # (batch,1,2*feat) return qstar