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')
