def train(net, bins, alpha, beta, batch_size): """ params: bins: number of bins for classification alpha: regression loss weight beta: ortho loss weight """ # create model if net == "resnet50": model = ResNet(torchvision.models.resnet50(pretrained=True), num_classes=bins) lr = args.lr_resnet else: model = MobileNetV2(bins) lr = args.lr_mobilenet # loading data logger.logger.info("Loading data".center(100, '=')) train_data_loader = loadData(args.train_data, args.input_size, batch_size, bins) valid_data_loader = loadData(args.valid_data, args.input_size, batch_size, bins, False) # initialize cls loss function if args.cls_loss == "KLDiv": cls_criterion = nn.KLDivLoss(reduction='batchmean').cuda(0) elif args.cls_loss == "BCE": cls_criterion = nn.BCELoss().cuda(0) # initialize reg loss function reg_criterion = nn.MSELoss().cuda(0) softmax = nn.Softmax(dim=1).cuda(0) model.cuda(0) # training log logger.logger.info("Training".center(100, '=')) # initialize learning rate and step lr = lr step = 0 # validation error min_avg_error = 1000. # start training for epoch in range(args.epochs): print("Epoch:", epoch) # learning rate initialization if net == 'resnet50': if epoch >= args.unfreeze: optimizer = torch.optim.Adam( [{ "params": get_non_ignored_params(model, net), "lr": lr }, { "params": get_cls_fc_params(model), "lr": lr * 10 }], lr=args.lr_resnet) else: optimizer = torch.optim.Adam( [{ "params": get_non_ignored_params(model, net), "lr": lr }, { "params": get_cls_fc_params(model), "lr": lr * 10 }], lr=args.lr_resnet) else: if epoch >= args.unfreeze: optimizer = torch.optim.Adam( [{ "params": get_non_ignored_params(model, net), "lr": lr }, { "params": get_cls_fc_params(model), "lr": lr }], lr=args.lr_mobilenet) else: optimizer = torch.optim.Adam( [{ "params": get_non_ignored_params(model, net), "lr": lr * 10 }, { "params": get_cls_fc_params(model), "lr": lr * 10 }], lr=args.lr_mobilenet) # reduce lr by lr_decay factor for each epoch lr = lr * args.lr_decay print("------------") for i, (images, cls_v1, cls_v2, cls_v3, reg_v1, reg_v2, reg_v3, name) in enumerate(train_data_loader): step += 1 images = images.cuda(0) # get classified labels cls_v1 = cls_v1.cuda(0) cls_v2 = cls_v2.cuda(0) cls_v3 = cls_v3.cuda(0) # get continuous labels reg_v1 = reg_v1.cuda(0) reg_v2 = reg_v2.cuda(0) reg_v3 = reg_v3.cuda(0) # inference x_pred_v1, y_pred_v1, z_pred_v1, x_pred_v2, y_pred_v2, z_pred_v2, x_pred_v3, y_pred_v3, z_pred_v3 = model( images) logits = [ x_pred_v1, y_pred_v1, z_pred_v1, x_pred_v2, y_pred_v2, z_pred_v2, x_pred_v3, y_pred_v3, z_pred_v3 ] loss, degree_error_v1, degree_error_v2, degree_error_v3 = utils.computeLoss( cls_v1, cls_v2, cls_v3, reg_v1, reg_v2, reg_v3, logits, softmax, cls_criterion, reg_criterion, [ bins, alpha, beta, args.cls_loss, args.reg_loss, args.ortho_loss ]) # backward grad = [torch.tensor(1.0).cuda(0) for _ in range(3)] optimizer.zero_grad() torch.autograd.backward(loss, grad) optimizer.step() # save training log and weight if (i + 1) % 100 == 0: msg = "Epoch: %d/%d | Iter: %d/%d | x_loss: %.6f | y_loss: %.6f | z_loss: %.6f | degree_error_f:%.3f | degree_error_r:%.3f | degree_error_u:%.3f" % ( epoch, args.epochs, i + 1, len(train_data_loader.dataset) // batch_size, loss[0].item(), loss[1].item(), loss[2].item(), degree_error_v1.item(), degree_error_v2.item(), degree_error_v3.item()) logger.logger.info(msg) # Test on validation dataset error_v1, error_v2, error_v3 = valid(model, valid_data_loader, softmax, bins) print("Epoch:", epoch) print("Validation Error:", error_v1.item(), error_v2.item(), error_v3.item()) logger.logger.info("Validation Error(l,d,f)_{},{},{}".format( error_v1.item(), error_v2.item(), error_v3.item())) # save model if achieve better validation performance if error_v1.item() + error_v2.item() + error_v3.item() < min_avg_error: min_avg_error = error_v1.item() + error_v2.item() + error_v3.item() print("Training Info:") print("Model:", net, " ", "Number of bins:", bins, " ", "Alpha:", alpha, " ", "Beta:", beta) print("Saving Model......") torch.save( model.state_dict(), os.path.join(snapshot_dir, output_string + '_Best_' + '.pkl')) print("Saved")
class Test: def __init__(self, model_name, snapshot, num_classes): self.num_classes = num_classes if model_name == "resnet50": self.model = ResNet(torchvision.models.resnet50(pretrained=False), num_classes) elif model_name == "mobilenetv2": self.model = MobileNetV2(num_classes=num_classes) else: print("No such model...") exit(0) self.saved_state_dict = torch.load(snapshot) self.model.load_state_dict(self.saved_state_dict) self.model.cuda(0) self.model.eval() # Change model to 'eval' mode self.softmax = nn.Softmax(dim=1).cuda(0) self.optimizer = Optimize() def draw_vectors(self, pred_vector1, pred_vector2, pred_vector3, img, center, width): optimize_v = self.optimizer.Get_Ortho_Vectors(np.array(pred_vector1), np.array(pred_vector2), np.array(pred_vector3)) v1, v2, v3 = optimize_v[0], optimize_v[1], optimize_v[2] # draw vector in blue color predx, predy, predz = v1 utils.draw_front(img, predx, predy, width, tdx=center[0], tdy=center[1], size=100, color=(255, 0, 0)) # draw vector in green color predx, predy, predz = v2 utils.draw_front(img, predx, predy, width, tdx=center[0], tdy=center[1], size=100, color=(0, 255, 0)) # draw vector in red color predx, predy, predz = v3 utils.draw_front(img, predx, predy, width, tdx=center[0], tdy=center[1], size=100, color=(0, 0, 255)) cv.imshow("pose visualization", img) def test_per_img(self, cv_img, draw_img, center, w): with torch.no_grad(): images = cv_img.cuda(0) # get x,y,z cls predictions x_v1, y_v1, z_v1, x_v2, y_v2, z_v2, x_v3, y_v3, z_v3 = self.model( images) # get prediction vector(get continue value from classify result) _, _, _, pred_vector1 = utils.classify2vector( x_v1, y_v1, z_v1, self.softmax, self.num_classes) _, _, _, pred_vector2 = utils.classify2vector( x_v2, y_v2, z_v2, self.softmax, self.num_classes) _, _, _, pred_vector3 = utils.classify2vector( x_v3, y_v3, z_v3, self.softmax, self.num_classes) #visualize vectors self.draw_vectors(pred_vector1[0].cpu().tolist(), pred_vector2[0].cpu().tolist(), pred_vector3[0].cpu().tolist(), draw_img, center, w)
args = parse_args() utils.mkdir(args.save_dir) # cls and sord print("Creating model......") if args.model_name == "mobilenetv2": model = MobileNetV2(num_classes=args.num_classes) else: model = ResNet(torchvision.models.resnet50(pretrained=False), args.num_classes) print("Loading weight......") saved_state_dict = torch.load(args.snapshot) model.load_state_dict(saved_state_dict) model.cuda(0) model.eval() # Change model to 'eval' mode (BN uses moving mean/var). softmax = nn.Softmax(dim=1).cuda(0) # test dataLoader test_loader = loadData(args.test_data, args.input_size, args.batch_size, args.num_classes, False) # testing print('Start testing......') if args.collect_score: utils.mkdir(os.path.join(args.save_dir, "collect_score")) test(model, test_loader, softmax, args)
for i in range(4): label = labels[i] class_correct[label] += c[i].item() class_total[label] += 1 for i in range(10): print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i])) if __name__ == '__main__': # Loading and normalizing CIFAR10 trainloader, testloader, classes = data_prepare() # Define a Convolutional Neural Network net = ResNet() net = net.cuda() x = torch.autograd.Variable(torch.rand(64, 3, 32, 32)) writer.add_graph(net, x.cuda(), verbose=True) if args.load_model == True: net.load_state_dict(torch.load(args.model_path)) # Define a Loss function and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-08, weight_decay=0., amsgrad=True) # Train the network for epoch in range(args.epochs): train(net, criterion, optimizer, trainloader, epoch)
def train(img_dir, xml_dir, epochs, input_size, batch_size, num_classes): """ params: bins: number of bins for classification alpha: regression loss weight beta: ortho loss weight """ # create model model = ResNet(torchvision.models.resnet50(pretrained=True), num_classes=num_classes) cls_criterion = nn.CrossEntropyLoss().cuda(1) softmax = nn.Softmax(dim=1).cuda(1) model.cuda(1) # initialize learning rate and step lr = 0.001 step = 0 optimizer = torch.optim.Adam(model.parameters(), lr=lr) #load data train_data_loader = loadData(img_dir, xml_dir, input_size, batch_size, True) test_loader = loadData('../yolov3/data/test_imgs', '../yolov3/data/test_anns', 224, 8, False) #variables history = [] best_acc = 0.0 best_epoch = 0 # start training for epoch in range(epochs): print("Epoch:", epoch) print("------------") # reduce lr by lr_decay factor for each epoch if epoch % 10 == 0: lr = lr * 0.9 train_loss = 0.0 train_acc = 0 val_acc = 0 model.train() for i, (images, labels) in enumerate(train_data_loader): if i % 10 == 0: print("batch: {}/{}".format( i, len(train_data_loader.dataset) // batch_size)) images = images.cuda(1) labels = labels.cuda(1) # backward optimizer.zero_grad() outputs = model(images) loss = cls_criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() ret, predictions = torch.max(outputs.data, 1) correct_counts = predictions.eq(labels.data.view_as(predictions)) acc = torch.mean(correct_counts.type(torch.FloatTensor)) train_acc += acc.item() * images.size(0) print("epoch: {:03d}, Training loss: {:.4f}, Accuracy: {:.4f}%".format( epoch + 1, train_loss, train_acc / 3096 * 100)) #if (epoch+1) % 3 == 0: # torch.save(model, 'models/'+'model_'+str(epoch+1)+'.pt') print("Start testing...") with torch.no_grad(): model.eval() for j, (images, labels) in enumerate(test_loader): images = images.cuda(1) labels = labels.cuda(1) outputs = model(images) ret, preds = torch.max(outputs.data, 1) cnt = preds.eq(labels.data.view_as(preds)) acc = torch.mean(cnt.type(torch.FloatTensor)) val_acc += acc.item() * images.size(0) if val_acc > best_acc: print("correct testing samples:", val_acc) best_acc = val_acc torch.save(model, 'models/' + 'model_' + str(epoch + 1) + '.pt')