buffer_mtt.py

import os
import argparse
import torch
import torch.nn as nn
from tqdm import tqdm
from utils import get_dataset, get_network, get_daparam, TensorDataset, epoch, ParamDiffAug
import copy

import torch.nn.utils.prune as prune

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

def get_params_with_prune(teacher_net):
    params = []
    state_dict = teacher_net.state_dict()
    for name, p in teacher_net.named_parameters():
        if 'orig' in name:
            params.append(p.detach().cpu() * state_dict[name[:-5] + "_mask"].detach().cpu())
        elif 'mask' not in name:
            params.append(p.detach().cpu())
        else:
            pass
    
    return params

def main(args):

    args.dsa = True if args.dsa == 'True' else False
    args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    args.dsa_param = ParamDiffAug()

    channel, im_size, num_classes, _, _, _, dst_train, _, testloader, _, class_map, _ = get_dataset(args.dataset, args.data_path, args.batch_real, args.subset, args=args)

    # print('\n================== Exp %d ==================\n '%exp)
    print('Hyper-parameters: \n', args.__dict__)

    save_dir = os.path.join(args.buffer_path, args.dataset)
    if args.dataset == "ImageNet":
        save_dir = os.path.join(save_dir, args.subset, str(args.res))
    if args.dataset in ["CIFAR10", "CIFAR100"] and not args.zca:
        save_dir += "_NO_ZCA"
    if args.use_random_pruning:
        save_dir += f"_random_pruning_{args.sparse_ratio}"
    save_dir = os.path.join(save_dir, args.model)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)


    images_all = []
    labels_all = []
    indices_class = [[] for c in range(num_classes)]
    print("BUILDING DATASET")
    for i in tqdm(range(len(dst_train))):
        sample = dst_train[i]
        images_all.append(torch.unsqueeze(sample[0], dim=0))
        labels_all.append(class_map[torch.tensor(sample[1]).item()])

    for i, lab in tqdm(enumerate(labels_all)):
        indices_class[lab].append(i)
    images_all = torch.cat(images_all, dim=0).to("cpu")
    labels_all = torch.tensor(labels_all, dtype=torch.long, device="cpu")

    for c in range(num_classes):
        print(f'class c = {c}: {len(indices_class[c])} real images')

    for ch in range(channel):
        print(f'real images channel {ch}, mean = {torch.mean(images_all[:, ch])}, std = {torch.std(images_all[:, ch])})')

    criterion = nn.CrossEntropyLoss().to(args.device)

    trajectories = []

    dst_train = TensorDataset(copy.deepcopy(images_all.detach()),
                              copy.deepcopy(labels_all.detach()))
    trainloader = torch.utils.data.DataLoader(dst_train,
                                              batch_size=args.batch_train,
                                              shuffle=True,
                                              num_workers=0)

    ''' set augmentation for whole-dataset training '''
    args.dc_aug_param = get_daparam(args.dataset, args.model, args.model, None)
    args.dc_aug_param['strategy'] = 'crop_scale_rotate'  # for whole-dataset training
    print('DC augmentation parameters: \n', args.dc_aug_param)

    for it in range(0, args.num_experts):

        teacher_net = get_network(args.model,
                                  channel,
                                  num_classes,
                                  im_size).to(args.device) # get a random model
        teacher_net.train()
        if args.use_random_pruning:
            for m in teacher_net.modules():
                if isinstance(m, nn.Conv2d):
                    prune.random_unstructured(m, 'weight', args.sparse_ratio)
        lr = args.lr_teacher
        teacher_optim = torch.optim.SGD(teacher_net.parameters(), lr=lr, momentum=args.mom, weight_decay=args.l2)  # optimizer_img for synthetic data
        teacher_optim.zero_grad()

        timestamps = []

        timestamps.append(get_params_with_prune(teacher_net))
        print(timestamps)
        lr_schedule = [args.train_epochs // 2 + 1]

        for e in range(args.train_epochs):

            train_loss, train_acc = epoch("train",
                                          dataloader=trainloader,
                                          net=teacher_net,
                                          optimizer=teacher_optim,
                                          criterion=criterion,
                                          args=args,
                                          aug=True)

            test_loss, test_acc = epoch("test",
                                        dataloader=testloader,
                                        net=teacher_net,
                                        optimizer=None,
                                        criterion=criterion,
                                        args=args,
                                        aug=False)

            print(f"Itr: {it}\tEpoch: {e}\tTrain Acc: {train_acc}\tTest Acc: {test_acc}")
            timestamps.append(get_params_with_prune(teacher_net))

            if e in lr_schedule and args.decay:
                lr *= 0.1
                teacher_optim = torch.optim.SGD(
                    teacher_net.parameters(), lr=lr, momentum=args.mom, weight_decay=args.l2)
                teacher_optim.zero_grad()

        trajectories.append(timestamps)
        if len(trajectories) == args.save_interval:
            n = 0
            while os.path.exists(os.path.join(save_dir, f"replay_buffer_{n}.pt")):
                n += 1
            torch.save(trajectories, os.path.join(save_dir, f"replay_buffer_{n}.pt"))
            trajectories = []


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Parameter Processing')
    parser.add_argument('--dataset', type=str, default='CIFAR10', help='dataset')
    parser.add_argument('--subset', type=str, default='imagenette', help='subset')
    parser.add_argument('--model', type=str, default='ConvNet', help='model')
    parser.add_argument('--num_experts', type=int, default=100, help='training iterations')
    parser.add_argument('--lr_teacher', type=float, default=0.01, help='learning rate for updating network parameters')
    parser.add_argument('--batch_train', type=int, default=256, help='batch size for training networks')
    parser.add_argument('--batch_real', type=int, default=256, help='batch size for real loader')
    parser.add_argument('--dsa', type=str, default='True', choices=['True', 'False'],
                        help='whether to use differentiable Siamese augmentation.')
    parser.add_argument('--dsa_strategy', type=str, default='color_crop_cutout_flip_scale_rotate',
                        help='differentiable Siamese augmentation strategy')
    parser.add_argument('--data_path', type=str, default='data', help='dataset path')
    parser.add_argument('--buffer_path', type=str, default='./buffers', help='buffer path')
    parser.add_argument('--train_epochs', type=int, default=50)
    parser.add_argument('--zca', action='store_true')
    parser.add_argument('--decay', action='store_true')
    parser.add_argument('--mom', type=float, default=0, help='momentum')
    parser.add_argument('--l2', type=float, default=0, help='l2 regularization')
    parser.add_argument('--save_interval', type=int, default=10)
    parser.add_argument('--use-random-pruning', action='store_true')
    parser.add_argument('--sparse-ratio', type=float, default=0.0)

    args = parser.parse_args()
    main(args)