-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
211 lines (158 loc) · 6.64 KB
/
main.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
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
import sys, os, json
import argparse
import torch
from dassl.utils import setup_logger, set_random_seed, collect_env_info
from dassl.config import get_cfg_default
from dassl.engine import build_trainer
import datasets.kather
import datasets.pannuke
import datasets.digestpath
import datasets.wsss4luad
import datasets.covid
import datasets.rsna18
import datasets.mimic
import trainers.coop
import trainers.coop_medclip
import trainers.coop_biomedclip
import trainers.zsclip
from trainers.utils import save_results, print_results
def print_args(args, cfg):
print("***************")
print("** Arguments **")
print("***************")
optkeys = list(args.__dict__.keys())
optkeys.sort()
for key in optkeys:
print("{}: {}".format(key, args.__dict__[key]))
print("************")
print("** Config **")
print("************")
print(cfg)
def reset_cfg(cfg, args):
if args.root:
cfg.DATASET.ROOT = args.root
if args.output_dir:
cfg.OUTPUT_DIR = args.output_dir
if args.resume:
cfg.RESUME = args.resume
if args.seed:
cfg.SEED = args.seed
if args.source_domains:
cfg.DATASET.SOURCE_DOMAINS = args.source_domains
if args.target_domains:
cfg.DATASET.TARGET_DOMAINS = args.target_domains
if args.transforms:
cfg.INPUT.TRANSFORMS = args.transforms
if args.trainer:
cfg.TRAINER.NAME = args.trainer
if args.backbone:
cfg.MODEL.BACKBONE.NAME = args.backbone
if args.head:
cfg.MODEL.HEAD.NAME = args.head
def extend_cfg(cfg):
"""
Add new config variables.
E.g.
from yacs.config import CfgNode as CN
cfg.TRAINER.MY_MODEL = CN()
cfg.TRAINER.MY_MODEL.PARAM_A = 1.
cfg.TRAINER.MY_MODEL.PARAM_B = 0.5
cfg.TRAINER.MY_MODEL.PARAM_C = False
"""
from yacs.config import CfgNode as CN
cfg.TRAINER.COOP = CN()
cfg.TRAINER.COOP.N_CTX = 16 # number of context vectors
cfg.TRAINER.COOP.CSC = False # class-specific context
cfg.TRAINER.COOP.CTX_INIT = "" # initialization words
cfg.TRAINER.COOP.PREC = "fp16" # fp16, fp32, amp
cfg.TRAINER.COOP.CLASS_TOKEN_POSITION = "end"
cfg.TRAINER.COCOOP = CN()
cfg.TRAINER.COCOOP.N_CTX = 16 # number of context vectors
cfg.TRAINER.COCOOP.CTX_INIT = "" # initialization words
cfg.TRAINER.COCOOP.PREC = "fp16" # fp16, fp32, amp
cfg.DATASET.SUBSAMPLE_CLASSES = "all" # all, base or new
cfg.BACKDOOR = CN()
cfg.BACKDOOR.POISON_PERCENTAGE = 5
cfg.BACKDOOR.TARGET_CLASS = 0
cfg.BACKDOOR.NOISE_EPS = 8
cfg.BACKDOOR.PATCH_TYPE = "text" # "text", "random-patch", "image-patch"
cfg.BACKDOOR.POSITION = "bottom-left" # "top-left", "top-center", "top-right", "center-left", "center-center", "center-right", "bottom-left", "bottom-center", "bottom-right"
cfg.BACKDOOR.TRIGGER_SIZE = 24
cfg.BACKDOOR.TRIGGER_IMG_PATH = "<PATH>"
cfg.MODEL_NAME = CN()
cfg.MODEL_NAME = ""
cfg.MODEL_ROOT = CN()
cfg.MODEL_ROOT = ""
cfg.DATASET_NAME = CN()
cfg.DATASET_NAME = ""
cfg.DEVICE = CN()
cfg.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
cfg.DTYPE = CN()
cfg.DTYPE = "float32"
def setup_cfg(args):
cfg = get_cfg_default()
extend_cfg(cfg)
# 1. From the dataset config file
if args.dataset_config_file:
cfg.merge_from_file(args.dataset_config_file)
# 2. From the method config file
if args.config_file:
cfg.merge_from_file(args.config_file)
# 3. From input arguments
reset_cfg(cfg, args)
# 4. From optional input arguments
cfg.merge_from_list(args.opts)
cfg.freeze()
return cfg
def main(args):
cfg = setup_cfg(args)
if cfg.SEED >= 0:
print("Setting fixed seed: {}".format(cfg.SEED))
set_random_seed(cfg.SEED)
setup_logger(cfg.OUTPUT_DIR)
if torch.cuda.is_available() and cfg.USE_CUDA:
torch.backends.cudnn.benchmark = True
print_args(args, cfg)
print("Collecting env info ...")
print("** System info **\n{}\n".format(collect_env_info()))
trainer = build_trainer(cfg)
if args.eval_only:
trainer.load_model(args.model_dir, epoch=args.load_epoch)
num_samples = trainer.test_loader.dataset.num_samples
trainer.test_loader.dataset.backdoor_tags = torch.zeros(num_samples)
print(f"\nTest (CLEAN)")
results_clean = trainer.test()
trainer.test_loader.dataset.backdoor_tags = torch.ones(num_samples)
print(f"\nTest (BACKDOOR)")
results_backdoor = trainer.test()
results = [results_clean, results_backdoor]
print_results(cfg, results)
return
if not args.no_train:
results = trainer.train()
print_results(cfg, results)
var_name = "target_class"
var_value = str(cfg.BACKDOOR.TARGET_CLASS)
save_results(cfg, results, var_name, var_value)
return
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--root", type=str, default="", help="path to dataset")
parser.add_argument("--output-dir", type=str, default="", help="output directory")
parser.add_argument( "--resume", type=str, default="", help="checkpoint directory (from which the training resumes)")
parser.add_argument("--seed", type=int, default=-1, help="only positive value enables a fixed seed")
parser.add_argument("--source-domains", type=str, nargs="+", help="source domains for DA/DG")
parser.add_argument("--target-domains", type=str, nargs="+", help="target domains for DA/DG")
parser.add_argument("--transforms", type=str, nargs="+", help="data augmentation methods")
parser.add_argument("--config-file", type=str, default="", help="path to config file")
parser.add_argument("--dataset-config-file",type=str,default="",help="path to config file for dataset setup",)
parser.add_argument("--trainer", type=str, default="", help="name of trainer")
parser.add_argument("--backbone", type=str, default="", help="name of CNN backbone")
parser.add_argument("--head", type=str, default="", help="name of head")
parser.add_argument("--eval-only", action="store_true", help="evaluation only")
parser.add_argument("--model-dir",type=str,default="",help="load model from this directory for eval-only mode",)
parser.add_argument("--load-epoch", type=int, help="load model weights at this epoch for evaluation")
parser.add_argument("--no-train", action="store_true", help="do not call trainer.train()")
parser.add_argument("opts",default=None,nargs=argparse.REMAINDER,help="modify config options using the command-line",)
args = parser.parse_args()
main(args)