train_base.py

import argparse
import os
import sys
from datetime import datetime

import torch.nn as nn
import torch
from torch.utils.data import DataLoader

# tensorboard
from torch.utils.tensorboard import SummaryWriter

from models.maps_net import get_maps
from datasets.dataset import *


# for dice loss
def dice_loss(out, gt, smooth=1.0):
    gt = gt.view(-1)
    out = out.view(-1)

    intersection = (gt * out).sum()
    dice = (2.0 * intersection + smooth) / (torch.square(gt).sum() + torch.square(
        out).sum() + smooth)  # TODO: need to confirm this matches what the paper says, and also the calculation/result is correct

    return 1.0 - dice


# for multiprocessing
def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


# for removing damaged images
def collate_fn(batch):
    batch = list(filter(lambda x: x is not None, batch))
    return torch.utils.data.dataloader.default_collate(batch)


def parse_args():
    parser = argparse.ArgumentParser()

    ## job
    parser.add_argument("--id", type=int, help="unique ID from Slurm")
    parser.add_argument("--run_name", type=str, default="MAPS-Net", help="run name")

    ## multiprocessing
    parser.add_argument('--dist_backend', default='nccl', choices=['gloo', 'nccl'], help='multiprocessing backend')
    parser.add_argument('--master_addr', type=str, default="localhost", help='address')
    parser.add_argument('--master_port', type=int, default=3721, help='address')
    parser.add_argument('--local_rank', default=0, type=int, help='local rank')

    ## dataset
    parser.add_argument("--paths_file", type=str, default="/root/autodl-tmp/MAPS-Net-main/datasets_txt/Casiav2.txt",
                        help="path to the file with input paths")  # each line of this file should contain "/path/to/image.ext /path/to/mask.ext /path/to/edge.ext 1 (for fake)/0 (for real)"; for real image.ext, set /path/to/mask.ext and /path/to/edge.ext as a string None
    parser.add_argument("--val_paths_file", type=str, help="path to the validation set")
    parser.add_argument("--n_c_samples", type=int, help="samples per classes (None for non-controlled)")
    parser.add_argument("--val_n_c_samples", type=int,
                        help="samples per classes for validation set (None for non-controlled)")

    parser.add_argument("--workers", type=int, default=1, help="number of cpu threads to use during batch generation")

    parser.add_argument("--image_size", type=int, default=512, help="size of the images")
    parser.add_argument("--channels", type=int, default=3, help="number of image channels")

    parser.add_argument("--batch_size", type=int, default=8,
                        help="size of the batches")  # no default value given by paper

    ## model
    parser.add_argument('--load_path', type=str, help='pretrained model or checkpoint for continued training')

    ## optimizer and scheduler
    parser.add_argument("--optim", choices=['adam', 'sgd'], default='adam', help="optimizer")
    parser.add_argument("--b1", type=float, default=0.9, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--momentum", type=float, default=0.9, help="sgd: momentum of gradient")

    parser.add_argument('--patience', type=int, default=5,
                        help='numbers of epochs to decay for ReduceLROnPlateau scheduler (None to disable)')

    parser.add_argument('--decay_epoch', type=int, default=40,
                        help='numbers of epochs to decay for StepLR scheduler (low priority, None to disable)')

    ## training
    parser.add_argument("--lr", type=float, default=1e-4, help="adam: learning rate")

    parser.add_argument("--n_epochs", type=int, default=150, help="number of epochs of training")

    parser.add_argument("--cond_epoch", type=int, default=0, help="epoch to start training from")

    parser.add_argument("--n_early", type=int, default=10, help="number of epochs for early stopping")

    ## losses
    parser.add_argument("--lambda_seg", type=float, default=0.16, help="pixel-scale loss weight (alpha)")
    parser.add_argument("--lambda_clf", type=float, default=0.04, help="image-scale loss weight (beta)")
    parser.add_argument("--lambda_redge", type=float, default=10, help="redge-scale loss weight (omega)")
    ## log
    parser.add_argument("--log_interval", type=int, default=100, help="interval between saving image samples")
    parser.add_argument("--checkpoint_interval", type=int, default=1000,
                        help="batch interval between model checkpoints")

    args = parser.parse_args()

    return args


def init_env(args, local_rank, global_rank):
    # for debug only
    torch.autograd.set_detect_anomaly(True)

    if (args.id is None):
        args.id = datetime.now().strftime("%Y%m%d%H%M%S")

    torch.cuda.set_device(local_rank)
    setup_for_distributed(global_rank == 0)

    # finalizing args, print here
    print(args)

    return args


def init_models(args):
    model = get_maps().cuda()

    return model


def init_dataset(args, global_rank, world_size, val=False):
    # return None if no validation set provided
    if (val and args.val_paths_file is None):
        print('No val set!')
        return None, None

    dataset = DeepfakeDataset((args.paths_file if not val else args.val_paths_file),
                              args.image_size,
                              args.id,
                              (args.n_c_samples if not val else args.val_n_c_samples),
                              val)
    sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=global_rank,
                                                              shuffle=True)

    local_batch_size = args.batch_size // world_size

    if (not val):
        print('Local batch size is {} ({}//{})!'.format(local_batch_size, args.batch_size, world_size))

    dataloader = DataLoader(dataset=dataset, batch_size=local_batch_size, num_workers=args.workers, pin_memory=True,
                            drop_last=False, sampler=sampler, collate_fn=collate_fn)

    print('{} set size is {}!'.format(('Train' if not val else 'Val'), len(dataloader) * args.batch_size))

    return sampler, dataloader


