-
Notifications
You must be signed in to change notification settings - Fork 1
/
fed_dacs.py
182 lines (162 loc) · 7.38 KB
/
fed_dacs.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
import argparse
import os.path as osp
import random
import numpy as np
import time
from datetime import timedelta
import torch
from torch import nn
from torch.backends import cudnn
from reid import models
from reid.server import FedDomainMemoTrainer
from reid.evaluators import Evaluator
from reid.utils.serialization import load_checkpoint, save_checkpoint
from reid.utils.tools import get_test_loader, get_data
from reid import datasets
start_epoch = best_mAP = 0
def create_model(args, num_cls=0):
# we only use triplet loss, remember to turn off 'norm'
model = models.create(
args.arch, num_features=args.features, norm=False,
dropout=args.dropout, num_classes=num_cls
)
# use CUDA
model = model.cuda()
model = nn.DataParallel(model) if args.is_parallel else model
return model
def main():
args = parser.parse_args()
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
main_worker(args)
def main_worker(args):
global start_epoch, best_mAP
start_time = time.monotonic()
cudnn.benchmark = True
all_datasets = datasets.names()
test_set_name = args.test_dataset
all_datasets.remove(test_set_name)
if args.exclude_dataset is not '':
exclude_set_name = args.exclude_dataset.split(',')
[all_datasets.remove(name) for name in exclude_set_name]
train_sets_name = sorted(all_datasets)
print("==========\nArgs:{}\n==========".format(args))
# Create datasets
print("==> Building Datasets")
test_set = get_data(args)
test_loader = get_test_loader(test_set, args.height, args.width,
args.batch_size, args.workers)
train_sets = get_data(args, train_sets_name)
num_users = len(train_sets)
# Create model
model = create_model(args)
# sub models on different servers
sub_models = [create_model(args) for key in range(num_users)]
aug_mods = [
models.create('aug', num_features=3, width=args.width, height=args.height).cuda()
for idx in range(num_users)
]
# Evaluator
evaluator = Evaluator(model)
trainer = FedDomainMemoTrainer(args, train_sets, model)
if args.resume:
checkpoint = load_checkpoint(args.resume)
for idx in range(num_users):
sub_models[idx].load_state_dict(checkpoint['sub_models'][idx])
trainer.classifier[idx].load_state_dict(checkpoint['cls_params'][idx])
start_epoch = checkpoint['epoch'] - 1
model.load_state_dict(checkpoint['state_dict'])
if args.evaluate:
evaluator.evaluate(test_loader, test_set.query, test_set.gallery, cmc_flag=True)
return
# start training
for epoch in range(start_epoch, args.epochs): # number of epochs
w_locals = []
torch.cuda.empty_cache()
for index in range(num_users):
# avg model boosted
w = trainer.train_dacs(
sub_models[index], model, aug_mods[index],
epoch, index, op_type='sgd'
)
w_locals.append(w)
# update server-side global model
w_global = trainer.fed_avg(w_locals)
model.load_state_dict(w_global)
# evaluate, no exchange, do not modify local models (sub_models)
if epoch % args.eval_step == 0:
cur_map, rank1 = evaluator.evaluate(test_loader, test_set.query,
test_set.gallery, cmc_flag=True)
# save
if cur_map > best_mAP:
print('best model saved!')
save_checkpoint({
'state_dict': w_global,
# 'cls_params': [cls_layer.state_dict() for cls_layer in trainer.classifier],
# 'sub_models': [mod.state_dict() for mod in sub_models],
'epoch': epoch + 1, 'best_mAP': best_mAP,
}, 1, fpath=osp.join(args.logs_dir, f'checkpoint_{epoch}.pth.tar'))
best_mAP = cur_map
# # save some transformed samples
# trainer.save_images(aug_mods[train_sets_name.index('msmt17')], epoch)
# # save tsne
# name_list = [cur_set.__class__.__name__ for cur_set in train_sets]
# trainer.save_tsne(
# [os.path.join(args.logs_dir, f"{pth_name}.pth") for pth_name in name_list], model,
# aug_mods[train_sets_name.index('msmt17')], epoch
# )
end_time = time.monotonic()
print('Total running time: ', timedelta(seconds=end_time - start_time))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Domain-level Fed Learning")
# data
parser.add_argument('-td', '--test-dataset', type=str, default='market1501',
choices=datasets.names())
parser.add_argument('-ed', '--exclude-dataset', type=str, default='')
parser.add_argument('-b', '--batch-size', type=int, default=64)
parser.add_argument('-j', '--workers', type=int, default=8)
parser.add_argument('--height', type=int, default=256, help="input height")
parser.add_argument('--width', type=int, default=128, help="input width")
# model
parser.add_argument('-a', '--arch', type=str, default='resnet50',
choices=models.names())
parser.add_argument('--features', type=int, default=0)
parser.add_argument('--dropout', type=float, default=0)
parser.add_argument('--lam', type=float, default=1)
# optimizer
parser.add_argument('--temp', type=float, default=0.05, help="temperature")
parser.add_argument('--rho', type=float, default=0.05, help="rho")
parser.add_argument('--momentum', type=float, default=0.9,
help="momentum to update model")
parser.add_argument('--weight-decay', type=float, default=5e-4)
parser.add_argument('--warmup-step', type=int, default=10)
parser.add_argument('--milestones', nargs='+', type=int,
default=[20, 30], help='milestones for the learning rate decay')
parser.add_argument('--lr', type=float, default=1e-3,
help="learning rate")
parser.add_argument('--epochs', type=int, default=41)
parser.add_argument('--max-iter', type=int, default=200)
parser.add_argument('--num-workers', type=int, default=8)
parser.add_argument('--num-instances', type=int, default=4,
help="each minibatch consist of "
"(batch_size // num_instances) identities, and "
"each identity has num_instances instances, "
"default: 4")
# training configs
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--print-freq', type=int, default=10)
parser.add_argument('--eval-step', type=int, default=10)
# path
working_dir = osp.dirname(osp.abspath(__file__))
parser.add_argument('--data-dir', type=str, metavar='PATH',
default=osp.join(working_dir, 'data'))
parser.add_argument('--logs-dir', type=str, metavar='PATH',
default=osp.join(working_dir, 'logs'))
parser.add_argument('--resume', type=str, default='')
parser.add_argument('--evaluate', action='store_true',
help="evaluation only")
parser.add_argument('--is_parallel', type=int, default=1)
main()