Function main

cnn-CGANCode/train.py:95–121

(args)

def main(args):
    #加载数据集
    dataset = torchvision.datasets.MNIST(root='data/',train=True,transform=config.transform,
                                         download=True)
    dataloader = DataLoader(
        dataset=dataset,
        batch_size=config.BATCH_SIZE,
        shuffle=config.SHUFFLE,
        drop_last=config.DROP_LAST,

    )
    #加载模型
    generator = Generator().to(config.DEVICE)
    discriminator = Discriminator().to(config.DEVICE)
    generator.apply(weights_init_normal)
    discriminator.apply(weights_init_normal)

    #定义损失函数和优化器
    adversarial_loss = torch.nn.BCELoss()
    optimizer_G = torch.optim.Adam(generator.parameters(),lr=config.LEARING_RATIO,betas=(config.BETA1,config.BETA2))
    optimizer_D = torch.optim.Adam(discriminator.parameters(),lr = config.LEARING_RATIO,betas=(config.BETA1,config.BETA2))

    for epoch in range(config.NUM_EPOCHS):
        train_fn(generator,discriminator,optimizer_G,optimizer_D,adversarial_loss,dataloader,epoch)
        if epoch % 10 == 0:
            saveImage(generator, epoch)
            save_checkpoint(generator,optimizer_G,config.CHECKPOINT_GEN)


if __name__ == '__main__':