pretrain_treegan.py

import torch
import torch.nn as nn
import torch.optim as optim

from data.CRN_dataset import CRNShapeNet
from model.scGAN import Generator, Discriminator
from utils.gradient_penalty import GradientPenalty

from arguments import Arguments

import time
import numpy as np
from utils.loss import *
from utils.metrics import *
import os
import os.path as osp
from eval_treegan import checkpoint_eval


class TreeGAN:
    def __init__(self, args):
        self.args = args

        ### dataset
        self.data = CRNShapeNet(args)
        self.dataLoader = torch.utils.data.DataLoader(
            self.data,
            batch_size=args.batch_size,
            shuffle=True,
            pin_memory=True,
            num_workers=16,
        )
        print("Training Dataset : {} prepared.".format(len(self.data)))

        ### Model
        self.G = Generator(
            features=args.G_FEAT,
            degrees=args.DEGREE,
            support=args.support,
            args=self.args,
        ).to(args.device)
        self.D = Discriminator(features=args.D_FEAT).to(args.device)
        self.optimizerG = optim.Adam(self.G.parameters(), lr=args.lr, betas=(0, 0.99))
        self.optimizerD = optim.Adam(self.D.parameters(), lr=args.lr, betas=(0, 0.99))
        self.GP = GradientPenalty(args.lambdaGP, gamma=1, device=args.device)

        ### uniform losses
        if self.args.expansion_penality:
            # MSN
            self.expansion = expansionPenaltyModule()
        if self.args.krepul_loss:
            # PU-net
            self.krepul_loss = kNNRepulsionLoss(
                k=self.args.krepul_k,
                n_seeds=self.args.krepul_n_seeds,
                h=self.args.krepul_h,
            )
        if self.args.knn_loss:
            # PatchVariance
            self.knn_loss = kNNLoss(k=self.args.knn_k, n_seeds=self.args.knn_n_seeds)

        print("Network prepared.")

        # ----------------------------------------------------------------------------------------------------- #
        if len(args.w_train_ls) == 1:
            self.w_train_ls = args.w_train_ls * 4
        else:
            self.w_train_ls = args.w_train_ls

    def run(self, save_ckpt=None, load_ckpt=None):

        epoch_log = 0

        loss_log = {"G_loss": [], "D_loss": []}
        loss_legend = list(loss_log.keys())

        metric = {"FPD": []}
        if load_ckpt is not None:
            checkpoint = torch.load(load_ckpt, map_location=self.args.device)
            self.D.load_state_dict(checkpoint["D_state_dict"])
            self.G.load_state_dict(checkpoint["G_state_dict"])

            epoch_log = checkpoint["epoch"]

            loss_log["G_loss"] = checkpoint["G_loss"]
            loss_log["D_loss"] = checkpoint["D_loss"]
            loss_legend = list(loss_log.keys())

            metric["FPD"] = checkpoint["FPD"]

            print("Checkpoint loaded.")
        # parallel after loading
        self.G = nn.DataParallel(self.G)
        self.D = nn.DataParallel(self.D)

        for epoch in range(epoch_log, self.args.epochs):
            epoch_g_loss = []
            epoch_d_loss = []
            epoch_time = time.time()
            self.w_train = self.w_train_ls[min(3, int(epoch / 500))]
            for _iter, data in enumerate(self.dataLoader):
                # Start Time
                start_time = time.time()
                point, _, _ = data
                point = point.to(self.args.device)

                # -------------------- Discriminator -------------------- #
                tic = time.time()
                for d_iter in range(self.args.D_iter):
                    self.D.zero_grad()
                    z = torch.randn(point.shape[0], 1, 96).to(self.args.device)

                    tree = [z]

                    with torch.no_grad():
                        fake_point = self.G(tree)

                    D_real, _ = self.D(point)
                    D_fake, _ = self.D(fake_point)
                    gp_loss = self.GP(self.D, point.data, fake_point.data)

                    # compute D loss
                    D_realm = D_real.mean()
                    D_fakem = D_fake.mean()
                    d_loss = -D_realm + D_fakem
                    d_loss_gp = d_loss + gp_loss
                    # times weight before backward
                    d_loss *= self.w_train
                    d_loss_gp.backward()
                    self.optimizerD.step()

                loss_log["D_loss"].append(d_loss.item())
                epoch_d_loss.append(d_loss.item())
                toc = time.time()
                # ---------------------- Generator ---------------------- #
                self.G.zero_grad()

                z = torch.randn(point.shape[0], 1, 96).to(self.args.device)

                tree = [z]

                fake_point = self.G(tree)
                G_fake, _ = self.D(fake_point)

                G_fakem = G_fake.mean()
                g_loss = -G_fakem
                if self.args.expansion_penality:
                    dist, _, mean_mst_dis = self.expansion(
                        fake_point,
                        self.args.expan_primitive_size,
                        self.args.expan_alpha,
                    )
                    expansion = torch.mean(dist)
                    g_loss = -G_fakem + self.args.expan_scalar * expansion
                if self.args.krepul_loss:
                    krepul_loss = self.krepul_loss(fake_point)
                    g_loss = -G_fakem + self.args.krepul_scalar * krepul_loss
                if self.args.knn_loss:
                    knn_loss = self.knn_loss(fake_point)
                    g_loss = -G_fakem + self.args.knn_scalar * knn_loss

                g_loss *= self.w_train
                g_loss.backward()
                self.optimizerG.step()

                loss_log["G_loss"].append(g_loss.item())
                epoch_g_loss.append(g_loss.item())
                tac = time.time()
                # --------------------- Visualization -------------------- #
                verbose = None
                if verbose is not None:
                    print(
                        "[Epoch/Iter] ",
                        "{:3} / {:3}".format(epoch, _iter),
                        "[ D_Loss ] ",
                        "{: 7.6f}".format(d_loss),
                        "[ G_Loss ] ",
                        "{: 7.6f}".format(g_loss),
                        "[ Time ] ",
                        "{:4.2f}s".format(time.time() - start_time),
                        "{:4.2f}s".format(toc - tic),
                        "{:4.2f}s".format(tac - toc),
                    )

            # ---------------- Epoch everage loss   --------------- #
            d_loss_mean = np.array(epoch_d_loss).mean()
            g_loss_mean = np.array(epoch_g_loss).mean()

            print(
                "[Epoch] ",
                "{:3}".format(epoch),
                "[ D_Loss ] ",
                "{: 7.6f}".format(d_loss_mean),
                "[ G_Loss ] ",
                "{: 7.6f}".format(g_loss_mean),
                "[ Time ] ",
                "{:.2f}s".format(time.time() - epoch_time),
            )
            epoch_time = time.time()

            ### call abstracted eval, which includes FPD
            if self.args.eval_every_n_epoch > 0:
                if epoch % self.args.eval_every_n_epoch == 0:
                    checkpoint_eval(
                        self.G,
                        self.args.device,
                        n_samples=5000,
                        batch_size=100,
                        conditional=False,
                        ratio="even",
                        FPD_path=self.args.FPD_path,
                        class_choices=self.args.class_choice,
                    )

            # ---------------------- Save checkpoint --------------------- #
            if epoch % self.args.save_every_n_epoch == 0 and not save_ckpt == None:
                if len(args.class_choice) == 1:
                    class_name = args.class_choice[0]
                else:
                    class_name = "multi"
                torch.save(
                    {
                        "epoch": epoch,
                        "D_state_dict": self.D.module.state_dict(),
                        "G_state_dict": self.G.module.state_dict(),
                        "D_loss": loss_log["D_loss"],
                        "G_loss": loss_log["G_loss"],
                        "FPD": metric["FPD"],
                    },
                    save_ckpt + str(epoch) + "_" + class_name + ".pt",
                )


if __name__ == "__main__":

    args = Arguments(stage="pretrain").parser().parse_args()
    args.device = torch.device(
        "cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu"
    )
    torch.cuda.set_device(args.device)

    if not osp.isdir("./pretrain_checkpoints"):
        os.mkdir("./pretrain_checkpoints")
        print("pretrain_checkpoints parent directory created.")

    if not osp.isdir(args.ckpt_path):
        os.mkdir(args.ckpt_path)

    SAVE_CHECKPOINT = (
        args.ckpt_path + args.ckpt_save if args.ckpt_save is not None else None
    )
    LOAD_CHECKPOINT = args.ckpt_load if args.ckpt_load is not None else None
    # print(args)
    model = TreeGAN(args)
    model.run(save_ckpt=SAVE_CHECKPOINT, load_ckpt=LOAD_CHECKPOINT)