train_mesm.py

import argparse
import os
import yaml

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter

import datasets
import models
import utils
import utils.few_shot as fs
from datasets.samplers import CategoriesSampler
from clock_driven import functional


def main(config):
    svname = args.name
    if svname is None:
        svname = 'meta_{}-{}shot'.format(
            config['train_dataset'], config['n_shot'])
        svname += '_' + config['model'] + '-' + config['model_args']['encoder']
    if args.tag is not None:
        svname += '_' + args.tag
    save_path = os.path.join('./save1',
                             svname + f"_{config['n_way']}way_{args.method}_{config['optimizer']}")  # meta实验结果保存在save1文件夹下
    utils.ensure_path(save_path)
    utils.set_log_path(save_path)
    writer = SummaryWriter(os.path.join(save_path, 'tensorboard'))

    yaml.dump(config, open(os.path.join(save_path, 'config.yaml'), 'w'))

    #### Dataset ####

    n_way, n_shot = config['n_way'], config['n_shot']
    n_query = config['n_query']

    if config.get('n_train_way') is not None:
        n_train_way = config['n_train_way']
    else:
        n_train_way = n_way
    if config.get('n_train_shot') is not None:
        n_train_shot = config['n_train_shot']
    else:
        n_train_shot = n_shot
    # if config.get('ep_per_batch') is not None:
    #     ep_per_batch = config['ep_per_batch']
    # else:
    #     ep_per_batch = 1

    # train
    train_dataset = datasets.make(config['train_dataset'], **config['train_dataset_args'])
    utils.log(
        'train dataset: {} (x{}), {}'.format(train_dataset[0][0].shape, len(train_dataset), train_dataset.n_classes))
    # if config.get('visualize_datasets'):
    #     utils.visualize_dataset(train_dataset, 'train_dataset', writer)
    train_sampler = CategoriesSampler(train_dataset.label, config['train_batches'], n_train_way, n_train_shot, n_query,
                                      ep_per_batch=config['train_ep_per_batch'])
    train_loader = DataLoader(train_dataset, batch_sampler=train_sampler, num_workers=8, pin_memory=True)

    # tval
    if config.get('tval_dataset'):
        tval_dataset = datasets.make(config['tval_dataset'], **config['tval_dataset_args'])
        utils.log(
            'tval dataset: {} (x{}), {}'.format(tval_dataset[0][0].shape, len(tval_dataset), tval_dataset.n_classes))
        # if config.get('visualize_datasets'):
        #     utils.visualize_dataset(tval_dataset, 'tval_dataset', writer)
        tval_sampler = CategoriesSampler(tval_dataset.label, config['val_batches'], n_way, n_shot, n_query,
                                         ep_per_batch=config['val_ep_per_batch'])
        tval_loader = DataLoader(tval_dataset, batch_sampler=tval_sampler, num_workers=8, pin_memory=True)
    else:
        tval_loader = None

    # val
    val_dataset = datasets.make(config['val_dataset'], **config['val_dataset_args'])
    utils.log('val dataset: {} (x{}), {}'.format(val_dataset[0][0].shape, len(val_dataset), val_dataset.n_classes))
    # if config.get('visualize_datasets'):
    #     utils.visualize_dataset(val_dataset, 'val_dataset', writer)
    val_sampler = CategoriesSampler(val_dataset.label, config['val_batches'], n_way, n_shot, n_query,
                                    ep_per_batch=config['val_ep_per_batch'])
    val_loader = DataLoader(val_dataset, batch_sampler=val_sampler, num_workers=8, pin_memory=True)

    ########

    #### Model and optimizer ####

    if config.get('load'):
        model_sv = torch.load(config['load'])
        model = models.load(model_sv)
    else:
        model = models.make(config['model'], **config['model_args'])

        if config.get('load_encoder'):
            encoder = models.load(torch.load(config['load_encoder'])).encoder
            model.encoder.load_state_dict(encoder.state_dict())

    model.set_method(args.method)

    if config.get('_parallel'):
        print(config.get('_parallel'))
        model = nn.DataParallel(model)

    utils.log('num params: {}'.format(utils.compute_n_params(model)))
    optimizer, lr_scheduler = utils.make_optimizer(model.parameters(), config['optimizer'], **config['optimizer_args'])
    # encoder1 = encoding.PoissonEncoder()  # 泊松编码

    ########

    max_epoch = config['max_epoch']
    save_epoch = config.get('save_epoch')
    max_va = 0.
    timer_used = utils.Timer()
    timer_epoch = utils.Timer()

    aves_keys = ['tl', 'ta', 'tvl', 'tva', 'vl', 'va']
    trlog = dict()
    for k in aves_keys:
        trlog[k] = []

    for epoch in range(1, max_epoch + 1):
        timer_epoch.s()
        aves = {k: utils.Averager() for k in aves_keys}

        # train
        model.train()
        if config.get('freeze_bn'):
            utils.freeze_bn(model)
        writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)

        np.random.seed(epoch)
        for data, _ in tqdm(train_loader, desc='train', leave=False):

            x_shot, x_query = fs.split_shot_query(data.cuda(), n_train_way, n_train_shot, n_query,
                                                  ep_per_batch=config['train_ep_per_batch'])
            label = fs.make_nk_label(n_train_way, n_query, ep_per_batch=config['train_ep_per_batch']).cuda()
            label = label.reshape(x_query.shape[:2])

            # 前向无梯度计算x_shot
            shot_shape = x_shot.shape[:3]  # 5：torch.Size([1, 5, 3])   2：[1, 2, 3]
            img_shape = x_shot.shape[-3:]  # 5：torch.Size([1, 80, 80])   2：[1, 80, 80]
            x_shot = x_shot.view(-1, *img_shape)  # 5：[15, 1, 80, 80]  2:[6,1,80,80]
            with torch.no_grad():
                if isinstance(model, torch.nn.DataParallel):
                    x_shot = model.module.encoder(x_shot)
                else:
                    x_shot = model.encoder(x_shot)
                functional.reset_net(model)
            channel_dim = x_shot.shape[-3]
            x_shot = x_shot.view(*shot_shape, channel_dim, -1)

            model.train()
            query_num = x_query.shape[1]
            for i in range(0, query_num, config['training_batch']):
                # print(x_shot.shape, x_query[:, i: min(i+config['training_batch'], query_num), ...].shape)
                logits = model(x_shot, x_query[:, i: min(i + config['training_batch'], query_num), ...]).view(-1,
                                                                                                              n_train_way).requires_grad_()
                loss = F.cross_entropy(logits, label[:, i: min(i + config['training_batch'], query_num)].reshape(-1))
                acc = utils.compute_acc(logits, label[:, i: min(i + config['training_batch'], query_num)].reshape(-1))

                loss.backward()
                functional.reset_net(
                    model)  ##############################################################################

                aves['tl'].add(loss.item())
                aves['ta'].add(acc)
            optimizer.step()
            optimizer.zero_grad()
            # break

        # eval
        model.eval()
        functional.reset_net(model)  ##############################################################################
        for name, loader, name_l, name_a in [
            ('tval', tval_loader, 'tvl', 'tva'),
            ('val', val_loader, 'vl', 'va')]:

            if (config.get('tval_dataset') is None) and name == 'tval':
                continue

            np.random.seed(0)
            for data, _ in tqdm(loader, desc=name, leave=False):
                x_shot, x_query = fs.split_shot_query(data.cuda(), n_way, n_shot, n_query,
                                                      ep_per_batch=config['val_ep_per_batch'])
                label = fs.make_nk_label(n_way, n_query, ep_per_batch=config['val_ep_per_batch']).cuda()

                # 前向无梯度计算x_shot
                shot_shape = x_shot.shape[:-3]  # 5：torch.Size([1, 5, 3])   2：[1, 2, 3]
                img_shape = x_shot.shape[-3:]  # 5：torch.Size([1, 80, 80])   2：[1, 80, 80]
                x_shot = x_shot.view(-1, *img_shape)  # 5：[15, 1, 80, 80]  2:[6,1,80,80]
                with torch.no_grad():
                    if isinstance(model, torch.nn.DataParallel):
                        x_shot = model.module.encoder(x_shot)
                    else:
                        x_shot = model.encoder(x_shot)
                    functional.reset_net(model)
                channel_dim = x_shot.shape[-3]
                x_shot = x_shot.view(*shot_shape, channel_dim, -1)

                with torch.no_grad():  ####################################################
                    # print(x_shot.shape, x_query.shape)
                    logits = model(x_shot, x_query).view(-1, n_way)
                loss = F.cross_entropy(logits, label)
                acc = utils.compute_acc(logits, label)
                functional.reset_net(model)
                aves[name_l].add(loss.item())
                aves[name_a].add(acc)
        _sig = int(_[-1])
        # post
        if lr_scheduler is not None:
            lr_scheduler.step()
        for k, v in aves.items():
            aves[k] = v.item()
            trlog[k].append(aves[k])
        t_epoch = utils.time_str(timer_epoch.t())
        t_used = utils.time_str(timer_used.t())
        t_estimate = utils.time_str(timer_used.t() / epoch * max_epoch)
        utils.log('epoch {}, train {:.4f}|{:.4f}, tval {:.4f}|{:.4f}, '
                  'val {:.4f}|{:.4f}, {} {}/{} (@{})'.format(
            epoch, aves['tl'], aves['ta'], aves['tvl'], aves['tva'],
            aves['vl'], aves['va'], t_epoch, t_used, t_estimate, _sig))

        writer.add_scalars('loss', {
            'train': aves['tl'],
            'tval': aves['tvl'],
            'val': aves['vl'],
        }, epoch)
        writer.add_scalars('acc', {
            'train': aves['ta'],
            'tval': aves['tva'],
            'val': aves['va'],
        }, epoch)

        if config.get('_parallel'):
            model_ = model.module
        else:
            model_ = model

        training = {
            'epoch': epoch,
            'optimizer': config['optimizer'],
            'optimizer_args': config['optimizer_args'],
            'optimizer_sd': optimizer.state_dict(),
        }
        save_obj = {
            'file': __file__,
            'config': config,
            'model': config['model'],
            'model_args': config['model_args'],
            'model_sd': model_.state_dict(),
            'training': training,
        }
        torch.save(save_obj, os.path.join(save_path, 'epoch-last.pth'))
        torch.save(trlog, os.path.join(save_path, 'trlog.pth'))

        if (save_epoch is not None) and epoch % save_epoch == 0:
            torch.save(save_obj,
                       os.path.join(save_path, 'epoch-{}.pth'.format(epoch)))
        if aves['va'] > max_va:
            max_va = aves['va']
            torch.save(save_obj, os.path.join(save_path, 'max-va.pth'))

        writer.flush()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config')
    parser.add_argument('--name', default=None)
    parser.add_argument('--tag', default=None)
    parser.add_argument('--gpu', default='0')
    parser.add_argument('--method', default='CKA')
    args = parser.parse_args()

    config = yaml.load(open(args.config, 'r'), Loader=yaml.FullLoader)
    if len(args.gpu.split(',')) > 1:
        config['_parallel'] = True
        config['_gpu'] = args.gpu

    utils.set_gpu(args.gpu)
    main(config)