RNN_OriginalFedAvg

class fl_sim.models.RNN_OriginalFedAvg(embedding_dim: int = 8, vocab_size: int = 90, hidden_size: int = 256)

Bases: 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.

Parameters:

• 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

[1]

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficient Learning of Deep Networks from Decentralized Data. In Artificial Intelligence and Statistics, 1273–1282. PMLR, 2017.

[2]

Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive Federated Optimization. In International Conference on Learning Representations. 2021.

forward(input_seq: Tensor) → Tensor

Forward pass.

Parameters:

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

Returns:

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

Return type:

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.

Parameters:

• 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.

Returns:

The predicted next characters.

Return type:

str