code: def main(
      # Architectural hyperparameters. These are quite small for illustrative purposes.
      initial_noise_size=5,  # How many noise dimensions to sample at the start of the SDE.
      noise_size=3,  # How many dimensions the Brownian motion has.
      hidden_size=16,  # How big the hidden size of the generator SDE and the discriminator CDE are.
      mlp_size=16,  # How big the layers in the various MLPs are.
      num_layers=1,  # How many hidden layers to have in the various MLPs.
      # Training hyperparameters. Be prepared to tune these very carefully, as with any GAN.
      generator_lr=2e-4,  # Learning rate often needs careful tuning to the problem.
      discriminator_lr=1e-3,  # Learning rate often needs careful tuning to the problem.
      batch_size=1024,  # Batch size.
      steps=10000,  # How many steps to train both generator and discriminator for.
      init_mult1=3,  # Changing the initial parameter size can help.
      init_mult2=0.5,  #
      weight_decay=0.01,  # Weight decay.
      swa_step_start=5000,  # When to start using stochastic weight averaging.
      # Evaluation and plotting hyperparameters
      steps_per_print=10,  # How often to print the loss.
      num_plot_samples=50,  # How many samples to use on the plots at the end.
      plot_locs=(
          0.1,
          0.3,
          0.5,
          0.7,
          0.9,
      ),  # Plot some marginal distributions at this proportion of the way along.
  ):
      is_cuda = torch.cuda.is_available()
      device = "cuda" if is_cuda else "cpu"
      if not is_cuda:
          print(
              "Warning: CUDA not available; falling back to CPU but this is likely to be very slow."
          )
  
      # Data
      ts, data_size, train_dataloader = get_data(
          batch_size=batch_size, device=device
      )
      infinite_train_dataloader = (
          elem
          for it in iter(lambda: train_dataloader, None)
          for elem in it
      )
  
      # Models
      generator = Generator(
          data_size,
          initial_noise_size,
          noise_size,
          hidden_size,
          mlp_size,
          num_layers,
      ).to(device)
      discriminator = Discriminator(
          data_size, hidden_size, mlp_size, num_layers
      ).to(device)
      # Weight averaging really helps with GAN training.
      averaged_generator = swa_utils.AveragedModel(generator)
      averaged_discriminator = swa_utils.AveragedModel(discriminator)
  
      # Picking a good initialisation is important!
      # In this case these were picked by making the parameters for the t=0 part of the generator be roughly the right
      # size that the untrained t=0 distribution has a similar variance to the t=0 data distribution.
      # Then the func parameters were adjusted so that the t>0 distribution looked like it had about the right variance.
      # What we're doing here is very crude -- one can definitely imagine smarter ways of doing things.
      # (e.g. pretraining the t=0 distribution)
      with torch.no_grad():
          for param in generator._initial.parameters():
              param *= init_mult1
          for param in generator._func.parameters():
              param *= init_mult2
  
      # Optimisers. Adadelta turns out to be a much better choice than SGD or Adam, interestingly.
      generator_optimiser = torch.optim.Adadelta(
          generator.parameters(),
          lr=generator_lr,
          weight_decay=weight_decay,
      )
      discriminator_optimiser = torch.optim.Adadelta(
          discriminator.parameters(),
          lr=discriminator_lr,
          weight_decay=weight_decay,
      )
  
      # Train both generator and discriminator.
      trange = tqdm.tqdm(range(steps))
      for step in trange:
          (real_samples,) = next(infinite_train_dataloader)
  
          generated_samples = generator(ts, batch_size)
          generated_score = discriminator(generated_samples)
          real_score = discriminator(real_samples)
          loss = generated_score - real_score
          loss.backward()
  
          for param in generator.parameters():
              param.grad *= -1
          generator_optimiser.step()
          discriminator_optimiser.step()
          generator_optimiser.zero_grad()
          discriminator_optimiser.zero_grad()
  
          ###################
          # We constrain the Lipschitz constant of the discriminator using carefully-chosen clipping (and the use of
          # LipSwish activation functions).
          ###################
          with torch.no_grad():
              for module in discriminator.modules():
                  if isinstance(module, torch.nn.Linear):
                      lim = 1 / module.out_features
                      module.weight.clamp_(-lim, lim)
  
          # Stochastic weight averaging typically improves performance.
          if step > swa_step_start:
              averaged_generator.update_parameters(generator)
              averaged_discriminator.update_parameters(discriminator)
  
          if (step % steps_per_print) == 0 or step == steps - 1:
              total_unaveraged_loss = evaluate_loss(
                  ts,
                  batch_size,
                  train_dataloader,
                  generator,
                  discriminator,
              )
              if step > swa_step_start:
                  total_averaged_loss = evaluate_loss(
                      ts,
                      batch_size,
                      train_dataloader,
                      averaged_generator.module,
                      averaged_discriminator.module,
                  )
                  trange.write(
                      f"Step: {step:3} Loss (unaveraged): {total_unaveraged_loss:.4f} "
                      f"Loss (averaged): {total_averaged_loss:.4f}"
                  )
              else:
                  trange.write(
                      f"Step: {step:3} Loss (unaveraged): {total_unaveraged_loss:.4f}"
                  )
      generator.load_state_dict(averaged_generator.module.state_dict())
      discriminator.load_state_dict(
          averaged_discriminator.module.state_dict()
      )
  
      _, _, test_dataloader = get_data(
          batch_size=batch_size, device=device
      )
  
      plot(ts, generator, test_dataloader, num_plot_samples, plot_locs)
  
imports: import fire
  import matplotlib.pyplot
  import torch
  import torch.optim.swa_utils
  import torchcde
  import torchsde
  import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: main(initial_noise_size, noise_size, hidden_size, mlp_size, num_layers, generator_lr, discriminator_lr, batch_size, steps, init_mult1, init_mult2, weight_decay, swa_step_start, steps_per_print, num_plot_samples, plot_locs) without return types