code: """
  Copyright 2021 Keisuke Izumiya
  
  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 torch
  import torch.distributed as distd
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  
  import mmle.distributed as mdistd
  import mmle.nn as mnn
  from mmle.utils.manager import Manager
  
  
  BATCH_SIZE = 400
  DATA_NUM = 40000
  DATA_DIM = 10
  EPOCH_NUM = 50
  LABEL_DIM = 5
  LOG_DIR = "log"
  MIDDLE_DIM = 16
  
  
  class Dataset(torch.utils.data.Dataset):
      def __init__(self):
          super().__init__()
  
          self.data = torch.randn(DATA_NUM, DATA_DIM)
          self.label = torch.randn(DATA_NUM, LABEL_DIM)
  
      def __len__(self):
          return DATA_NUM
  
      def __getitem__(self, idx):
          return self.data[idx], self.label[idx]
  
  
  def main():
      world_size = distd.get_world_size()
      rank = distd.get_rank()
  
      model = nn.Sequential(
          mnn.FC(DATA_DIM, MIDDLE_DIM), mnn.FC(MIDDLE_DIM, LABEL_DIM, bn=False, activ="id")
      )
      model = mdistd.distribute_module(model)
  
      optimizer = optim.Adam(model.parameters())
  
      dataset = Dataset()
      sampler = DistributedSampler(dataset, world_size, rank)
      loader = DataLoader(dataset, BATCH_SIZE, sampler=sampler, drop_last=True)
  
      if rank == 0:
          manager = Manager(LOG_DIR, use_tensorboard=True)
          manager.tensorboard_writer.add_graph(model, dataset[0][0].repeat(BATCH_SIZE, 1))
          step = 0
  
      for epoch in range(EPOCH_NUM):
          model.train()
          for data, label in loader:
              loss = F.mse_loss(model(data), label.to(rank))
  
              mnn.zero_grad(model)
              loss.backward()
              optimizer.step()
              distd.barrier()
  
              if rank == 0:
                  step += world_size
                  manager.tensorboard_writer.plot("loss", "train", loss.item(), step)
  
          if rank == 0:
              print(f"Finish epoch {epoch}: loss={loss.item():.3f}")
          distd.barrier()
  
      if rank == 0:
          manager.save({"model": model.state_dict()}, "model")
          manager.close()
  
  
  if __name__ == "__main__":
      mdistd.spawn(main, nprocs=torch.cuda.device_count())
  
apis: torch.utils.data.DataLoader torch.distributed.get_rank torch.distributed.get_world_size torch.utils.data.distributed.DistributedSampler torch.randn torch.cuda.device_count torch.distributed.barrier