crested.tl.zoo.borzoi

Contents

crested.tl.zoo.borzoi#

crested.tl.zoo.borzoi(seq_len, num_classes, num_conv_blocks=6, num_transformer_blocks=8, num_transformer_heads=8, target_length=6144, start_filters=512, tower_start_filters=None, filters=1536, pointwise_filters=1920, unet_connections=[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')#

Construct an fully replicated Borzoi model.

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: 6144)) – The target length in bins to crop to. Default is 6144, cropping away 5120 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 (Optional[int] (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: 1920)) – Number of filters of the post-transformer/upsampling final pointwise convolution. If None, block is not included.

  • unet_connections (Sequence[int] (default: [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.

Return type:

Model

Returns:

A Keras model.

Contents