def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
acc = 0.0
total = 0
if t_net and cfg.dis_feature:
hooks = f_distill.get_hooks()
# 自动混合精度
scaler = torch.cuda.amp.GradScaler()
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
with torch.cuda.amp.autocast(): # 自动混合精度 (pytorch1.6之后)
outputs = net(inputs)
if t_net:
loss = torch.cuda.FloatTensor(
[0]) if outputs.is_cuda else torch.Tensor([0])
with torch.no_grad():
teacher_outputs = t_net(inputs)
if cfg.dis_feature:
t_out = []
s_out = []
for k, v in f_distill.activation.items():
g, k, n = k.split("_")
# 一一配对feature, 进行loss 计算
if g == "s":
s_out.append(v)
else:
t_out.append(v)
# 选定的 feature 分别计算loss
fs_loss = fs_criterion(s_out, t_out)
loss += fs_loss
s_loss = criterion(outputs, targets)
do_loss = t_criterion(outputs, teacher_outputs)
loss += (s_loss * (1 - cfg.alpha) + do_loss * cfg.alpha)
else:
loss = criterion(outputs, targets)
# Scales loss. 放大梯度.
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
train_loss += loss.item()
total += targets.size(0)
acc += 1 - torch.mean(
torch.sum(
(outputs.detach() - targets.detach())**2, dim=0) / torch.sum(
(torch.mean(targets, dim=0) - targets)**2, dim=0))
progress_bar(
batch_idx, len(trainloader),
'Current lr: %f | Loss: %.5f | Acc: %.3f%% (%d)' %
(optimizer.state_dict()['param_groups'][0]['lr'], train_loss /
(batch_idx + 1), 100. * acc / (batch_idx + 1), total))
if t_net and cfg.dis_feature:
for hook in hooks:
hook.remove()
def test(epoch, test_loader, save=True):
global best_acc
net.eval()
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(test_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
outputs = net(inputs)
loss = criterion(outputs, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# timing
batch_time.update(time.time() - end)
end = time.time()
progress_bar(
batch_idx, len(test_loader),
'Loss: {:.3f} | Acc1: {:.3f}% | Acc5: {:.3f}%'.format(
losses.avg, top1.avg, top5.avg))
if save:
writer.add_scalar('loss/test', losses.avg, epoch)
writer.add_scalar('acc/test_top1', top1.avg, epoch)
writer.add_scalar('acc/test_top5', top5.avg, epoch)
is_best = False
if top1.avg > best_acc:
best_acc = top1.avg
is_best = True
print('Current best acc: {}'.format(best_acc))
save_checkpoint(
{
'epoch':
epoch,
'model':
args.model,
'dataset':
args.dataset,
'state_dict':
net.module.state_dict()
if isinstance(net, nn.DataParallel) else net.state_dict(),
'acc':
top1.avg,
'optimizer':
optimizer.state_dict(),
},
is_best,
checkpoint_dir=log_dir)
def test(epoch):
global best_acc
net.eval()
test_loss = 0
acc = 0
total = 0
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
with torch.no_grad():
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
total += targets.size(0)
acc += 1 - torch.mean(
torch.sum(
(outputs.detach() - targets.detach())**2, dim=0) /
torch.sum(
(torch.mean(targets, dim=0) - targets)**2, dim=0))
progress_bar(
batch_idx, len(testloader), 'Loss: %.5f | Acc: %.3f%% (%d)' %
(test_loss / (batch_idx + 1), 100. * acc /
(batch_idx + 1), total))
# Save checkpoint.
acc = 100. * acc / (batch_idx + 1)
if acc > best_acc:
print('Saving..')
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
'lr_scheduler': lr_scheduler,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(
state,
os.path.join(
cfg.save_checkpoint, "best_%s_%s_%dx%d.pth" %
(cfg.model, cfg.data_name, cfg.input_size[0],
cfg.input_size[1])))
best_acc = acc
def test(model_name, model_path, val_path, device='cuda'):
# create dataloader
transform = fast_transform()
testset = KeypointsDetDataSet(root=cfg.val_root,
input_size=cfg.input_size,
is_train=False,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers)
# loading model
net = build_network_by_name(model_name,
None,
num_classes=cfg.num_classes,
deploy=True)
model_info = torch.load(model_path)
# net.load_state_dict(model_info['net'])
net.load_state_dict(model_info)
net = net.to(device)
net.eval()
total = 0
acc = 0
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
with torch.no_grad():
outputs = net(inputs)
total += targets.size(0)
acc += 1 - torch.mean(
torch.sum(
(outputs.detach() - targets.detach())**2, dim=0) /
torch.sum(
(torch.mean(targets, dim=0) - targets)**2, dim=0))
progress_bar(
batch_idx, len(testloader), 'Acc(R square): %.3f%% (%d)' %
(100. * acc / (batch_idx + 1), total))
def visualize_nn_output(raw, player=0):
# utils.progress_bar(t,total_steps),t-trajectory_start)
length = 30
blank = ' ' * length
pref = blank if player else ''
for prob_rtp, piece_value, piece in zip(*raw):
pi = prob_rtp[:, :, piece]
unif = np.ones_like(pi) / (pi.shape[0] * pi.shape[1])
entropy = utils.entropy(pi)
max_entropy = utils.entropy(unif)
entr_vis = utils.progress_bar(entropy, max_entropy, length=length)
print(pref + entr_vis, piece_value[0, 0, piece].round(decimals=3))
def train(epoch, train_loader):
print('\nEpoch: %d' % epoch)
net.train()
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
for batch_idx, (inputs, targets) in enumerate(train_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# timing
batch_time.update(time.time() - end)
end = time.time()
progress_bar(
batch_idx, len(train_loader),
'Loss: {:.3f} | Acc1: {:.3f}% | Acc5: {:.3f}%'.format(
losses.avg, top1.avg, top5.avg))
writer.add_scalar('loss/train', losses.avg, epoch)
writer.add_scalar('acc/train_top1', top1.avg, epoch)
writer.add_scalar('acc/train_top5', top5.avg, epoch)
if render:
env.render()
if not fast:
time.sleep(s["time_elapsed_each_action"] / 1000)
#Reset the envs that reach terminal states
for i, d in enumerate(done):
if d:
#End current round, resets etc
assert n_envs == 1, "Scoreboard will not be correct unless n_envs == 1"
for p, dead in enumerate(env.envs[0].get_info()["is_dead"]):
if not dead:
game_score.declare_winner(name[p])
print(game_score.score_table(frac=frac))
print("{} Round ended. {} steps.".format(
utils.progress_bar(t, total_steps), t - trajectory_start))
if wait:
input('<enter> to start the next round')
env.reset(env=i)
#Prepare next round!
agent, name, weight = random_match(
all_agents, all_names,
all_weights) #change who's go it is!
game_score.set_current_players(name)
for a, w in zip(agent, weight):
if reload_weights:
if debug:
print("[*]agent loaded:", w, "(", a, ")")
else:
print("[*]", end='')
try:
def test(model_name,
model_path,
val_path,
device='cuda',
out_err="data/error"):
# class map
class_dict = {v: k for k, v in dict(enumerate(cfg.classes)).items()}
idx2classes = dict(enumerate(cfg.classes))
# create dataloader
transform_test = data_transform(False)
testset = ImageDataSet(root=cfg.val_root,
classes_dict=class_dict,
transform=transform_test,
is_train=False)
testloader = torch.utils.data.DataLoader(testset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=4)
net = build_network_by_name(model_name, None, num_classes=len(cfg.classes))
model_info = torch.load(model_path)
net.load_state_dict(model_info["net"])
net = net.to(device)
net.eval()
correct = 0
total = 0
with torch.no_grad():
for batch_idx, (inputs, targets, im_paths) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)
_, predicted = outputs.max(1)
total += targets.size(0)
result = predicted.eq(targets)
correct += result.sum().item()
# 保存预测错误的图片
if out_err:
if not os.path.exists(out_err):
os.makedirs(out_err)
im_paths = np.array(im_paths)
err_indexes = ~np.array(result.to('cpu'))
err_im_paths = im_paths[err_indexes]
right_result = [
idx2classes[i]
for i in np.array(targets.to('cpu'))[err_indexes]
]
err_result = [
idx2classes[i]
for i in np.array(predicted.to('cpu'))[err_indexes]
]
for (p, r, e) in zip(err_im_paths, right_result, err_result):
out_dir = os.path.join(out_err, e, r)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
shutil.copy(p, os.path.join(out_dir, os.path.basename(p)))
progress_bar(
batch_idx, len(testloader), 'Acc: %.3f%% (%d/%d)' %
(100. * correct / total, correct, total))