Example #1
0
from ToadVAE import LevelPatchesDataset, ToadPatchVAE, ToadLevelVAE

################################################################################
# TOAD-GAN

# Logger init
logger.remove()
logger.add(sys.stdout,
           colorize=True,
           format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
           "<level>{level}</level> | " +
           "<light-black>{file.path}:{line}</light-black> | " + "{message}")

# Parse arguments
opt = get_arguments().parse_args()
opt = post_config(opt)

# Init game specific inputs
replace_tokens = {}
sprite_path = opt.game + '/sprites'
if opt.game == 'mario':
    opt.ImgGen = MarioLevelGen(sprite_path)
    replace_tokens = MARIO_REPLACE_TOKENS
    downsample = special_mario_downsampling
elif opt.game == 'mariokart':
    opt.ImgGen = MariokartLevelGen(sprite_path)
    replace_tokens = MARIOKART_REPLACE_TOKENS
    downsample = special_mariokart_downsampling
else:
    NameError("name of --game not recognized. Supported: mario, mariokart")
Example #2
0
from datetime import datetime
from config import get_arguments, post_config
from pokeGAN.training import train
from pokeGAN.functions import *
import matplotlib.pyplot as plt
import torchvision.utils as vutils
import torch

if __name__ == "__main__":
    parser = get_arguments()
    args = parser.parse_args()
    args = post_config(args)
    G_losses, D_losses, img_list = train(args=args)

    # plot
    device = get_device()
    # Grab a batch of real images from the dataloader
    dataloader = get_dataloader(args.image_dir, args.image_size, args.nc,
                                args.batch_size)
    real_batch = next(iter(dataloader))[0]
    size = real_batch.size(0)

    # create grid of 64 real images
    while size < 64:
        next_batch = next(iter(dataloader))[0]
        real_batch = torch.cat((real_batch, next_batch), dim=0)
        size = real_batch.size(0)

    # Plot the real images
    plt.figure(figsize=(14, 7))
    plt.subplot(1, 2, 1)
Example #3
0
def main():
    """ Main Training funtion. Parses inputs, inits logger, trains, and then generates some samples. """

    # torch.autograd.set_detect_anomaly(True)

    # Logger init
    logger.remove()
    logger.add(sys.stdout,
               colorize=True,
               format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
               "<level>{level}</level> | " +
               "<light-black>{file.path}:{line}</light-black> | " +
               "{message}")

    # Parse arguments
    opt = get_arguments().parse_args()
    opt = post_config(opt)

    # Init wandb
    run = wandb.init(project="mario",
                     tags=get_tags(opt),
                     config=opt,
                     dir=opt.out)
    opt.out_ = run.dir

    # Init game specific inputs
    replace_tokens = {}
    sprite_path = opt.game + '/sprites'
    if opt.game == 'mario':
        opt.ImgGen = MarioLevelGen(sprite_path)
        replace_tokens = MARIO_REPLACE_TOKENS
        downsample = special_mario_downsampling
    elif opt.game == 'mariokart':
        opt.ImgGen = MariokartLevelGen(sprite_path)
        replace_tokens = MARIOKART_REPLACE_TOKENS
        downsample = special_mariokart_downsampling
    else:
        NameError("name of --game not recognized. Supported: mario, mariokart")

    # Read level according to input arguments
    real = read_level(opt, None, replace_tokens).to(opt.device)

    # Train!
    generators, noise_maps, reals, noise_amplitudes = train(real, opt)

    # Generate Samples of same size as level
    logger.info("Finished training! Generating random samples...")
    in_s = None
    generate_samples(generators,
                     noise_maps,
                     reals,
                     noise_amplitudes,
                     opt,
                     in_s=in_s)

    # Generate samples of smaller size than level
    logger.info("Generating arbitrary sized random samples...")
    scale_v = 0.8  # Arbitrarily chosen scales
    scale_h = 0.4
    real_down = downsample(1, [[scale_v, scale_h]], real, opt.token_list)
    real_down = real_down[0]
    # necessary for correct input shape
    in_s = torch.zeros(real_down.shape, device=opt.device)
    generate_samples(generators,
                     noise_maps,
                     reals,
                     noise_amplitudes,
                     opt,
                     in_s=in_s,
                     scale_v=scale_v,
                     scale_h=scale_h,
                     save_dir="arbitrary_random_samples")