code: #!/usr/bin/env python3
  # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import argparse
  from collections import Counter
  from pathlib import Path
  from statistics import mean
  
  import torch
  import torch.nn as nn
  from opacus import PrivacyEngine
  from opacus.layers import DPGRU, DPLSTM, DPRNN
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import DataLoader, Dataset
  from tqdm import tqdm
  
  
  parser = argparse.ArgumentParser(
      description="PyTorch Name language classification DP Training",
      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  )
  parser.add_argument(
      "--data-root",
      required=True,
      type=str,
      help="Path to training set of names (ie. ~/data/names/)",
  )
  parser.add_argument(
      "--device",
      type=str,
      default="cuda",
      help="GPU ID for this process",
  )
  parser.add_argument(
      "-b",
      "--batch-size",
      default=800,
      type=int,
      metavar="N",
      help="mini-batch size",
  )
  parser.add_argument(
      "--mode",
      default="lstm",
      choices=["lstm", "gru", "rnn"],
      help="recursive network type",
  )
  parser.add_argument(
      "--embedding-size", default=64, type=int, help="Character embedding dimension"
  )
  parser.add_argument(
      "--hidden-size", default=128, type=int, help="hidden state dimensions"
  )
  parser.add_argument("--n-layers", default=1, type=int, help="How many layers to use")
  parser.add_argument(
      "--test-every",
      default=0,
      type=int,
      help="Run evaluation on the test every these many epochs",
  )
  parser.add_argument(
      "--bidirectional",
      action="store_true",
      default=False,
      help="If turned on, makes the RNN bidirectional",
  )
  parser.add_argument(
      "--learning-rate",
      default=2.0,
      type=float,
      metavar="LR",
      help="initial learning rate",
  )
  parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
  parser.add_argument(
      "--train-split",
      type=float,
      default=0.8,
      help="Fraction of data to utilize for training (rest for evaluation)",
  )
  parser.add_argument(
      "--sigma",
      type=float,
      default=1.0,
      metavar="S",
      help="Noise multiplier",
  )
  parser.add_argument(
      "-c",
      "--max-per-sample-grad-norm",
      type=float,
      default=1.5,
      metavar="C",
      help="Clip per-sample gradients to this norm",
  )
  parser.add_argument(
      "--disable-dp",
      action="store_true",
      default=False,
      help="Disable privacy training and just train with vanilla SGD",
  )
  parser.add_argument(
      "--secure-rng",
      action="store_true",
      default=False,
      help="Enable Secure RNG to have trustworthy privacy guarantees. Comes at a performance cost",
  )
  parser.add_argument(
      "--delta",
      type=float,
      default=8e-5,
      metavar="D",
      help="Target delta",
  )
  parser.add_argument(
      "--print-every",
      type=int,
      default=5,
      help="Print the evaluation accuracy every these many iterations",
  )
  
  
  class CharByteEncoder(nn.Module):
      """
      This encoder takes a UTF-8 string and encodes its bytes into a Tensor. It can also
      perform the opposite operation to check a result.
  
      Examples:
  
      >>> encoder = CharByteEncoder()
      >>> t = encoder('Ślusàrski')  # returns tensor([256, 197, 154, 108, 117, 115, 195, 160, 114, 115, 107, 105, 257])
      >>> encoder.decode(t)  # returns "<s>Ślusàrski</s>"
      """
  
      def __init__(self):
          super().__init__()
          self.start_token = "<s>"
          self.end_token = "</s>"
          self.pad_token = "<pad>"
  
          self.start_idx = 256
          self.end_idx = 257
          self.pad_idx = 258
  
      def forward(self, s: str, pad_to=0) -> torch.LongTensor:
          """
          Encodes a string. It will append a start token <s> (id=self.start_idx) and an end token </s>
          (id=self.end_idx).
  
          Args:
              s: The string to encode.
              pad_to: If not zero, pad by appending self.pad_idx until string is of length `pad_to`.
                  Defaults to 0.
  
          Returns:
              The encoded LongTensor of indices.
          """
          encoded = s.encode()
          n_pad = pad_to - len(encoded) if pad_to > len(encoded) else 0
          return torch.LongTensor(
              [self.start_idx]
              + [c for c in encoded]  # noqa
              + [self.end_idx]
              + [self.pad_idx for _ in range(n_pad)]
          )
  
      def decode(self, char_ids_tensor: torch.LongTensor) -> str:
          """
          The inverse of `forward`. Keeps the start, end and pad indices.
          """
          char_ids = char_ids_tensor.cpu().detach().tolist()
  
          out = []
          buf = []
          for c in char_ids:
              if c < 256:
                  buf.append(c)
              else:
                  if buf:
                      out.append(bytes(buf).decode())
                      buf = []
                  if c == self.start_idx:
                      out.append(self.start_token)
                  elif c == self.end_idx:
                      out.append(self.end_token)
                  elif c == self.pad_idx:
                      out.append(self.pad_token)
  
          if buf:  # in case some are left
              out.append(bytes(buf).decode())
          return "".join(out)
  
      def __len__(self):
          """
          The length of our encoder space. This is fixed to 256 (one byte) + 3 special chars
          (start, end, pad).
  
          Returns:
              259
          """
          return 259
  
  
  class NamesDataset(Dataset):
      def __init__(self, root):
          self.root = Path(root)
  
          self.labels = list({langfile.stem for langfile in self.root.iterdir()})
          self.labels_dict = {label: i for i, label in enumerate(self.labels)}
          self.encoder = CharByteEncoder()
          self.samples = self.construct_samples()
  
      def __getitem__(self, i):
          return self.samples[i]
  
      def __len__(self):
          return len(self.samples)
  
      def construct_samples(self):
          samples = []
          for langfile in self.root.iterdir():
              label_name = langfile.stem
              label_id = self.labels_dict[label_name]
              with open(langfile, "r") as fin:
                  for row in fin:
                      samples.append(
                          (self.encoder(row.strip()), torch.tensor(label_id).long())
                      )
          return samples
  
      def label_count(self):
          cnt = Counter()
          for _x, y in self.samples:
              label = self.labels[int(y)]
              cnt[label] += 1
          return cnt
  
  
  VOCAB_SIZE = 256 + 3  # 256 alternatives in one byte, plus 3 special characters.
  
  
  class CharNNClassifier(nn.Module):
      def __init__(
          self,
          rnn_type,
          embedding_size,
          hidden_size,
          output_size,
          num_layers=1,
          bidirectional=False,
          vocab_size=VOCAB_SIZE,
      ):
          super().__init__()
  
          self.embedding_size = embedding_size
          self.hidden_size = hidden_size
          self.output_size = output_size
          self.vocab_size = vocab_size
  
          self.embedding = nn.Embedding(vocab_size, embedding_size)
          self.rnn = rnn_type(
              embedding_size,
              hidden_size,
              num_layers=num_layers,
              bidirectional=bidirectional,
              batch_first=True,
          )
          self.out_layer = nn.Linear(hidden_size, output_size)
  
      def forward(self, x, hidden=None):
          x = self.embedding(x)  # -> [B, T, D]
          x, _ = self.rnn(x, hidden)  # -> [B, T, H]
          x = x[:, -1, :]  # -> [B, H]
          x = self.out_layer(x)  # -> [B, C]
          return x
  
  
  def padded_collate(batch, padding_idx=0):
      x = pad_sequence(
          [elem[0] for elem in batch], batch_first=True, padding_value=padding_idx
      )
      y = torch.stack([elem[1] for elem in batch]).long()
  
      return x, y
  
  
  def train(
      model,
      criterion,
      optimizer,
      train_loader,
      epoch,
      privacy_engine,
      target_delta,
      device="cuda:0",
  ):
      model.train()
  
      accs = []
      losses = []
      for x, y in tqdm(train_loader):
          x = x.to(device)
          y = y.to(device)
  
          logits = model(x)
          loss = criterion(logits, y)
          loss.backward()
  
          optimizer.step()
          optimizer.zero_grad()
  
          preds = logits.argmax(-1)
          n_correct = float(preds.eq(y).sum())
          batch_accuracy = n_correct / len(y)
  
          accs.append(batch_accuracy)
          losses.append(float(loss))
  
      printstr = (
          f"\t Epoch {epoch}. Accuracy: {mean(accs):.6f} | Loss: {mean(losses):.6f}"
      )
      try:
          epsilon, best_alpha = privacy_engine.accountant.get_privacy_spent(
              delta=target_delta
          )
          printstr += f" | (ε = {epsilon:.2f}, δ = {target_delta}) for α = {best_alpha}"
      except AttributeError:
          pass
      print(printstr)
      return
  
  
  def test(model, test_loader, privacy_engine, target_delta, device="cuda:0"):
      model.eval()
  
      accs = []
      with torch.no_grad():
          for x, y in tqdm(test_loader):
              x = x.to(device)
              y = y.to(device)
  
              preds = model(x).argmax(-1)
              n_correct = float(preds.eq(y).sum())
              batch_accuracy = n_correct / len(y)
  
              accs.append(batch_accuracy)
      mean_acc = mean(accs)
      printstr = "\n----------------------------\n" f"Test Accuracy: {mean_acc:.6f}"
      if privacy_engine:
          epsilon, best_alpha = privacy_engine.accountant.get_privacy_spent(
              delta=target_delta
          )
          printstr += f" (ε = {epsilon:.2f}, δ = {target_delta}) for α = {best_alpha}"
      print(printstr + "\n----------------------------\n")
      return mean_acc
  
  
  def main():
      args = parser.parse_args()
      device = torch.device(args.device)
      ds = NamesDataset(args.data_root)
      train_len = int(args.train_split * len(ds))
      test_len = len(ds) - train_len
  
      print(f"{train_len} samples for training, {test_len} for testing")
  
      if args.secure_rng:
          try:
              import torchcsprng as prng
          except ImportError as e:
              msg = (
                  "To use secure RNG, you must install the torchcsprng package! "
                  "Check out the instructions here: https://github.com/pytorch/csprng#installation"
              )
              raise ImportError(msg) from e
  
          generator = prng.create_random_device_generator("/dev/urandom")
  
      else:
          generator = None
  
      train_ds, test_ds = torch.utils.data.random_split(
          ds, [train_len, test_len], generator=generator
      )
  
      if args.mode == "rnn":
          rnn_type = DPRNN
      elif args.mode == "gru":
          rnn_type = DPGRU
      elif args.mode == "lstm":
          rnn_type = DPLSTM
      else:
          raise ValueError(f"Invalid network type: {args.mode}")
  
      model = CharNNClassifier(
          rnn_type,
          args.embedding_size,
          args.hidden_size,
          len(ds.labels),
          args.n_layers,
          args.bidirectional,
      )
      model = model.to(device)
  
      train_ds, test_ds = torch.utils.data.random_split(
          ds, [train_len, test_len], generator=generator
      )
  
      train_loader = DataLoader(
          train_ds,
          batch_size=args.batch_size,
          num_workers=1,
          pin_memory=True,
          collate_fn=padded_collate,
      )
  
      test_loader = DataLoader(
          test_ds,
          batch_size=2 * args.batch_size,
          shuffle=False,
          num_workers=1,
          pin_memory=True,
          collate_fn=padded_collate,
      )
  
      criterion = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
  
      if not args.disable_dp:
          privacy_engine = PrivacyEngine(secure_mode=args.secure_rng)
          model, optimizer, train_loader = privacy_engine.make_private(
              module=model,
              optimizer=optimizer,
              data_loader=train_loader,
              noise_multiplier=args.sigma,
              max_grad_norm=args.max_per_sample_grad_norm,
          )
      else:
          privacy_engine = None
  
      print(f"Train stats ({args.mode}): \n")
      for epoch in tqdm(range(args.epochs)):
          train(
              model,
              criterion,
              optimizer,
              train_loader,
              epoch,
              privacy_engine,
              args.delta,
              device=device,
          )
          if args.test_every:
              if epoch % args.test_every == 0:
                  test(model, test_loader, privacy_engine, args.delta, device=device)
  
      mean_acc = test(model, test_loader, privacy_engine, args.delta, device=device)
      torch.save(mean_acc, f"run_results_chr_{args.mode}_classification.pt")
  
  
  if __name__ == "__main__":
      main()
  
apis: torch.nn.utils.rnn.pad_sequence torch.utils.data.DataLoader torch.stack torch.nn.Linear torch.save torch.no_grad torch.nn.Embedding torch.nn.CrossEntropyLoss torch.tensor torch.utils.data.random_split torch.device