-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain_BCT.py
143 lines (118 loc) · 4.66 KB
/
train_BCT.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
#
# Copyright (C) 2022 Apple Inc. All rights reserved.
#
from accelerate import Accelerator
from utils.getters import get_model, get_optimizer
from utils.net_utils import LabelSmoothing, backbone_to_torchscript
from dataset import SubImageFolder
from trainers import BCTTrainer
import pickle
import torch.nn as nn
import torch
from typing import Dict
from argparse import ArgumentParser
from collections import Counter
import yaml
from PIL import ImageFile
import tqdm
ImageFile.LOAD_TRUNCATED_IMAGES = True
def main(config: Dict) -> None:
"""Run training.
:param config: A dictionary with all configurations to run training.
:return:
"""
model = get_model(config.get('arch_params'))
accelerator = Accelerator()
device = accelerator.device
torch.backends.cudnn.benchmark = True
old_model = torch.jit.load(config.get('old_model_path')).eval().to(device)
if torch.cuda.is_available():
model = torch.nn.DataParallel(model)
old_model = torch.nn.DataParallel(old_model)
model.to(device)
alpha = float(config.get('alpha'))
print(f'alpha is {alpha}')
trainer = BCTTrainer()
optimizer = get_optimizer(model, **config.get('optimizer_params'))
data = SubImageFolder(**config.get('dataset_params'))
# lr_policy = get_policy(optimizer, **config.get('lr_policy_params'))
if config.get('label_smoothing') is None:
criterion = nn.CrossEntropyLoss()
else:
criterion = LabelSmoothing(smoothing=config.get('label_smoothing'))
optimizer, train_loader, val_loader =\
accelerator.prepare(optimizer, data.train_loader, data.val_loader)
old_model = accelerator.prepare(old_model)
print("==>Preparing pesudo classifier")
num_classes = int(config.get('arch_params')['num_classes'])
embedding_dim = int(config.get('arch_params')['embedding_dim'])
pseudo_classifier = torch.zeros(num_classes, embedding_dim)
old_embedding_dim = embedding_dim
if config.get('old_embedding_dim') is not None:
old_embedding_dim = int(config.get('old_embedding_dim'))
label_count = Counter()
for i, ((paths, images), target) in tqdm.tqdm(
enumerate(train_loader), ascii=True, total=len(train_loader)
):
images = images.to(device, non_blocking=True)
target = target.cpu()
with torch.no_grad():
_, features = old_model(images)
for feature, label in zip(features, target):
pseudo_classifier[int(
label)][:old_embedding_dim] += feature.flatten().cpu()
label_count.update([int(label)])
for i in range(num_classes):
pseudo_classifier[i] = pseudo_classifier[i]/label_count[i]
# old_classifier = old_model.fc.weight.data
# old_classifier = old_classifier.view(old_classifier.size(0),-1)
# pseudo_classifier[:500,:] = old_classifier[:500,:]
old_model = old_model.cpu()
# print(pseudo_classifier.size())
model = accelerator.prepare(model)
# Training loop
for epoch in range(config.get('epochs')):
# lr_policy(epoch, iteration=None)
train_acc1, train_acc5, train_loss = trainer.train(
train_loader=train_loader,
model=model,
alpha=alpha,
pseudo_classifier=pseudo_classifier,
criterion=criterion,
optimizer=optimizer,
device=device,
accelerator=accelerator
)
print(
"Train: epoch = {}, Loss = {}, Top 1 = {}, Top 5 = {}".format(
epoch, train_loss, train_acc1, train_acc5
))
test_acc1, test_acc5, test_loss = trainer.validate(
val_loader=val_loader,
model=model,
alpha=alpha,
pseudo_classifier=pseudo_classifier,
criterion=criterion,
device=device,
accelerator=accelerator
)
print(
"Test: epoch = {}, Loss = {}, Top 1 = {}, Top 5 = {}".format(
epoch, test_loss, test_acc1, test_acc5
))
if (epoch+1) % 5 == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': train_loss,
}, config.get('output_model_path')+f'.checkpoint')
# backbone_to_torchscript(model, config.get('output_model_path'))
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument('--config', type=str, required=True,
help='Path to config file for this pipeline.')
args = parser.parse_args()
with open(args.config) as f:
read_config = yaml.safe_load(f)
main(read_config)