def run_training_test(root_dir, device='cuda:0'):
    real_images = sorted(glob(os.path.join(root_dir, 'img*.nii.gz')))
    train_files = [{'reals': img} for img in zip(real_images)]
    train_transforms = Compose([LoadImaged(keys=['reals']), AsChannelFirstd(keys=['reals']), ScaleIntensityd(keys=['reals']), RandFlipd(keys=['reals'], prob=0.5), ToTensord(keys=['reals'])])
    train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5)
    train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
    learning_rate = 0.0002
    betas = (0.5, 0.999)
    real_label = 1
    fake_label = 0
    disc_net = Discriminator(in_shape=(1, 64, 64), channels=(8, 16, 32, 64, 1), strides=(2, 2, 2, 2, 1), num_res_units=1, kernel_size=5).to(device)
    disc_net.apply(normal_init)
    disc_opt = torch.optim.Adam(disc_net.parameters(), learning_rate, betas=betas)
    disc_loss_criterion = torch.nn.BCELoss()

    def discriminator_loss(gen_images, real_images):
        real = real_images.new_full((real_images.shape[0], 1), real_label)
        gen = gen_images.new_full((gen_images.shape[0], 1), fake_label)
        realloss = disc_loss_criterion(disc_net(real_images), real)
        genloss = disc_loss_criterion(disc_net(gen_images.detach()), gen)
        return torch.div(torch.add(realloss, genloss), 2)
    latent_size = 64
    gen_net = Generator(latent_shape=latent_size, start_shape=(latent_size, 8, 8), channels=[32, 16, 8, 1], strides=[2, 2, 2, 1])
    gen_net.apply(normal_init)
    gen_net.conv.add_module('activation', torch.nn.Sigmoid())
    gen_net = gen_net.to(device)
    gen_opt = torch.optim.Adam(gen_net.parameters(), learning_rate, betas=betas)
    gen_loss_criterion = torch.nn.BCELoss()

    def generator_loss(gen_images):
        output = disc_net(gen_images)
        cats = output.new_full(output.shape, real_label)
        return gen_loss_criterion(output, cats)
    key_train_metric = None
    train_handlers = [StatsHandler(name='training_loss', output_transform=(lambda x: {Keys.GLOSS: x[Keys.GLOSS], Keys.DLOSS: x[Keys.DLOSS]})), TensorBoardStatsHandler(log_dir=root_dir, tag_name='training_loss', output_transform=(lambda x: {Keys.GLOSS: x[Keys.GLOSS], Keys.DLOSS: x[Keys.DLOSS]})), CheckpointSaver(save_dir=root_dir, save_dict={'g_net': gen_net, 'd_net': disc_net}, save_interval=2, epoch_level=True)]
    disc_train_steps = 2
    num_epochs = 5
    trainer = GanTrainer(device, num_epochs, train_loader, gen_net, gen_opt, generator_loss, disc_net, disc_opt, discriminator_loss, d_train_steps=disc_train_steps, latent_shape=latent_size, key_train_metric=key_train_metric, train_handlers=train_handlers)
    trainer.run()
    return trainer.state
def run_training_test(root_dir, device='cuda:0'):
    real_images = sorted(glob(os.path.join(root_dir, 'img*.nii.gz')))
    train_files = [{'reals': img} for img in zip(real_images)]
    train_transforms = Compose([LoadImaged(keys=['reals']), AsChannelFirstd(keys=['reals']), ScaleIntensityd(keys=['reals']), RandFlipd(keys=['reals'], prob=0.5), ToTensord(keys=['reals'])])
    train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5)
    train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
    learning_rate = 0.0002
    betas = (0.5, 0.999)
    real_label = 1
    fake_label = 0
    disc_net = Discriminator(in_shape=(1, 64, 64), channels=(8, 16, 32, 64, 1), strides=(2, 2, 2, 2, 1), num_res_units=1, kernel_size=5).to(device)
    disc_net.apply(normal_init)
    disc_opt = torch.optim.Adam(disc_net.parameters(), learning_rate, betas=betas)
    disc_loss_criterion = torch.nn.BCELoss()

    def discriminator_loss(gen_images, real_images):
        real = real_images.new_full((real_images.shape[0], 1), real_label)
        gen = gen_images.new_full((gen_images.shape[0], 1), fake_label)
        realloss = disc_loss_criterion(disc_net(real_images), real)
        genloss = disc_loss_criterion(disc_net(gen_images.detach()), gen)
        return torch.div(torch.add(realloss, genloss), 2)
    latent_size = 64
    gen_net = Generator(latent_shape=latent_size, start_shape=(latent_size, 8, 8), channels=[32, 16, 8, 1], strides=[2, 2, 2, 1])
    gen_net.apply(normal_init)
    gen_net.conv.add_module('activation', torch.nn.Sigmoid())
    gen_net = gen_net.to(device)
    gen_opt = torch.optim.Adam(gen_net.parameters(), learning_rate, betas=betas)
    gen_loss_criterion = torch.nn.BCELoss()

    def generator_loss(gen_images):
        output = disc_net(gen_images)
        cats = output.new_full(output.shape, real_label)
        return gen_loss_criterion(output, cats)
    key_train_metric = None
    train_handlers = [StatsHandler(name='training_loss', output_transform=(lambda x: {Keys.GLOSS: x[Keys.GLOSS], Keys.DLOSS: x[Keys.DLOSS]})), TensorBoardStatsHandler(log_dir=root_dir, tag_name='training_loss', output_transform=(lambda x: {Keys.GLOSS: x[Keys.GLOSS], Keys.DLOSS: x[Keys.DLOSS]})), CheckpointSaver(save_dir=root_dir, save_dict={'g_net': gen_net, 'd_net': disc_net}, save_interval=2, epoch_level=True)]
    disc_train_steps = 2
    num_epochs = 5
    trainer = GanTrainer(device, num_epochs, train_loader, gen_net, gen_opt, generator_loss, disc_net, disc_opt, discriminator_loss, d_train_steps=disc_train_steps, latent_shape=latent_size, key_train_metric=key_train_metric, train_handlers=train_handlers)
    trainer.run()
    return trainer.state