def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
optimizer.zero_grad()
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.data.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum().item()
progress_bar(
batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
def print_training_result(self, batch_idx, n_batch, append=""):
if self.dataset_type == 'small':
progress_bar(batch_idx, n_batch,
"Loss: %.3f, Acc: %.3f%% | %s"
% (self.train_loss / (batch_idx + 1), self.top1.avg, append))
else:
raise NotImplementedError
def print_training_result(self, batch_idx, n_batch, monitor_freq=100, append=""):
if self.dataset_type == 'small':
progress_bar(batch_idx, n_batch, "Loss: %.3f (%.3f), Acc: %.3f%% (%.3f%%) | %s"
% (self.loss.val, self.loss.avg, self.top1.val, self.top1.avg, append))
else:
# raise NotImplementedError
if batch_idx % monitor_freq == 0:
print('Training: [%d / %d] \t Time %.3f (%.3f) \t Loss %.4f(%.4f)\n'
'Prec@1 %.4f(%.4f) \t Prec@5 %.4f(%.4f) \n'\
%(batch_idx, n_batch, self.batch_time.val, self.batch_time.sum,
self.loss_ImageNet.val, self.loss_ImageNet.avg,
self.top1.val, self.top1.avg, self.top5.val, self.top5.avg))
if append is not None:
print(append+'\n')
def train(
_net, _train_loader, _optimizer, _criterion, _device = 'cpu',
_recorder: Recorder = None,
_weight_quantization_error_collection = None, _input_quantization_error_collection = None,
_weight_bit_allocation_collection = None, _input_bit_allocation_collection = None
):
_net.train()
_train_loss = 0
_correct = 0
_total = 0
for batch_idx, (inputs, targets) in enumerate(_train_loader):
inputs, targets = inputs.to(_device), targets.to(_device)
_optimizer.zero_grad()
outputs = _net(inputs)
losses = _criterion(outputs, targets)
losses.backward()
_optimizer.step()
_train_loss += losses.data.item()
_, predicted = torch.max(outputs.data, 1)
_total += targets.size(0)
_correct += predicted.eq(targets.data).cpu().sum().item()
progress_bar(
batch_idx, len(_train_loader),
'Loss: %.3f | Acc: %.3f%% (%d/%d)' % (_train_loss / (batch_idx + 1), 100. * _correct / _total, _correct, _total)
)
if _recorder is not None:
_recorder.update(loss=losses.data.item(), acc=[_correct / _total], batch_size=inputs.size(0), is_train=True)
if _weight_quantization_error_collection and _input_quantization_error_collection is not None:
for name, layer in _net.quantized_layer_collections.items():
_weight_quantization_error = torch.abs(layer.quantized_weight - layer.pre_quantized_weight).mean().item()
_input_quantization_error = torch.abs(layer.quantized_input - layer.pre_quantized_input).mean().item()
_weight_quantization_error_collection[name].write('%.8e\n' % _weight_quantization_error)
_input_quantization_error_collection[name].write('%.8e\n' % _input_quantization_error)
if _weight_bit_allocation_collection and _input_bit_allocation_collection is not None:
for name, layer in _net.quantized_layer_collections.items():
_weight_bit_allocation_collection[name].write('%.2f\n' % (torch.abs(layer.quantized_weight_bit).mean().item()))
_input_bit_allocation_collection[name].write('%.2f\n' % (torch.abs(layer.quantized_input_bit).mean().item()))
return _train_loss / (len(_train_loader)), _correct / _total
def test(net, target_path, test_loader, use_cuda=True):
correct = 0
total = 0
net.eval()
for batch_idx, (inputs, targets) in enumerate(test_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
_, output = net(inputs, target_path)
_, predicted = torch.max(output, dim=1)
correct += predicted.eq(targets).cpu().sum().item()
total += targets.size(0)
progress_bar(batch_idx, len(test_loader),
"Acc: %.3f%%" % (100.0 * correct / total))
def test(epoch):
global best_acc
net.eval()
test_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(testloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
# inputs, targets = Variable(inputs, volatile=True), Variable(targets)
with torch.no_grad():
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.data.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum().item()
progress_bar(
batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(test_loss /
(batch_idx + 1), 100. * float(correct) / total, correct, total))
# Save checkpoint.
acc = 100. * correct / total
if acc > best_acc:
print('Saving..')
state = {
'net': net.module if use_cuda else net,
'acc': acc,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, './checkpoint/%s_ckpt.t7' % model_name)
best_acc = acc
if not os.path.exists('./%s' % model_name):
os.makedirs('./%s' % model_name)
torch.save(net.module.state_dict(),
'./%s/%s_pretrain.p' % (model_name, model_name))
def test(net, testloader):
net.eval()
test_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(testloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
with torch.no_grad():
outputs = net(inputs)
loss = nn.CrossEntropyLoss()(outputs, targets)
test_loss += loss.data.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum().item()
progress_bar(
batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(test_loss /
(batch_idx + 1), 100. * float(correct) / total, correct, total))
# losses.backward()
losses_s.backward()
losses_t.backward()
optimizer_s.step()
optimizer_t.step()
_, predicted = torch.max(out_s.data, dim=1)
correct_s += predicted.eq(y_s.data).cpu().sum().item()
_, predicted = torch.max(out_t.data, dim=1)
correct_t += predicted.eq(y_t.data).cpu().sum().item()
loss_t += losses_t.item()
total += y_s.size(0)
progress_bar(batch_idx, min(len(source_loader), len(target_loader)), "[Training] Source acc: %.3f%% | Target acc: %.3f%%"
%(100.0 * correct_s / total, 100.0 * correct_t / total))
#######################
# Record Training log #
#######################
source_recorder.update(loss=losses_s.item(), acc=accuracy(out_s.data, y_s.data, (1, 5)),
batch_size=out_s.shape[0], cur_lr=optimizer_s.param_groups[0]['lr'], end=end)
target_recorder.update(loss=losses_t.item(), acc=accuracy(out_t.data, y_t.data, (1, 5)),
batch_size=out_t.shape[0], cur_lr=optimizer_t.param_groups[0]['lr'], end=end)
# Test target acc
test_acc = mask_test(target_net, target_mask_dict, target_test_loader)
print('\n[Epoch %d] Test Acc: %.3f' % (epoch, test_acc))
target_recorder.update(loss=None, acc=test_acc, batch_size=0, end=None, is_train=False)
for epoch in range(n_epoch):
net.train()
loss = 0
correct = 0
total = 0
print('\n[Epoch: %d] \nTraining' % (epoch))
for batch_idx, (inputs, targets) in enumerate(target_train_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
outputs = net(inputs)
losses = nn.CrossEntropyLoss()(outputs, targets)
optimizer.zero_grad()
losses.backward()
optimizer.step()
loss += losses.item()
_, predicted = torch.max(outputs, dim=1)
correct += predicted.eq(targets.data).cpu().sum().item()
total += targets.size(0)
progress_bar(
batch_idx, len(target_train_loader), "Loss: %.3f | Acc: %.3f%%" %
(loss / (batch_idx + 1), 100.0 * correct / total))
print('Test')
validate(net, target_test_loader)
def test(net,
quantized_type,
test_loader,
use_cuda=True,
dataset_name='CIFAR10',
n_batches_used=None):
net.eval()
if dataset_name != 'ImageNet':
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(test_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
with torch.no_grad():
outputs = net(inputs, quantized_type)
_, predicted = torch.max(outputs.data, dim=1)
correct += predicted.eq(targets.data).cpu().sum().item()
total += targets.size(0)
progress_bar(batch_idx, len(test_loader),
"Test Acc: %.3f%%" % (100.0 * correct / total))
return 100.0 * correct / total
else:
batch_time = AverageMeter()
train_loss = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
with torch.no_grad():
end = time.time()
for batch_idx, (inputs, targets) in enumerate(test_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
outputs = net(inputs, quantized_type)
losses = nn.CrossEntropyLoss()(outputs, targets)
prec1, prec5 = accuracy(outputs.data,
targets.data,
topk=(1, 5))
train_loss.update(losses.data.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % 200 == 0:
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
batch_idx,
len(test_loader),
batch_time=batch_time,
loss=train_loss,
top1=top1,
top5=top5))
if n_batches_used is not None and batch_idx >= n_batches_used:
break
print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1,
top5=top5))
return top1.avg, top5.avg
def DNS_train(net,
CR,
optimizer,
train_loader,
validate_loader,
max_epoch=100,
criterion=nn.CrossEntropyLoss(),
save_path=None,
min_lr=1e-6,
max_descent_count=3,
use_cuda=True):
small_train_loss = 1e9
descend_count = 0
stop_flag = False
best_test_acc = 0
for epoch in range(max_epoch):
if stop_flag: break
net.train()
total = 0
correct = 0
train_loss = 0
for batch_idx, (inputs, targets) in enumerate(train_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
outputs = net(inputs, CR)
losses = criterion(outputs, targets)
optimizer.zero_grad()
losses.backward()
optimizer.step()
train_loss += losses.item()
_, predicted = torch.max(outputs, dim=1)
correct += predicted.eq(targets.data).cpu().sum().item()
total += targets.size(0)
test_acc = DNS_test(net, CR, validate_loader)
progress_bar(epoch, max_epoch, "Acc: %.3f%%" % test_acc)
if best_test_acc < test_acc:
best_test_acc = test_acc
if save_path is not None:
# print('Saving...')
try:
torch.save(net.module.state_dict(), save_path)
except:
torch.save(net.state_dict(), save_path)
if train_loss < small_train_loss:
small_train_loss = train_loss
descend_count = 0
else:
descend_count += 1
if descend_count >= max_descent_count:
descend_count = 0
optimizer.param_groups[0]['lr'] *= 0.1
if optimizer.param_groups[0]['lr'] < min_lr:
stop_flag = True
print('\nBest acc: %.3f' % best_test_acc)
break
return best_test_acc
# Calculate Classes Balance #
#############################
marginal_prob = torch.mean(output_t, dim=0)
cls_balance_losses = torch.sum(marginal_prob *
torch.log(marginal_prob + 1e-7))
optimizer.zero_grad()
# (losses_t + losses_s).backward()
losses = losses_s
# (losses_s + cond_entropy_param * cond_entropy + clas_balance_param * cls_balance_losses).backward()
losses += cond_entropy_param * cond_entropy
losses += clas_balance_param * cls_balance_losses
losses += MMD_param * MMD_losses
losses.backward()
optimizer.step()
# Get record
_, predicted = torch.max((logits_s).data, dim=1)
correct_s += predicted.eq(y_s).cpu().sum().item()
_, predicted = torch.max((logits_t).data, dim=1)
correct_t += predicted.eq(y_t).cpu().sum().item()
total += y_s.size(0)
progress_bar(
batch_idx, n_ite, "Src acc: %.3f%% | Tgt acc: %.3f%%" %
(100.0 * correct_s / total, 100.0 * correct_t / total))
print('\nTest with source data')
test(net, False, source_test_loader)
print('\nTest with target data')
test(net, True, target_test_loader)
# Training target task
_, output = net(inputs, target_path=True)
losses = nn.CrossEntropyLoss()(output, targets)
train_loss += losses.item()
optimizer.zero_grad()
losses.backward()
optimizer.step()
# Get record
_, predicted = torch.max(output.data, dim=1)
correct_t += predicted.eq(targets).cpu().sum().item()
total += targets.size(0)
progress_bar(
batch_idx, len(train_laoder), "Loss: %.3f | Acc: %.3f%%" %
(100.0 * train_loss / (batch_idx + 1), 100.0 * correct_t / total))
writer.add_scalar('Train/Loss', train_loss / (batch_idx + 1), niter)
writer.add_scalar('Train/Accuracy',
100.0 * train_loss / (batch_idx + 1), niter)
niter += 1
if train_loss < small_train_loss:
# if 100. * float(correct) / float(total) > best_train_acc:
small_train_loss = train_loss
best_train_acc = 100. * float(correct_t) / float(total)
descend_count = 0
else:
