cbcc3f315c653706ceded6ba42f222616322a3f2,ch03/03_atari_gan.py,,,#,132

Before Change

    net_discr = Discriminator(input_shape=input_shape)
    net_gener = Generator(output_shape=input_shape)
    if args.cuda:
        net_discr.cuda()
        net_gener.cuda()

    objective = nn.BCELoss()
    gen_optimizer = optim.Adam(params=net_gener.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
    dis_optimizer = optim.Adam(params=net_discr.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
    writer = SummaryWriter()

    gen_losses = []
    dis_losses = []
    iter_no = 0

    true_labels_v = Variable(torch.FloatTensor([1.0] * BATCH_SIZE))
    fake_labels_v = Variable(torch.FloatTensor([0.0] * BATCH_SIZE))
    if args.cuda:
        true_labels_v = true_labels_v.cuda()
        fake_labels_v = fake_labels_v.cuda()

    for batch_v in iterate_batches(envs):
        // generate extra fake samples, input is 4D: batch, filters, x, y
        gen_input_v = Variable(torch.FloatTensor(BATCH_SIZE, LATENT_VECTOR_SIZE, 1, 1).normal_(0, 1))
        if args.cuda:
            batch_v = batch_v.cuda()
            gen_input_v = gen_input_v.cuda()
        gen_output_v = net_gener(gen_input_v)

        // train discriminator
        dis_optimizer.zero_grad()

After Change

    parser.add_argument("--cuda", default=False, action="store_true", help="Enable cuda computation")
    args = parser.parse_args()

    device = torch.device("cuda" if args.cuda else "cpu")
    envs = [InputWrapper(gym.make(name)) for name in ("Breakout-v0", "AirRaid-v0", "Pong-v0")]
    input_shape = envs[0].observation_space.shape

Instances