crested.tl.zoo.borzoi_prime

Contents

crested.tl.zoo.borzoi_prime#

crested.tl.zoo.borzoi_prime(seq_len, num_classes, num_conv_blocks=6, num_transformer_blocks=8, num_transformer_heads=8, target_length=12288, start_filters=512, tower_start_filters=None, filters=1536, pointwise_filters=None, unet_connections=[4, 5, 6], unet_filters=1536, upsample_conv=True, conv_activation='gelu_approx', transformer_activation='relu', output_activation='softplus', pool_type='max', first_kernel_size=15, kernel_size=5, transformer_dropout=0.2, pointwise_dropout=0.1, bn_sync=False, name='Borzoi_prime')#

Construct an fully replicated Borzoi Prime model.

Borzoi Prime is effectively the Borzoi architecture, except with one more upsampling layer, no pointwise filters and a larger target length to account for the higher resolution.

Note that unlike other CREsted model zoo architectures, this architecture is not suited for predicting individual regions out of the box.

Parameters:

  • seq_len (int) – Width of the input region. Borzoi default is 524288.

  • num_classes (int | Sequence[int]) – Number of classes to predict. If an int, creates a single head with num_classes classes. If a list of integers, creates multiple heads in a list.

  • num_conv_blocks (int (default: 6)) – Number of convolution blocks to include in the tower, after the stem block.

  • num_transformer_blocks (int (default: 8)) – Number of transformer blocks to include in the transformer stack.

  • target_length (int (default: 12288)) – The target length in bins to crop to. Default is 12288, cropping away 10240 bins (164kb) on each side.

  • start_filters (int (default: 512)) – Starting number of filters for the first DNA-facing block, exponentially increasing towards filters through the conv tower.

  • tower_start_filters (int | None (default: None)) – Optional: Different number of filters to start the conv tower with after the stem conv. By default, inferred starting from start_filters to filters.

  • filters (int (default: 1536)) – Number of filters at the end of the conv tower and in the upsampling.

  • pointwise_filters (int | None (default: None)) – Number of filters of the post-transformer/upsampling final pointwise convolution. If None, block is not included.

  • unet_connections (Sequence[int] (default: [4, 5, 6])) – Levels in the convolution tower to add U-net skip connections past the transformer tower. 1-indexed, so [5, 6] means after the 5th and 6th block.

  • unet_filters (int (default: 1536)) – Number of filters to use for the U-net connection skip blocks.

  • upsample_conv (bool (default: True)) – Whether to include an extra convolution during the upsampling. Used in release Borzoi, removed later.

  • conv_activation (str (default: 'gelu_approx')) – Activation function to use in the conv tower, in the upsampling, and in the final pointwise block.

  • transformer_activation (str (default: 'relu')) – Activation function to use in the feedforward section of the transformer blocks.

  • output_activation (str (default: 'softplus')) – Final activation to use on the output heads, just before predicting the tracks.

  • pool_type (str (default: 'max')) – What kind of pooling type to use, one of ‘max’ or ‘attention’.

  • first_kernel_size (int (default: 15)) – Kernel size of the first conv layer, directly interfacing the sequence.

  • kernel_size (int (default: 5)) – Kernel size of the convolutions in the conv tower.

  • transformer_dropout (default: 0.2) – Dropout rate used in the transformer blocks, both MHA and feed-forward.

  • pointwise_dropout (float (default: 0.1)) – Dropout rate of the post-transformer final pointwise layer.

  • bn_sync (bool (default: False)) – Whether to use synchronized cross-GPU BatchNormalisations. Default is False to preserve TensorFlow-PyTorch compatibility.

Returns:

A Keras model.

Contents