mindflow.cell.AttentionBlock

class mindflow.cell.AttentionBlock(in_channels, num_heads, drop_mode='dropout', dropout_rate=0.0, compute_dtype=mstype.float32)

AttentionBlock comprises an MultiHeadAttention and an MLP layer.

Parameters

• in_channels (int) – The input channels.

• num_heads (int) – The number of attention heads.

• drop_mode (str) – Dropout method. Default: dropout. Support dropout or droppath.

• dropout_rate (float) – The drop rate of dropout layer, greater than 0 and less equal than 1. Default: 0.0.

• compute_dtype (mindspore.dtype) – Compute dtype. Default: mstype.float32, indicates mindspore.float32.

Inputs:

• x (Tensor) - Tensor with shape \((batch\_size, sequence\_len, in\_channels)\).

• mask (Tensor) - Tensor with shape \((batch\_size, sequence\_len, sequence\_len)\) or \((sequence\_len, sequence\_len)\) or \((batch\_size, num_heads, sequence\_len, sequence\_len)\).

Outputs:

• output (Tensor) - Tensor with shape \((batch\_size, sequence\_len, in\_channels)\).

Supported Platforms:

Ascend CPU

Examples

>>> from mindspore import ops
>>> from mindflow.cell import AttentionBlock
>>> model = AttentionBlock(in_channels=256, num_heads=4)
>>> x = ops.rand((4, 100, 256))
>>> output = model(x)
>>> print(output.shape)
(4, 100, 256)