CCFF/train.py at master · Undefinedx-Zhang/CCFF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import os
import json
import argparse
import torch
import dataloaders
import models
import math
from trainer import Trainer
import torch.nn.functional as F
from utils.losses import  CE_loss, Alg_loss

def main(config, resume, gpu, aug_type='all'):
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu
    torch.manual_seed(42)

    config['val_loader']['aug_type'] = aug_type
    method = config['model']['method'] + '_' + aug_type
    config['train_supervised']['aug_type'] = aug_type
    config['train_unsupervised']['aug_type'] = aug_type

    backbone = config['model']['backbone']
    config['experim_name'] = config['experim_name'].replace('method', method)
    config['experim_name'] = config['experim_name'].replace('percent', str(config['percent'] ))
    config['trainer']['save_dir'] = config['trainer']['save_dir'].replace('backbone', backbone)
    print (config)

    # DATA LOADERS
    config['train_supervised']['percnt_lbl'] = config["percent"]
    config['train_unsupervised']['percnt_lbl'] = config["percent"]
    supervised_loader = dataloaders.CDDataset(config['train_supervised'])
    unsupervised_loader = dataloaders.CDDataset(config['train_unsupervised'])
    print ('supervised: ', len(supervised_loader))
    print ('unsupervised: ', len(unsupervised_loader))
    val_loader = dataloaders.CDDataset(config['val_loader'])
    iter_per_epoch = len(unsupervised_loader)

    # SUPERVISED LOSS
    loss_l = CE_loss
    loss_alg = Alg_loss

    # MODEL
    if backbone=='resnet50':
        model = models.CCFF(num_classes=val_loader.dataset.num_classes,
                                         conf=config['model'],
                                         loss_l=loss_l,
                                         loss_alg=loss_alg,
                                         len_unsper=len(unsupervised_loader))

    print(f'\n{model}\n')

    # TRAINING
    trainer = Trainer(
        model=model,
        resume=resume,
        config=config,
        supervised_loader=supervised_loader,
        unsupervised_loader=unsupervised_loader,
        val_loader=val_loader,
        iter_per_epoch=iter_per_epoch)

    trainer.train()

if __name__=='__main__':
    # PARSE THE ARGS
    parser = argparse.ArgumentParser(description='PyTorch Training')
    parser.add_argument('-c', '--config', default='configs/config_WHU.json',type=str,
                        help='Path to the config file')
    parser.add_argument('-a', '--aug_type', default='all', type=str,
                        help='augmentation type')
    parser.add_argument('-r', '--resume', default=None, type=str,
                        help='Path to the .pth model checkpoint to resume training')
    parser.add_argument('-g', '--gpu', default="0", type=str,
                        help='indices of GPUs to enable (default: all)')
    parser.add_argument('--local', action='store_true', default=False)
    args = parser.parse_args()

    torch.backends.cudnn.benchmark = True

    config = json.load(open(args.config))
    gpu = args.gpu
    aug_type = args.aug_type
    main(config, args.resume, gpu, aug_type)