forked from occludedReIDICCV/HolisitcGuidanceOccReID
-
Notifications
You must be signed in to change notification settings - Fork 0
/
source_training.py
73 lines (61 loc) · 1.6 KB
/
source_training.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
import torchreid
import torch.nn as nn
source = 'market1501'
target = source
batch_size = 128
datamanager = torchreid.data.ImageDataManager(
root='reid-data',
sources=source,
targets=target,
height=384,
width=128,
batch_size_train=batch_size,
batch_size_test=100,
transforms=['random_flip', 'random_crop'],
num_instances=4,
train_sampler='RandomIdentitySampler',
load_train_targets=False
)
model = torchreid.models.build_model(
name='multigrain',
num_classes=datamanager.num_train_pids,
loss='triplet',
pretrained=True
)
#model = model.cuda()
model = nn.DataParallel(model).cuda() # Comment previous line and uncomment this line for multi-gpu use
optimizer = torchreid.optim.build_optimizer(
model,
optim='adam',
lr=0.0003
)
scheduler = torchreid.optim.build_lr_scheduler(
optimizer,
lr_scheduler='single_step',
stepsize=50
)
engine = torchreid.engine.ImageTripletAEEngine(
datamanager,
model,
optimizer=optimizer,
scheduler=scheduler,
label_smooth=True,
)
"""start_epoch = torchreid.utils.resume_from_checkpoint(
'log/source_training_parts_clean/market1501/model.pth.tar-25',
model,
optimizer,
)"""
#engine.run(
# test_only=True,
# dist_metric='cosine'
#)
engine.run(
save_dir='log/source_training_parts_clean/' + source, #_Duke_Maks_Based
max_epoch=80,
eval_freq=5,
print_freq=10,
test_only=False,
visrank=False,
start_epoch=0
)