DVS_Filtering/train.py at master · vision-agh/DVS_Filtering · 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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import lightning as L
import argparse
import multiprocessing as mp
import wandb

from lightning.pytorch.loggers.wandb import WandbLogger
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor

from omegaconf import OmegaConf

from data.ncaltech101.dataset import NCaltech101
from data.ncars.dataset import NCars
from data.nimagenet.dataset import NImageNet

from model.recognition import LNRecognition

def main(args):
    L.seed_everything(42, workers=True)

    cfg_dataset = OmegaConf.load(args.config_data)
    cfg_dataset.train.all_noisy = args.all_noisy

    if args.dataset is not None:
        cfg_dataset.path = args.dataset

    if 'cnn' in args.model:
        model_type = 'cnn'
        cfg_dataset.representation.type = 'event_frame'
        max_epochs = 50
    elif 'vit' in args.model:
        model_type = 'vit'
        cfg_dataset.representation.type = 'event_voxel'
        max_epochs = 50
    elif 'snn' in args.model:
        model_type = 'snn'
        cfg_dataset.representation.type = 'event_spikes'
        max_epochs = 50
    elif 'gcn' in args.model:
        model_type = 'gcn'
        cfg_dataset.representation.type = 'event_graph'
        max_epochs = 50

    cfg_model = cfg_dataset.model[model_type]

    print(cfg_dataset)
    print(cfg_model)

    if 'ncaltech' in args.config_data:
        dm = NCaltech101(cfg_dataset, cfg_model)
    elif 'ncars' in args.config_data:
        dm = NCars(cfg_dataset, cfg_model)
    elif 'nimagenet' in args.config_data:
        dm = NImageNet(cfg_dataset, cfg_model)

    dm.setup()

    model = LNRecognition(cfg_dataset, cfg_model)

    lr_monitor = LearningRateMonitor(logging_interval='step')

    folder_name = cfg_dataset.path.split('/')[-1]
    sub_folder_name = cfg_dataset.path.split('/')[-2]

    wandb_logger = WandbLogger(project=f'dvs_filtering',
                                group=f'{cfg_dataset.name}',
                                name=f'{cfg_dataset.name}_{cfg_dataset.representation.type}_{folder_name}_all_noise={args.all_noisy}')
    wandb_logger.watch(model)

    checkpoint_callback = ModelCheckpoint(
        dirpath='checkpoints',
        filename=f'{model_type}_{sub_folder_name}_{folder_name}_all_noise={args.all_noisy}',
        monitor='val_acc',
        mode='max',
        save_top_k=1
    )

    trainer = L.Trainer(max_epochs=max_epochs,
                        log_every_n_steps=1,
                        gradient_clip_val=1.0,
                        logger=wandb_logger,
                        callbacks=[lr_monitor, checkpoint_callback],
                        deterministic=True,
                        check_val_every_n_epoch=2,
                        devices=1)

    trainer.fit(model, dm)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-cd', '--config_data', type=str, default='configs/dataset/ncaltech.yaml')
    parser.add_argument('-d', '--dataset', type=str)
    parser.add_argument('-m', '--model', type=str, default='cnn')
    parser.add_argument('-all', '--all_noisy', type=bool, default=False)

    args = parser.parse_args()
    main(args)