def run(pretrained_model): log(f' :: Start training with {pretrained_model}') net = RACNN(num_classes=200).cuda() net.load_state_dict(torch.load(pretrained_model)) cudnn.benchmark = True cls_params = list(net.b1.parameters()) + list(net.b2.parameters()) + list(net.b3.parameters()) + \ list(net.classifier1.parameters()) + list(net.classifier2.parameters()) + list(net.classifier3.parameters()) apn_params = list(net.apn1.parameters()) + list(net.apn2.parameters()) cls_opt = optim.SGD(cls_params, lr=0.001, momentum=0.9) apn_opt = optim.SGD(apn_params, lr=0.001, momentum=0.9) trainset = CUB200_loader('external/CUB_200_2011', split='train') testset = CUB200_loader('external/CUB_200_2011', split='test') trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, collate_fn=trainset.CUB_collate, num_workers=4) testloader = torch.utils.data.DataLoader(testset, batch_size=8, shuffle=False, collate_fn=testset.CUB_collate, num_workers=4) sample = random_sample(testloader) for epoch in range(50): cls_loss = train(net, trainloader, cls_opt, epoch, 'backbone') rank_loss = train(net, trainloader, apn_opt, epoch, 'apn') test(net, testloader) # visualize cropped inputs _, _, _, resized = net(sample.unsqueeze(0)) x1, x2 = resized[0].data, resized[1].data save_img( x1, path=f'build/.cache/epoch_{epoch}@2x.jpg', annotation=f'cls_loss = {cls_loss:.7f}, rank_loss = {rank_loss:.7f}' ) save_img( x2, path=f'build/.cache/epoch_{epoch}@4x.jpg', annotation=f'cls_loss = {cls_loss:.7f}, rank_loss = {rank_loss:.7f}' ) # save model per 10 epoches if epoch % 10 == 0 and epoch != 0: stamp = f'e{epoch}{int(time.time())}' torch.save(net.state_dict, f'build/racnn_mobilenetv2_cub200-e{epoch}s{stamp}.pt') log(f' :: Saved model dict as:\tbuild/racnn_mobilenetv2_cub200-e{epoch}s{stamp}.pt' ) torch.save(cls_opt.state_dict, f'build/cls_optimizer-s{stamp}.pt') torch.save(apn_opt.state_dict, f'build/apn_optimizer-s{stamp}.pt')
def run(): log(f'forge started') net = RACNN(num_classes=120).cuda() net.load_state_dict(torch.load('build/new_racnn_pretrained-1618222083.pt')) # load the model # net.load_state_dict(torch.load('build/racnn_mobilenetv2_cub200-e39se391618150738.pt')) # net.eval() cudnn.benchmark = True cls_params = list(net.b1.parameters()) + list(net.b2.parameters()) + list(net.b3.parameters()) + \ list(net.classifier1.parameters()) + list(net.classifier2.parameters()) + list(net.classifier3.parameters()) apn_params = list(net.apn1.parameters()) + list(net.apn2.parameters()) cls_opt = optim.SGD(cls_params, lr=0.001, momentum=0.9) apn_opt = optim.SGD(apn_params, lr=0.001, momentum=0.9) trainset = StanfordDogLoader('external/StanfordDog', split='train') testset = StanfordDogLoader('external/StanfordDog', split='test') trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, collate_fn=trainset.CUB_collate, num_workers=4) testloader = torch.utils.data.DataLoader(testset, batch_size=8, shuffle=False, collate_fn=testset.CUB_collate, num_workers=4) sample = random_sample(testloader) # load the optimizer # cls_opt.load_state_dict(torch.load('build/cls_optimizer-se391618150738.pt')) # apn_opt.load_state_dict(torch.load('build/apn_optimizer-se391618150738.pt')) # test(net, testloader) # print('loaded') for epoch in range(50): cls_loss = train(net, trainloader, cls_opt, epoch, 'backbone') rank_loss = train(net, trainloader, apn_opt, epoch, 'apn') test(net, testloader) # visualize cropped inputs _, _, _, resized = net(sample.unsqueeze(0)) x1, x2 = resized[0].data, resized[1].data save_img(x1, path=f'build/.cache/epoch_{epoch}@2x.jpg', annotation=f'cls_loss = {cls_loss:.7f}, rank_loss = {rank_loss:.7f}') save_img(x2, path=f'build/.cache/epoch_{epoch}@4x.jpg', annotation=f'cls_loss = {cls_loss:.7f}, rank_loss = {rank_loss:.7f}') # save model per 10 epoches if epoch % 5 == 4: stamp = f'e{epoch}{int(time.time())}' torch.save(net.state_dict(), f'build/new_racnn_mobilenetv2_-e{epoch}_{stamp}.pt') torch.save(net, f'build/main_net-e{epoch}_{stamp}.pt') log(f' :: Saved model dict') torch.save(cls_opt.state_dict(), f'build/new_cls_optimizer-s{stamp}.pt') torch.save(apn_opt.state_dict(), f'build/new_apn_optimizer-s{stamp}.pt')
def train(net, dataloader, optimizer, epoch, _type): assert _type in ['apn', 'backbone'] losses = 0 net.mode(_type), log(f' :: Switch to {_type}') # switch loss type for step, (inputs, targets) in enumerate(dataloader, 0): loss = net.echo(inputs, targets, optimizer) losses += loss if step % 20 == 0 and step != 0: avg_loss = losses/20 log(f':: loss @step({step:2d}/{len(dataloader)})-epoch{epoch}: {loss:.10f}\tavg_loss_20: {avg_loss:.10f}') losses = 0 return avg_loss
def test(net, dataloader): log(' :: Testing on test set ...') correct_summary = {'clsf-0': {'top-1': 0, 'top-5': 0}, 'clsf-1': {'top-1': 0, 'top-5': 0}, 'clsf-2': {'top-1': 0, 'top-5': 0}} for step, (inputs, labels) in enumerate(dataloader, 0): inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda() with torch.no_grad(): outputs, _, _, _ = net(inputs) for idx, logits in enumerate(outputs): correct_summary[f'clsf-{idx}']['top-1'] += torch.eq(logits.topk(max((1, 1)), 1, True, True)[1], labels.view(-1, 1)).sum().float().item() # top-1 correct_summary[f'clsf-{idx}']['top-5'] += torch.eq(logits.topk(max((1, 5)), 1, True, True)[1], labels.view(-1, 1)).sum().float().item() # top-5 if step > 200: for clsf in correct_summary.keys(): _summary = correct_summary[clsf] for topk in _summary.keys(): log(f'\tAccuracy {clsf}@{topk} ({step}/{len(dataloader)}) = {_summary[topk]/((step+1)*int(inputs.shape[0])):.5%}') return