def train(train_loader, augmentation_gpu, criterion, G_criterion, netG, netD, optimizer_g, optimizer_d, epoch, args,
writer):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses_g = AverageMeter('Loss_G', ':.4f')
losses_d = AverageMeter('Loss_D', ':.4f')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(
len(train_loader),
[batch_time, data_time, losses_g, losses_d, top1, top5],
prefix="Epoch: [{}]".format(epoch))
# switch to train mode
netG.train()
netD.train()
end = time.time()
for i, (video, audio) in enumerate(train_loader):
# measure data loading time
# print('========================================')
data_time.update(time.time() - end)
if args.gpu is not None:
video[0] = video[0].cuda(args.gpu, non_blocking=True)
video[1] = video[1].cuda(args.gpu, non_blocking=True)
video[0] = augmentation_gpu(video[0])
video[1] = augmentation_gpu(video[1])
im_q_fake = netG(video[0])
q_fake, q_real, output, target = netD(im_q_fake, im_q=video[0], im_k=video[1], t=args.temporal_decay)
set_requires_grad([netD], False) # Ds require no gradients when optimizing Gs
optimizer_g.zero_grad() # set generator's gradients to zero
loss_g = -100 * G_criterion(q_fake, q_real)
loss_g.backward(retain_graph=True)
set_requires_grad([netD], True)
optimizer_d.zero_grad() # set discriminator's gradients to zero
loss_d = criterion(output, target)
loss_d.backward()
optimizer_g.step() # update generator's weights
optimizer_d.step() # update discriminator's weights
# acc1/acc5 are (K+1)-way contrast classifier accuracy
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses_g.update(loss_g.item(), video[0].size(0))
losses_d.update(loss_d.item(), video[0].size(0))
top1.update(acc1[0], video[0].size(0))
top5.update(acc5[0], video[0].size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
if writer is not None:
total_iter = i + epoch * len(train_loader)
writer.add_scalar('moco_train/loss', loss_d, total_iter)
writer.add_scalar('moco_train/loss', loss_g, total_iter)
writer.add_scalar('moco_train/acc1', acc1, total_iter)
writer.add_scalar('moco_train/acc5', acc5, total_iter)
writer.add_scalar('moco_train_avg/loss_g', losses_g.avg, total_iter)
writer.add_scalar('moco_train_avg/loss_d', losses_d.avg, total_iter)
writer.add_scalar('moco_train_avg/acc1', top1.avg, total_iter)
writer.add_scalar('moco_train_avg/acc5', top5.avg, total_iter)
def validate(model, device, args, *, all_iters=None):
objs = AvgrageMeter()
top1 = AvgrageMeter()
top5 = AvgrageMeter()
loss_function = args.loss_function
val_dataprovider = args.val_dataprovider
model.eval()
max_val_iters = 250
t1 = time.time()
with torch.no_grad():
for _ in range(1, max_val_iters + 1):
data, target = val_dataprovider.next()
target = target.type(torch.LongTensor)
data, target = data.to(device), target.to(device)
output = model(data)
loss = loss_function(output, target)
prec1, prec5 = accuracy(output, target, topk=(1, 5))
n = data.size(0)
objs.update(loss.item(), n)
top1.update(prec1.item(), n)
top5.update(prec5.item(), n)
logInfo = 'TEST Iter {}: loss = {:.6f},\t'.format(all_iters, objs.avg) + \
'Top-1 err = {:.6f},\t'.format(1 - top1.avg / 100) + \
'Top-5 err = {:.6f},\t'.format(1 - top5.avg / 100) + \
'val_time = {:.6f}'.format(time.time() - t1)
logging.info(logInfo)
def train(model,
device,
args,
*,
val_interval,
bn_process=False,
all_iters=None):
optimizer = args.optimizer
loss_function = args.loss_function
scheduler = args.scheduler
train_dataprovider = args.train_dataprovider
t1 = time.time()
Top1_err, Top5_err = 0.0, 0.0
model.train()
for iters in range(1, val_interval + 1):
scheduler.step()
if bn_process:
adjust_bn_momentum(model, iters)
all_iters += 1
d_st = time.time()
data, target = train_dataprovider.next()
target = target.type(torch.LongTensor)
data, target = data.to(device), target.to(device)
data_time = time.time() - d_st
output = model(data)
loss = loss_function(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
prec1, prec5 = accuracy(output, target, topk=(1, 5))
Top1_err += 1 - prec1.item() / 100
Top5_err += 1 - prec5.item() / 100
if all_iters % args.display_interval == 0:
printInfo = 'TRAIN Iter {}: lr = {:.6f},\tloss = {:.6f},\t'.format(all_iters, scheduler.get_lr()[0], loss.item()) + \
'Top-1 err = {:.6f},\t'.format(Top1_err / args.display_interval) + \
'Top-5 err = {:.6f},\t'.format(Top5_err / args.display_interval) + \
'data_time = {:.6f},\ttrain_time = {:.6f}'.format(data_time, (time.time() - t1) / args.display_interval)
logging.info(printInfo)
t1 = time.time()
Top1_err, Top5_err = 0.0, 0.0
if all_iters % args.save_interval == 0:
save_checkpoint({
'state_dict': model.state_dict(),
}, all_iters)
return all_iters
def val_one_epoch(logger, val_loader, model, criterion, use_cuda):
model.eval()
losses = []
accs = []
model.eval()
for batch_idx, (inputs, targets) in enumerate(val_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
outputs = model(inputs)
loss = criterion(outputs, targets)
prec1, _ = accuracy(outputs.data, targets.data, topk=(1, 2))
losses.append(loss.item())
accs.append(prec1.item())
logger.val_losses.extend(losses)
logger.val_accs.extend(accs)
return np.mean(losses), np.mean(accs)
def validate(val_loader, model, criterion, args):
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(len(val_loader), [batch_time, losses, top1, top5],
prefix='Test: ')
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (video, audio, target) in enumerate(val_loader):
if args.gpu is not None:
video = video.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(video)
output = output.view(-1, args.clip_per_video, args.num_class)
target = target.view(-1, args.clip_per_video)
output = torch.mean(output, dim=1)
# make sure 10 clips belong to the same video
for j in range(1, args.clip_per_video):
assert all(target[:, 0] == target[:, j])
target = target[:, 0]
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), video.size(0))
top1.update(acc1[0], video.size(0))
top5.update(acc5[0], video.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
# TODO: this should also be done with the ProgressMeter
print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1,
top5=top5))
return losses.avg, top1.avg, top5.avg
def main():
parser = argparse.ArgumentParser(description='Voxelnet for semantic')
parser.add_argument('--lr',
default=0.001,
type=float,
help='Initial learning rate')
parser.add_argument('--epochs', default=100, help='epochs')
parser.add_argument('--batchsize', default=4, help='epochs')
parser.add_argument('--weight_file', default='', help='weights to load')
parser.add_argument(
'--test_area',
type=int,
default=5,
help='Which area to use for test, option: 1-6 [default: 6]')
parser.add_argument('--num_point',
type=int,
default=4096,
help='Point number [default: 4096]')
args = parser.parse_args()
NUM_POINT = args.num_point
BATCH_SIZE = args.batchsize
lr = args.lr
ALL_FILES = getDataFiles('indoor3d_sem_seg_hdf5_data/all_files.txt')
room_filelist = [
line.rstrip()
for line in open('indoor3d_sem_seg_hdf5_data/room_filelist.txt')
]
# Load ALL data
data_batch_list = []
label_batch_list = []
for h5_filename in ALL_FILES:
data_batch, label_batch = loadDataFile(h5_filename)
data_batch_list.append(data_batch)
label_batch_list.append(label_batch)
data_batches = np.concatenate(data_batch_list, 0)
label_batches = np.concatenate(label_batch_list, 0)
print(data_batches.shape)
print(label_batches.shape)
test_area = 'Area_' + str(args.test_area)
train_idxs = []
test_idxs = []
for i, room_name in enumerate(room_filelist):
if test_area in room_name:
test_idxs.append(i)
else:
train_idxs.append(i)
train_data = data_batches[
train_idxs, ...] # ... means ellipsis, the same as [train_idxs, :, :]
train_label = label_batches[train_idxs].astype(np.int64)
test_data = data_batches[test_idxs, ...]
test_label = label_batches[test_idxs].astype(np.int64)
print(train_data.shape, train_label.shape)
print(test_data.shape, test_label.shape)
time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
log_dir = os.path.join('log_ptn/train', test_area + '_' + time_string)
if not os.path.exists(log_dir): os.makedirs(log_dir)
checkpoint_dir = os.path.join(log_dir, 'checkpoint')
if not os.path.exists(checkpoint_dir): os.makedirs(checkpoint_dir)
#writer = SummaryWriter(log_dir=os.path.join( log_dir, 'tensorboard'))
start_epoch = 0
epochs = args.epochs
model = get_model()
model.cuda()
# print(model)
optimizer = torch.optim.Adam(model.parameters(), lr)
criterion = nn.CrossEntropyLoss().cuda()
if args.weight_file != '':
pre_trained_model = torch.load(args.weight_file)
start_epoch = pre_trained_model['epoch']
model_state = model.state_dict()
model_state.update(pre_trained_model['state_dict'])
model.load_state_dict(model_state)
global_counter = 0
for epoch in range(start_epoch, epochs):
learn_rate_now = adjust_learning_rate(optimizer, global_counter,
BATCH_SIZE, lr)
#writer.add_scalar('train/learning_rate', learn_rate_now, global_counter)
losses = AverageMeter()
top1 = AverageMeter()
model.train()
train_data_shuffled, train_label_shuffled, _ = shuffle_data(
train_data[:, 0:NUM_POINT, :], train_label)
file_size = train_data_shuffled.shape[0]
num_batches = file_size // BATCH_SIZE
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
feature = train_data_shuffled[start_idx:end_idx, :, :]
label = train_label_shuffled[start_idx:end_idx]
feature = np.expand_dims(feature, axis=1)
input = Variable(torch.from_numpy(feature).cuda(),
requires_grad=True)
input = torch.transpose(input, 3, 1)
target = Variable(torch.from_numpy(label).cuda(),
requires_grad=False)
target = target.view(-1, )
output = model(input)
output_reshaped = output.permute(0, 3, 2,
1).contiguous().view(-1, 13)
loss = criterion(output_reshaped, target)
prec1 = accuracy(output_reshaped.data, target.data, topk=(1, ))
#prec1[0] = prec1[0].cpu().numpy()[0]
prec1 = prec1[0].cpu().numpy()
#losses.update(loss.data[0], BATCH_SIZE)
losses.update(loss.data, BATCH_SIZE)
#top1.update(prec1[0], BATCH_SIZE)
top1.update(prec1, BATCH_SIZE)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch: [{0}][{1}]\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(epoch,
batch_idx,
loss=losses,
top1=top1))
with open(os.path.join(log_dir, 'train_log.txt'), 'a') as f:
f.write('Epoch: [{0}][{1}]\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) \n'.format(
epoch, batch_idx, loss=losses, top1=top1))
global_counter += 1
#writer.add_scalar('train/loss', losses.avg, global_counter)
#writer.add_scalar('train/accuracy', top1.avg, global_counter)
losses = AverageMeter()
top1 = AverageMeter()
model.eval()
file_size = test_data.shape[0]
num_batches = file_size // BATCH_SIZE
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
feature = test_data[start_idx:end_idx, :, :]
label = test_label[start_idx:end_idx]
feature = np.expand_dims(feature, axis=1)
input = Variable(torch.from_numpy(feature).cuda(),
requires_grad=True)
input = torch.transpose(input, 3, 1)
target = Variable(torch.from_numpy(label).cuda(),
requires_grad=False)
target = target.view(-1, )
output = model(input)
output_reshaped = output.permute(0, 3, 2,
1).contiguous().view(-1, 13)
loss = criterion(output_reshaped, target)
prec1 = accuracy(output_reshaped.data, target.data, topk=(1, ))
#prec1[0] = prec1[0].cpu().numpy()[0]
prec1 = prec1[0].cpu().numpy()
#losses.update(loss.data[0], BATCH_SIZE)
losses.update(loss.data, BATCH_SIZE)
#top1.update(prec1[0], BATCH_SIZE)
top1.update(prec1, BATCH_SIZE)
#writer.add_scalar('val/loss', losses.avg, global_counter)
#writer.add_scalar('val/accuracy', top1.avg, global_counter)
print('Epoch {} Val Loss {:.3f} Val Acc {:.3f} \t'.format(
epoch, losses.avg, top1.avg))
with open(os.path.join(log_dir, 'test_log.txt'), 'a') as f:
f.write('Epoch: [{0}]\t'
'Loss {loss.avg:.4f} \t'
'Prec@1 {top1.avg:.3f} \n'.format(epoch,
loss=losses,
top1=top1))
if (epoch % 5 == 0):
torch.save(
{
'epoch': epoch + 1,
'args': args,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
},
os.path.join(checkpoint_dir,
'checkpoint_' + str(epoch) + '.pth.tar'))
def evaluate(room_path, out_data_label_filename, out_gt_label_filename):
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
if args.visu:
fout = open(
os.path.join(DUMP_DIR,
os.path.basename(room_path)[:-4] + '_pred.obj'), 'w')
fout_gt = open(
os.path.join(DUMP_DIR,
os.path.basename(room_path)[:-4] + '_gt.obj'), 'w')
fout_data_label = open(out_data_label_filename, 'w')
fout_gt_label = open(out_gt_label_filename, 'w')
current_data, current_label = room2blocks_wrapper_normalized(
room_path, NUM_POINT)
current_data = current_data[:, 0:NUM_POINT, :].astype(np.float32)
current_label = np.squeeze(current_label).astype(np.int64)
# Get room dimension..
data_label = np.load(room_path)
data = data_label[:, 0:6]
max_room_x = max(data[:, 0])
max_room_y = max(data[:, 1])
max_room_z = max(data[:, 2])
file_size = current_data.shape[0]
num_batches = file_size // BATCH_SIZE
losses = AverageMeter()
top1 = AverageMeter()
model.eval()
for batch_idx in range(num_batches):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
cur_batch_size = end_idx - start_idx
feature = current_data[start_idx:end_idx, :, :]
label = current_label[start_idx:end_idx]
feature = np.expand_dims(feature, axis=1)
input = Variable(torch.from_numpy(feature).cuda(), requires_grad=True)
input = torch.transpose(input, 3, 1)
target = Variable(torch.from_numpy(label).cuda(), requires_grad=False)
target = target.view(-1, )
output = model(input)
output_reshaped = output.permute(0, 3, 2, 1).contiguous().view(-1, 13)
loss = criterion(output_reshaped, target)
prec1 = accuracy(output_reshaped.data, target.data, topk=(1, ))
prec1[0] = prec1[0].cpu().numpy()[0]
losses.update(loss.data[0], BATCH_SIZE)
top1.update(prec1[0], BATCH_SIZE)
pred_label = np.reshape(
np.argmax(output_reshaped.data.cpu().numpy(), axis=1),
(BATCH_SIZE, -1))
pred_val = np.reshape(output_reshaped.data.cpu().numpy(),
(BATCH_SIZE, -1, 13))
# Save prediction labels to OBJ file
for b in range(BATCH_SIZE):
pts = current_data[start_idx + b, :, :]
l = current_label[start_idx + b, :]
pts[:, 6] *= max_room_x
pts[:, 7] *= max_room_y
pts[:, 8] *= max_room_z
pts[:, 3:6] *= 255.0
pred = pred_label[b, :]
for i in range(NUM_POINT):
color = g_label2color[pred[i]]
color_gt = g_label2color[current_label[start_idx + b, i]]
if args.visu:
fout.write('v %f %f %f %d %d %d\n' %
(pts[i, 6], pts[i, 7], pts[i, 8], color[0],
color[1], color[2]))
fout_gt.write('v %f %f %f %d %d %d\n' %
(pts[i, 6], pts[i, 7], pts[i, 8],
color_gt[0], color_gt[1], color_gt[2]))
fout_data_label.write(
'%f %f %f %d %d %d %f %d\n' %
(pts[i, 6], pts[i, 7], pts[i, 8], pts[i, 3], pts[i, 4],
pts[i, 5], pred_val[b, i, pred[i]], pred[i]))
fout_gt_label.write('%d\n' % (l[i]))
correct = np.sum(pred_label == current_label[start_idx:end_idx, :])
total_correct += correct
total_seen += (cur_batch_size * NUM_POINT)
for i in range(start_idx, end_idx):
for j in range(NUM_POINT):
l = current_label[i, j]
total_seen_class[l] += 1
total_correct_class[l] += (pred_label[i - start_idx, j] == l)
print('eval accuracy: %f' % (total_correct / float(total_seen)))
fout_data_label.close()
fout_gt_label.close()
if args.visu:
fout.close()
fout_gt.close()
return total_correct, total_seen
def train(train_loader, model, criterion, optimizer, epoch, args, writer):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(len(train_loader),
[batch_time, data_time, losses, top1, top5],
prefix="Epoch: [{}]".format(epoch))
"""
Switch to eval mode:
Under the protocol of linear classification on frozen features/models,
it is not legitimate to change any part of the pre-trained model.
BatchNorm in train mode may revise running mean/std (even if it receives
no gradient), which are part of the model parameters too.
"""
model.eval()
end = time.time()
for i, (video, audio, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
video = video.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(video)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), video.size(0))
top1.update(acc1[0], video.size(0))
top5.update(acc5[0], video.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
if writer is not None:
total_iter = i + epoch * len(train_loader)
writer.add_scalar('lincls_train/loss', loss, total_iter)
writer.add_scalar('lincls_train/acc1', acc1, total_iter)
writer.add_scalar('lincls_train/acc5', acc5, total_iter)
writer.add_scalar('lincls_train_avg/lr',
optimizer.param_groups[0]['lr'], total_iter)
writer.add_scalar('lincls_train_avg/loss', losses.avg,
total_iter)
writer.add_scalar('lincls_train_avg/acc1', top1.avg,
total_iter)
writer.add_scalar('lincls_train_avg/acc5', top5.avg,
total_iter)