def main():
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch_size', '-bs',
type=int, default=64)
parser.add_argument('--nb_epoch', '-e',
type=int, default=100)
parser.add_argument('--save_steps', '-ss',
type=int, default=10)
parser.add_argument('--lr_d',
type=float, default=2e-4)
parser.add_argument('--lr_g',
type=float, default=2e-4)
parser.add_argument('--ngf', '-ngf',
type=int, default=64)
parser.add_argument('--ndf', '-ndf',
type=int, default=64)
parser.add_argument('--latent_dim', '-ld',
type=int, default=128)
parser.add_argument('--logdir',
type=str, default='logs')
parser.add_argument('--no-cuda',
dest='use_cuda',
action='store_false', default=True,
help='disables CUDA training')
args = parser.parse_args()
device = torch.device("cuda" if args.use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if args.use_cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data', train=True, download=True,
transform=transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Lambda(lambda x: (x-0.5)*2),
])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
generator = Generator(args.latent_dim, args.ngf)
discriminator = Discriminator(args.ndf)
gan = DCGAN(generator, discriminator, device)
gan.init_params()
if args.use_cuda:
generator.cuda()
discriminator.cuda()
gan.fit(train_loader,
nb_epoch=args.nb_epoch,
lr_d=args.lr_d,
lr_g=args.lr_g,
save_steps=args.save_steps,
logdir=args.logdir)
import argparse
parser = argparse.ArgumentParser(description='Main script')
parser.add_argument('--data_dir', type=str, default='C:/Users/Jonas/Documents/GitHub/pokemon-generation/data/sprites')
parser.add_argument('--name', type=str, default='gan')
parser.add_argument('--type', type=str, default='dcgan', help='GAN Type')
parser.add_argument('--spectral_norm', action='store_true')
parser.add_argument('--epochs', type=int, default=1000)
parser.add_argument('--batch_size', type=int, default=64)
parser.add_argument('--sample_interval', type=int, default=5)
parser.add_argument('--label_column', type=str, default='type_1')
args = parser.parse_args()
from dcgan import DCGAN
from acgan import ACGAN
config = {'spectral_norm': args.spectral_norm}
if args.type == 'dcgan':
model = DCGAN(name=args.name, config=config)
elif args.type == 'acgan':
model = ACGAN(name=args.name, label_column=args.label_column, config=config)
model.fit(args.data_dir, args.epochs, args.batch_size, args.sample_interval)
