Getting this error
Capture979×455 30.4 KB
def train(self):
    # Data iterator

    data_iter = iter(self.data_loader)

    step_per_epoch = len(self.data_loader)

    model_save_step = int(self.model_save_step * step_per_epoch)

    # Fixed input for debugging

    fixed_img, _ = next(data_iter)

    fixed_z = tensor2var(torch.randn(self.batch_size, self.z_dim))

    if self.use_tensorboard:

        self.writer.add_image('img/fixed_img', denorm(fixed_img.data), 0)

    else:

        save_image(denorm(fixed_img.data),

                    os.path.join(self.sample_path, 'fixed_img.png'))

    # Start with trained model

    if self.pretrained_model:

        start = self.pretrained_model + 1

    else:

        start = 0

    self.D.train()

    self.E.train()

    self.G.train()

    # Start time

    start_time = time.time()

    for step in range(start, self.total_step):

        self.reset_grad()

        # Sample from data and prior

        try:

            real_images, _ = next(data_iter)

        except:

            data_iter = iter(self.data_loader)

            real_images, _ = next(data_iter)

        real_images = tensor2var(real_images)

        fake_z = tensor2var(torch.randn(real_images.size(0), self.z_dim))

        

        noise1 = torch.Tensor(real_images.size()).normal_().cuda()

        noise2 = torch.Tensor(real_images.size()).normal_().cuda()

        # Sample from condition

        real_z, _, _ = self.E(real_images)

        fake_images, gf1, gf2 = self.G(fake_z)

        dr, dr5, dr4, dr3, drz, dra2, dra1 = self.D(real_images+noise1, real_z)

        df, df5, df4, df3, dfz, dfa2, dfa1 = self.D(fake_images+noise2, fake_z)

        # Compute loss with real and fake images

        # dr1, dr2, df1, df2, gf1, gf2 are attention scores

        if self.adv_loss == 'wgan-gp':

            d_loss_real = - torch.mean(dr)

            d_loss_fake = df.mean()

            g_loss_fake = - df.mean()

            e_loss_real = - dr.mean() 

        elif self.adv_loss == 'hinge1':

            d_loss_real = torch.nn.ReLU()(1.0 - dr).mean()

            d_loss_fake = torch.nn.ReLU()(1.0 + df).mean()

            g_loss_fake = - df.mean()

            e_loss_real = - dr.mean() 

        elif self.adv_loss == 'hinge':

            d_loss_real = - log(dr).mean()

            d_loss_fake = - log(1.0 - df).mean()

            g_loss_fake = - log(df).mean()

            e_loss_real = - log(1.0 - dr).mean() 

        elif self.adv_loss == 'inverse':

            d_loss_real = - log(1.0 - dr).mean()

            d_loss_fake = - log(df).mean()

            g_loss_fake = - log(1.0 - df).mean()

            e_loss_real = - log(dr).mean() 

        # ================== Train D ================== #

        d_loss = d_loss_real + d_loss_fake

        d_loss.backward(retain_graph=True)

        self.d_optimizer.step()

        if self.adv_loss == 'wgan-gp':

            # Compute gradient penalty

            alpha = torch.rand(real_images.size(0), 1, 1, 1).cuda().expand_as(real_images)

            interpolated = Variable(alpha * real_images.data + (1 - alpha) * fake_images.data, requires_grad=True)

            out,_,_ = self.D(interpolated)

            grad = torch.autograd.grad(outputs=out,

                                       inputs=interpolated,

                                       grad_outputs=torch.ones(out.size()).cuda(),

                                       retain_graph=True,

                                       create_graph=True,

                                       only_inputs=True)[0]

            grad = grad.view(grad.size(0), -1)

            grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))

            d_loss_gp = torch.mean((grad_l2norm - 1) ** 2)

            # Backward + Optimize

            d_loss = self.lambda_gp * d_loss_gp

            d_loss.backward()

            self.d_optimizer.step()

        # ================== Train G and E ================== #

        ge_loss = g_loss_fake + e_loss_real

        ge_loss.backward()

        self.ge_optimizer.step()