RNN_OriginalFedAvg

RNN_OriginalFedAvg#

class fl_sim.models.RNN_OriginalFedAvg(embedding_dim: int = 8, vocab_size: int = 90, hidden_size: int = 256)[源代码]#

基类：Module, CLFMixin, SizeMixin, DiffMixin

Creates a RNN model using LSTM layers for Shakespeare language models (next character prediction task).

This replicates the model structure in the paper [McMahan et al.[1]]. This is also recommended model by [Reddi et al.[2]].

Modified from FedML.

参数:

  • embedding_dim (int) -- The size of each embedding vector.

  • vocab_size (int) -- The number of different characters that can appear in the input.

  • hidden_size (int) -- The number of features in the hidden state h.

References

forward(input_seq: Tensor) → Tensor[源代码]#

Forward pass.

参数:

input_seq (torch.Tensor) -- Shape (batch_size, seq_len), dtype torch.long.

返回:

Shape (batch_size, vocab_size, seq_len), dtype torch.float32.

返回类型:

torch.Tensor

pipeline(truncated_sentence: str, char_to_id: Callable[[str], int] | Dict[str, int] | None = None, id_to_char: Callable[[int], str] | Dict[int, str] | None = None) → str[源代码]#

Predict the next character given a truncated sentence.

参数:

  • truncated_sentence (str) -- The truncated sentence.

  • char_to_id (Callable[[str], int] or Dict[str, int], optional) -- A function that maps a character to its id.

  • id_to_char (Callable[[int], str] or Dict[int, str], optional) -- A function that maps an id to its character.

返回:

The predicted next characters.

返回类型:

str

目录