def init_optims(args, world_size,
                model):
    # Optimizers
    local_lr = args.lr / world_size

    print('Local learning rate is {} ({}/{})!'.format(local_lr, args.lr, world_size))

    if (args.optim == 'adam'):
        print("Using optimizer adam")
        optimizer = torch.optim.Adam(model.parameters(), lr=local_lr, betas=(args.b1, args.b2))
    elif (args.optim == 'sgd'):
        print("Using optimizer sgd")
        optimizer = torch.optim.SGD(model.parameters(), lr=local_lr, momentum=args.momentum)
    else:
        print("Unrecognized optimizer %s" % args.optim)
        sys.exit()

    return optimizer


def init_schedulers(args, optimizer):
    lr_scheduler = None

    # high priority for ReduceLROnPlateau (validation set required)
    if (args.val_paths_file and args.patience):
        print("Using scheduler ReduceLROnPlateau")
        lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,
                                                                  factor=0.1,
                                                                  patience=args.patience)
    # low priority StepLR
    elif (args.decay_epoch):
        print("Using scheduler StepLR")
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,
                                                       step_size=args.decay_epoch,
                                                       gamma=0.5)

    else:
        print("No scheduler used")

    return lr_scheduler


def load_dicts(args, get_module,
               model):
    # Load pretrained models
    if args.load_path != None and args.load_path != 'timm':
        print('Load pretrained model: {}'.format(args.load_path))

        if (not get_module):
            model.load_state_dict(torch.load(args.load_path))
        else:
            model.module.load_state_dict(torch.load(args.load_path))

    return model


# for saving checkpoints
def save_checkpoints(checkpoint_dir, id, epoch, step, get_module,
                     model):
    if (get_module):
        net = model.module
    else:
        net = model
    
    if ((epoch % 10 == 9) or (epoch % 10 == 4)):
        torch.save(net.state_dict(),
               os.path.join(checkpoint_dir, str(id) + "_" + str(epoch) + '_' + str(step) + '.pth'))


