def main():
    loss_fp = open('losses.csv', 'w')
    video_d_net = d_net.DiscriminatorModel(kernel_sizes_list=d_net.SCALE_KERNEL_SIZES_D, conv_layer_fms_list=d_net.SCALE_CONV_FMS_D, scale_fc_layer_sizes_list=d_net.SCALE_FC_LAYER_SIZES_D)
    video_d_net.type(dtype)
    video_g_net = g_net.VideoGANGenerator()
    video_g_net.type(dtype)
    video_d_optimizer = optim.SGD(video_d_net.parameters(), lr=0.0001)
    video_g_optimizer = optim.SGD(video_g_net.parameters(), lr=0.0001)
    max_size = len(os.listdir('train'))
    pacman_dataloader = data_loader.DataLoader('train', max_size, 16, 32, 32, 4)
    count = 0
    for i in range(1, 5000):
        (clips_x, clips_y) = pacman_dataloader.get_train_batch()
        clips_x = torch.tensor(np.rollaxis(clips_x, 3, 1)).type(dtype)
        clips_y = torch.tensor(np.rollaxis(clips_y, 3, 1)).type(dtype)
        video_d_optimizer.zero_grad()
        video_g_optimizer.zero_grad()
        batch_size = 16
        video_images = video_g_net(clips_x)
        video_d_loss_real = (video_d_net(clips_y) - 1).pow(2).mean()
        video_d_loss_fake = video_d_net(video_images).pow(2).mean()
        labels = torch.zeros(batch_size, 4).t().unsqueeze(2).type(dtype)
        video_d_loss_fake = loss_funs.adv_loss(video_d_net(video_images), labels)
        video_d_optimizer.zero_grad()
        video_d_loss_fake.backward()
        video_d_optimizer.step()
        labels = torch.ones(batch_size, 4).t().unsqueeze(2).type(dtype)
        video_d_loss_real = loss_funs.adv_loss(video_d_net(clips_y), labels)
        video_d_optimizer.zero_grad()
        video_d_loss_real.backward()
        video_d_optimizer.step()
        batch_size = 16
        video_images = video_g_net(clips_x)
        d_preds = video_d_net(video_images).type(dtype)
        gt_frames = clips_y.type(dtype)
        gen_frames = video_images.type(dtype)
        video_g_loss = loss_funs.combined_loss(gen_frames, gt_frames, d_preds)
        video_g_loss.backward()
        video_g_optimizer.step()
        if ((count % 20) == 0):
            save_samples(clips_y, count, 'video_real')
            save_samples(video_images, count, 'video_fake')
            out_str = '{}, {}, {}, {}'.format(count, video_d_loss_real, video_d_loss_fake, video_g_loss)
            print(out_str)
            loss_fp.write(out_str)
            loss_fp.write('\n')
            loss_fp.flush()
            torch.save(video_g_net.state_dict(), 'generator_net.pth.tmp')
        count += 1
    loss_fp.close()
    torch.save(video_g_net.state_dict(), 'generator_net.pth')
def main():
    loss_fp = open('losses.csv', 'w')
    video_d_net = d_net.DiscriminatorModel(kernel_sizes_list=d_net.SCALE_KERNEL_SIZES_D, conv_layer_fms_list=d_net.SCALE_CONV_FMS_D, scale_fc_layer_sizes_list=d_net.SCALE_FC_LAYER_SIZES_D)
    video_d_net.type(dtype)
    video_g_net = g_net.VideoGANGenerator()
    video_g_net.type(dtype)
    video_d_optimizer = optim.SGD(video_d_net.parameters(), lr=0.0001)
    video_g_optimizer = optim.SGD(video_g_net.parameters(), lr=0.0001)
    max_size = len(os.listdir('train'))
    pacman_dataloader = data_loader.DataLoader('train', max_size, 16, 32, 32, 4)
    count = 0
    for i in range(1, 5000):
        (clips_x, clips_y) = pacman_dataloader.get_train_batch()
        clips_x = torch.tensor(np.rollaxis(clips_x, 3, 1)).type(dtype)
        clips_y = torch.tensor(np.rollaxis(clips_y, 3, 1)).type(dtype)
        video_d_optimizer.zero_grad()
        video_g_optimizer.zero_grad()
        batch_size = 16
        video_images = video_g_net(clips_x)
        video_d_loss_real = (video_d_net(clips_y) - 1).pow(2).mean()
        video_d_loss_fake = video_d_net(video_images).pow(2).mean()
        labels = torch.zeros(batch_size, 4).t().unsqueeze(2).type(dtype)
        video_d_loss_fake = loss_funs.adv_loss(video_d_net(video_images), labels)
        video_d_optimizer.zero_grad()
        video_d_loss_fake.backward()
        video_d_optimizer.step()
        labels = torch.ones(batch_size, 4).t().unsqueeze(2).type(dtype)
        video_d_loss_real = loss_funs.adv_loss(video_d_net(clips_y), labels)
        video_d_optimizer.zero_grad()
        video_d_loss_real.backward()
        video_d_optimizer.step()
        batch_size = 16
        video_images = video_g_net(clips_x)
        d_preds = video_d_net(video_images).type(dtype)
        gt_frames = clips_y.type(dtype)
        gen_frames = video_images.type(dtype)
        video_g_loss = loss_funs.combined_loss(gen_frames, gt_frames, d_preds)
        video_g_loss.backward()
        video_g_optimizer.step()
        if ((count % 20) == 0):
            save_samples(clips_y, count, 'video_real')
            save_samples(video_images, count, 'video_fake')
            out_str = '{}, {}, {}, {}'.format(count, video_d_loss_real, video_d_loss_fake, video_g_loss)
            print(out_str)
            loss_fp.write(out_str)
            loss_fp.write('\n')
            loss_fp.flush()
            torch.save(video_g_net.state_dict(), 'generator_net.pth.tmp')
        count += 1
    loss_fp.close()
    torch.save(video_g_net.state_dict(), 'generator_net.pth')