74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
from torchreid import metrics
|
|
from torchreid.engine import Engine
|
|
from torchreid.losses import CrossEntropyLoss
|
|
|
|
|
|
class ImageSoftmaxNASEngine(Engine):
|
|
|
|
def __init__(
|
|
self,
|
|
datamanager,
|
|
model,
|
|
optimizer,
|
|
scheduler=None,
|
|
use_gpu=False,
|
|
label_smooth=True,
|
|
mc_iter=1,
|
|
init_lmda=1.,
|
|
min_lmda=1.,
|
|
lmda_decay_step=20,
|
|
lmda_decay_rate=0.5,
|
|
fixed_lmda=False
|
|
):
|
|
super(ImageSoftmaxNASEngine, self).__init__(datamanager, use_gpu)
|
|
self.mc_iter = mc_iter
|
|
self.init_lmda = init_lmda
|
|
self.min_lmda = min_lmda
|
|
self.lmda_decay_step = lmda_decay_step
|
|
self.lmda_decay_rate = lmda_decay_rate
|
|
self.fixed_lmda = fixed_lmda
|
|
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
self.scheduler = scheduler
|
|
self.register_model('model', model, optimizer, scheduler)
|
|
|
|
self.criterion = CrossEntropyLoss(
|
|
num_classes=self.datamanager.num_train_pids,
|
|
use_gpu=self.use_gpu,
|
|
label_smooth=label_smooth
|
|
)
|
|
|
|
def forward_backward(self, data):
|
|
imgs, pids = self.parse_data_for_train(data)
|
|
|
|
if self.use_gpu:
|
|
imgs = imgs.cuda()
|
|
pids = pids.cuda()
|
|
|
|
# softmax temporature
|
|
if self.fixed_lmda or self.lmda_decay_step == -1:
|
|
lmda = self.init_lmda
|
|
else:
|
|
lmda = self.init_lmda * self.lmda_decay_rate**(
|
|
self.epoch // self.lmda_decay_step
|
|
)
|
|
if lmda < self.min_lmda:
|
|
lmda = self.min_lmda
|
|
|
|
for k in range(self.mc_iter):
|
|
outputs = self.model(imgs, lmda=lmda)
|
|
loss = self.compute_loss(self.criterion, outputs, pids)
|
|
self.optimizer.zero_grad()
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
|
|
loss_dict = {
|
|
'loss': loss.item(),
|
|
'acc': metrics.accuracy(outputs, pids)[0].item()
|
|
}
|
|
|
|
return loss_dict
|