def denormalize(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    """denormalize image with mean and std
    """
    image = image.clone().detach().cpu()
    image = image * torch.tensor(std).view(3, 1, 1)
    image = image + torch.tensor(mean).view(3, 1, 1)
    return image


# a single step of prediction and loss calculation (same for both training and validating)
def predict_loss(args, data, model,
                 criterion_BCE, criterion_MSE,
                 gmp):
    # load data
    in_imgs, in_masks, in_edges, in_labels = data

    in_imgs = in_imgs.to(device=torch.device('cuda'), dtype=torch.float32)

    in_masks = in_masks.to(device=torch.device('cuda'), dtype=torch.float32)
    in_edges = in_edges.to(device=torch.device('cuda'), dtype=torch.float32)
    in_labels = in_labels.to(device=torch.device('cuda'), non_blocking=True).float()

    # predict
    out_edges, out_masks, ret2 = model(in_imgs)
    out_edges = torch.sigmoid(out_edges)
    out_masks = torch.sigmoid(out_masks)
    out_masks = out_masks.to(device=torch.device('cuda'), dtype=torch.float32)

    # Pixel-scale loss
    loss_seg = dice_loss(out_masks, in_masks)

    # Edge loss
    loss_edg = dice_loss(out_edges, in_edges)
    # Redge loss
    zero_tensor = torch.zeros_like(ret2)
    loss_redge = criterion_MSE(ret2, zero_tensor.cuda())
    # Image-scale loss (with GMP)
    out_labels = gmp(out_masks).squeeze()
    in_labels = in_labels.view(out_labels.size())
    loss_clf = criterion_BCE(out_labels, in_labels)

    # Total loss
    alpha = args.lambda_seg
    beta = args.lambda_clf
    omega = args.lambda_redge

    weighted_loss_seg = alpha * loss_seg
    weighted_loss_clf = beta * loss_clf
    weighted_loss_redge = omega * loss_redge
    weighted_loss_edg = (1.0 - alpha - beta) * loss_edg

    loss = weighted_loss_seg + weighted_loss_clf + weighted_loss_edg + weighted_loss_redge

    return loss, weighted_loss_seg, weighted_loss_clf, weighted_loss_edg, weighted_loss_redge, in_imgs, in_masks, in_edges, out_masks, out_edges, ret2


def write_loss_to_file(loss_list, filename):
    with open(filename, 'w') as file:
        for i, loss in enumerate(loss_list):
            file.write(f'Epoch {i + 1}: {loss}\n')


# self convGem layer

class GemPool2d(nn.Module):
    def __init__(self, p=3.0, eps=1e-6):
        super(GemPool2d, self).__init__()
        self.p = p
        self.eps = eps

    def forward(self, x):
        # 对输入张量进行幂运算
        x = x.pow(self.p)

        # 计算均值池化
        x = x.mean(dim=(-2, -1)).clamp(min=self.eps).pow(1.0 / self.p)

        return x


class CustomPoolingBlock(nn.Module):
    def __init__(self, a):
        super(CustomPoolingBlock, self).__init__()
        self.a = a
        self.in_channels = 1
        self.out_channels = 32
        # 3x3卷积层
        self.conv1 = nn.Conv2d(self.in_channels, self.out_channels, kernel_size=3, padding=1)

        # ReLU激活函数
        self.relu = nn.ReLU()

        # Batch Norm
        self.batch_norm = nn.BatchNorm2d(self.out_channels)

        # 1x1卷积层
        self.conv2 = nn.Conv2d(self.out_channels, 1, kernel_size=1)

        # Sigmoid激活函数
        self.sigmoid = nn.Sigmoid()

        # Gem池化层
        self.gem_pool1 = GemPool2d()
        self.gem_pool2 = GemPool2d()

    def forward(self, x):
        # 分支1：卷积、激活、Batch Norm、卷积、激活
        out1 = self.conv1(x)
        out1 = self.relu(out1)
        out1 = self.batch_norm(out1)
        out1 = self.conv2(out1)
        out1 = self.sigmoid(out1)
        out1 = self.gem_pool1(out1)
        # 分支2：Gem池化
        out2 = self.gem_pool2(x)

        # 混合两个分支
        output = (1 - self.a) * out1 + self.a * out2

        return output


# 指定要保存的文件名
train_loss_filename = './loss/train_loss.txt'
val_loss_filename = './loss/val_loss.txt'
Pixel_scale_Loss_filename = './loss/pixel_scale_Loss.txt'
Edge_Loss_list_filename = './loss/edge_loss.txt'
Image_scale_Loss_list_filename = './loss/image_scale_loss.txt'
redge_scale_Loss_list_filename = './loss/redge_scale_loss.txt'


def train(args, global_rank, sync, get_module,
          model,
          train_sampler, dataloader, val_sampler, val_dataloader,
          optimizer,
          lr_scheduler):
    # Losses that are built-in in PyTorch
    criterion_BCE = nn.BCEWithLogitsLoss().cuda()
    criterion_MSE = nn.MSELoss().cuda()
    # tensorboard
    if global_rank == 0:
        os.makedirs("logs", exist_ok=True)
        writer = SummaryWriter("logs/" + str(args.id) + "_" + args.run_name)
        checkpoint_dir = "checkpoints/" + str(args.id) + "_" + args.run_name
        os.makedirs(checkpoint_dir, exist_ok=True)

    # for early stopping
    best_val_loss = float('inf')
    n_last_epochs = 0
    early_stopping = False

    # GMP layer
    #     gmp = nn.MaxPool2d(args.image_size)
    gmp = CustomPoolingBlock(a=0.9).cuda()
    # self convGem layer
    Total_Loss_list = []
    Pixel_scale_Loss_list = []
    Edge_Loss_list = []
    Image_scale_Loss_list = []
    Val_Loss_list = []
    best_total_loss = float('inf')
    for epoch in range(args.cond_epoch, args.n_epochs):

        train_sampler.set_epoch(epoch)

        print('Starting Epoch {}'.format(epoch))

        # loss sum for epoch
        epoch_total_seg = 0
        epoch_total_clf = 0
        epoch_total_edg = 0

        epoch_total_model = 0

        epoch_val_loss = 0

        # number of steps in one epoch
        # can be replaced by len(dataloader), but kept as warm-up epochs may be added
        epoch_steps = 0

        # ------------------
        #  Train step
        # ------------------
        for step, data in enumerate(dataloader):
            curr_steps = epoch * len(dataloader) + step

            model.train()

            if (sync): optimizer.synchronize()
            optimizer.zero_grad()

            loss, weighted_loss_seg, weighted_loss_clf, weighted_loss_edg, weighted_loss_redge, in_imgs, in_masks, in_edges, out_masks, out_edges, ret2 = predict_loss(
                args, data, model, criterion_BCE, criterion_MSE, gmp)

            # backward prop
            loss.backward()
            optimizer.step()

            # log losses for epoch
            epoch_steps += 1

            epoch_total_seg += weighted_loss_seg.item()
            epoch_total_clf += weighted_loss_clf.item()
            epoch_total_edg += weighted_loss_edg.item()
            #             epoch_total_redge +=weighted_loss_redge.item()
            epoch_total_model += loss.item()

            # --------------
            #  Log Progress (for certain steps)
            # --------------
            if step != 0 and step % args.log_interval == 0 and global_rank == 0:
                print(f"[Epoch {epoch}/{args.n_epochs - 1}] [Batch {step}/{len(dataloader)}] "
                      f"[Total Loss {loss:.3f}]"
                      f"[Pixel-scale Loss {weighted_loss_seg:.3e}]"
                      f"[Edge Loss {weighted_loss_edg:.3e}]"
                      f"[Image-scale Loss {weighted_loss_clf:.3e}]"
                      f"[Image-redge Loss {weighted_loss_redge:.3e}]"
                      f"")

                writer.add_scalar("LearningRate", optimizer.param_groups[0]['lr'], curr_steps)
                writer.add_scalar("Loss/Total Loss", loss, epoch * len(dataloader) + step)
                writer.add_scalar("Loss/Pixel-scale", weighted_loss_seg, curr_steps)
                writer.add_scalar("Loss/Edge", weighted_loss_edg, curr_steps)
                writer.add_scalar("Loss/Image-scale", weighted_loss_clf, curr_steps)
                writer.add_scalar("Loss/redge-scale", weighted_loss_redge, curr_steps)
                in_imgs = denormalize(in_imgs)
                writer.add_images('Input Img', in_imgs, epoch * len(dataloader) + step)
                in_masks = in_masks.unsqueeze(1)
                writer.add_images('Input Mask', in_masks, epoch * len(dataloader) + step)
                writer.add_images('Output Mask', out_masks, epoch * len(dataloader) + step)
                writer.add_images('Input Edge', in_edges, epoch * len(dataloader) + step)
                writer.add_images('Output Edge', out_edges, epoch * len(dataloader) + step)

            # save model parameters
            if step != 0 and step % args.checkpoint_interval == 0 and global_rank == 0:
                save_checkpoints(checkpoint_dir, args.id, epoch, step, 0,
                                 model)
        Total_Loss_list.append(loss.item())
        Pixel_scale_Loss_list.append(weighted_loss_seg.item())
        Edge_Loss_list.append(weighted_loss_edg.item())
        Image_scale_Loss_list.append(weighted_loss_clf.item())
        #  Validation
        # ------------------
        # 调用函数将损失值写入文件
        write_loss_to_file(Total_Loss_list, train_loss_filename)
        write_loss_to_file(Pixel_scale_Loss_list, Pixel_scale_Loss_filename)
        write_loss_to_file(Edge_Loss_list, Edge_Loss_list_filename)
        write_loss_to_file(Image_scale_Loss_list, Image_scale_Loss_list_filename)
        #         write_loss_to_file(redge_scale_Loss_list, redge_scale_Loss_list_filename)
        # ------------------
        #  Validation
        # ------------------
        if (args.val_paths_file and val_sampler and val_dataloader):

            val_sampler.set_epoch(epoch)

            model.eval()

            for step, data in enumerate(val_dataloader):
                with torch.no_grad():
                    loss, _, _, _, _, _, _, _, _, _, _ = predict_loss(args, data, model, criterion_BCE, gmp)

                    epoch_val_loss += loss.item()

            # early
            if epoch_val_loss <= best_val_loss:
                best_val_loss = epoch_val_loss
                n_last_epochs = 0
            else:
                n_last_epochs += 1

                if (n_last_epochs >= args.n_early):
                    early_stopping = False

        # ------------------
        #  Step
        # ------------------
        if (lr_scheduler):
            if (args.val_paths_file and args.patience):
                lr_scheduler.step(epoch_val_loss)  # ReduceLROnPlateau
            elif (args.decay_epoch):
                lr_scheduler.step()  # StepLR
            else:
                print("Error in scheduler step")
                sys.exit()

        # --------------
        #  Log Progress (for epoch)
        # --------------
        # loss average for epoch
        if epoch_steps != 0 and global_rank == 0:
            epoch_avg_seg = epoch_total_seg / epoch_steps
            epoch_avg_edg = epoch_total_edg / epoch_steps
            epoch_avg_clf = epoch_total_clf / epoch_steps
            epoch_avg_model = epoch_total_model / epoch_steps

            if (args.val_paths_file):
                epoch_val_loss_avg = epoch_val_loss / len(val_dataloader)
                best_val_loss_avg = best_val_loss / len(val_dataloader)
            else:
                epoch_val_loss_avg = 0
                best_val_loss_avg = 0

            print(f"[Epoch {epoch}/{args.n_epochs - 1}]"
                  f"[Epoch Total Loss {epoch_avg_model:.3f}]"
                  f"[Epoch Pixel-scale Loss {epoch_avg_seg:.3e}]"
                  f"[Epoch Edge Loss {epoch_avg_edg:.3e}]"
                  f"[Epoch Image-scale Loss {epoch_avg_clf:.3e}]"
                  f"[Epoch Val Loss {epoch_val_loss_avg:.3f} (best Val Loss {best_val_loss_avg:.3f} last for {n_last_epochs:d})]"
                  f"")

            writer.add_scalar("Epoch LearningRate", optimizer.param_groups[0]['lr'], epoch)
            writer.add_scalar("Epoch Loss/Total Loss", epoch_avg_model, epoch)
            writer.add_scalar("Epoch Loss/Pixel-scale", epoch_avg_seg, epoch)
            writer.add_scalar("Epoch Loss/Edge", epoch_avg_edg, epoch)
            writer.add_scalar("Epoch Loss/Image-scale", epoch_avg_clf, epoch)
            writer.add_scalar("Epoch Loss/Val", epoch_val_loss_avg, epoch)
            if torch.max(in_imgs) > 1 or torch.min(in_imgs) < 0:
                in_imgs = denormalize(in_imgs)
            writer.add_images('Epoch Input Img', in_imgs, epoch)
            if len(in_masks.shape) == 3:
                in_masks = in_masks.unsqueeze(1)
            writer.add_images('Epoch Input Mask', in_masks, epoch)
            writer.add_images('Epoch Output Mask', out_masks, epoch)
            writer.add_images('Epoch Input Edge', in_edges, epoch)
            writer.add_images('Epoch Output Edge', out_edges, epoch)

            # save model parameters
            if global_rank == 0:
                save_checkpoints(checkpoint_dir, args.id, epoch, 'end',  # set step to a string 'end'
                                 0,
                                 model)

        # check early_stopping
        if (early_stopping):
            print('Early stopping')
            break

    print('Finished training')

    if global_rank == 0:
        writer.close()

    pass