-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
67 lines (52 loc) · 1.72 KB
/
train.py
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
#!/usr/bin/env python3
"""
DAP: Generating Instance-level Prompts for Rehearsal-free Continual Learning (ICCV 2023 Oral)
main function
"""
import os
import torch
import warnings
import numpy as np
import random
from src.configs.config import get_cfg
from src.data import loader as data_loader
from src.engine.trainer import Trainer
from src.models.build_model import build_model
from src.utils.file_io import PathManager
from launch import default_argument_parser
warnings.filterwarnings("ignore")
def setup(args):
cfg = get_cfg()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
output_path = os.path.join(cfg.OUTPUT_DIR, cfg.DATA.NAME)
if PathManager.exists(output_path):
raise ValueError(f"Already run for {output_path}")
PathManager.mkdirs(output_path)
cfg.OUTPUT_DIR = output_path
return cfg
def get_datasets(cfg):
print("Loading training data...")
train_dataset = data_loader._construct_dataset(cfg, split='train')
print("Loading test data...")
test_dataset = data_loader._construct_dataset(cfg, split='test')
return train_dataset, test_dataset
def train(cfg):
if torch.cuda.is_available():
torch.cuda.empty_cache()
if cfg.SEED is not None:
torch.manual_seed(cfg.SEED)
np.random.seed(cfg.SEED)
random.seed(0)
train_dataset, test_dataset = get_datasets(cfg)
print("Constructing models...")
model, cur_device = build_model(cfg)
print("Setting up trainer...")
trainer = Trainer(cfg, model, cur_device)
trainer.train_classifier(train_dataset, test_dataset)
def main(args):
cfg = setup(args)
train(cfg)
if __name__ == '__main__':
args = default_argument_parser().parse_args()
main(args)