def linear(x: TensorType, size: int, name: str, initializer: Optional[Any] = None, bias_init: float = 0.0) -> TensorType: w = tf1.get_variable( name + "/w", [x.get_shape()[1], size], initializer=initializer) b = tf1.get_variable( name + "/b", [size], initializer=tf1.constant_initializer(bias_init)) return tf.matmul(x, w) + b
def conv2d( x: TensorType, num_filters: int, name: str, filter_size: Tuple[int, int] = (3, 3), stride: Tuple[int, int] = (1, 1), pad: str = "SAME", dtype: Optional[Any] = None, collections: Optional[Any] = None, ) -> TensorType: if dtype is None: dtype = tf.float32 with tf1.variable_scope(name): stride_shape = [1, stride[0], stride[1], 1] filter_shape = [ filter_size[0], filter_size[1], int(x.get_shape()[3]), num_filters, ] # There are "num input feature maps * filter height * filter width" # inputs to each hidden unit. fan_in = np.prod(filter_shape[:3]) # Each unit in the lower layer receives a gradient from: "num output # feature maps * filter height * filter width" / pooling size. fan_out = np.prod(filter_shape[:2]) * num_filters # Initialize weights with random weights. w_bound = np.sqrt(6 / (fan_in + fan_out)) w = tf1.get_variable( "W", filter_shape, dtype, tf1.random_uniform_initializer(-w_bound, w_bound), collections=collections, ) b = tf1.get_variable( "b", [1, 1, 1, num_filters], initializer=tf1.constant_initializer(0.0), collections=collections, ) return tf1.nn.conv2d(x, w, stride_shape, pad) + b
def add_time_dimension(padded_inputs: TensorType, *, max_seq_len: int, framework: str = "tf", time_major: bool = False): """Adds a time dimension to padded inputs. Args: padded_inputs (TensorType): a padded batch of sequences. That is, for seq_lens=[1, 2, 2], then inputs=[A, *, B, B, C, C], where A, B, C are sequence elements and * denotes padding. max_seq_len (int): The max. sequence length in padded_inputs. framework (str): The framework string ("tf2", "tf", "tfe", "torch"). time_major (bool): Whether data should be returned in time-major (TxB) format or not (BxT). Returns: TensorType: Reshaped tensor of shape [B, T, ...] or [T, B, ...]. """ # Sequence lengths have to be specified for LSTM batch inputs. The # input batch must be padded to the max seq length given here. That is, # batch_size == len(seq_lens) * max(seq_lens) if framework in ["tf2", "tf", "tfe"]: assert time_major is False, "time-major not supported yet for tf!" padded_batch_size = tf.shape(padded_inputs)[0] # Dynamically reshape the padded batch to introduce a time dimension. new_batch_size = padded_batch_size // max_seq_len new_shape = ([new_batch_size, max_seq_len] + padded_inputs.get_shape().as_list()[1:]) return tf.reshape(padded_inputs, new_shape) else: assert framework == "torch", "`framework` must be either tf or torch!" padded_batch_size = padded_inputs.shape[0] # Dynamically reshape the padded batch to introduce a time dimension. new_batch_size = padded_batch_size // max_seq_len if time_major: new_shape = (max_seq_len, new_batch_size) + padded_inputs.shape[1:] else: new_shape = (new_batch_size, max_seq_len) + padded_inputs.shape[1:] return torch.reshape(padded_inputs, new_shape)
def flatten(x: TensorType) -> TensorType: return tf.reshape(x, [-1, np.prod(x.get_shape().as_list()[1:])])