def main():
parser = argparse.ArgumentParser(description='Voxelnet for semantic')
parser.add_argument('--lr',
default=0.001,
type=float,
help='Initial learning rate') # default=0.001(good)
parser.add_argument('--epochs', default=2,
help='epochs') # default=100, 50, 30
parser.add_argument('--batchsize', default=4, help='epochs') # default=32
parser.add_argument('--weight_file', default='', help='weights to load')
# log_ptn/train/Area_2_2019-09-11-11-43-48/checkpoint/checkpoint_0_max_mIoU_test_25.17065278824228.pth.tar
parser.add_argument(
'--test_area',
type=int,
default=2,
help='Which area to use for test, option: 1-2 [default: 2]')
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') # .h5 file routes
room_filelist = [
line.rstrip()
for line in open('indoor3d_sem_seg_hdf5_data/room_filelist.txt')
]
# Load ALL data into a big data_batch & a big label_batch
data_batch_list = []
label_batch_list = []
print(ALL_FILES)
for h5_filename in ALL_FILES:
h5_dir = os.path.join(
'/home/chenkun/pointnet_pytorch-master/indoor3d_sem_seg_hdf5_data',
h5_filename)
f = h5py.File(h5_dir)
data_batch = f['data'][:]
label_batch = f['label'][:]
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, ...]
train_label = label_batches[train_idxs].astype(np.int64)
# test_data = data_batches[test_idxs, ...] # ZZC
# test_label = label_batches[test_idxs].astype(np.int64) # ZZC
test_data = train_data # ZZC
test_label = train_label # ZZC
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)
start_epoch = 0
epochs = args.epochs
model = get_model()
model.cuda()
# print(model)
optimizer = torch.optim.Adam(model.parameters(), lr)
# class_names = ["ground", "vegetation", "building", "clutter"] # ZZC
class_names = ["T2T", "B2B", "BH", "BL", "V2V", "OT"]
# Add weights to the loss function
# weightsTrain = [0.04, 0.20, 0.12, 0.64] # default
# weightsTrain = [0.25, 0.25, 0.25, 0.25]
# weightsTrain = [0.20, 0.50, 0.30, 0.50]
weightsTrain = [0.2, 0.4, 0.6, 1.00, 1.00, 1.00]
class_weights_Train = torch.FloatTensor(weightsTrain).cuda()
criterionTrain = nn.CrossEntropyLoss(weight=class_weights_Train,
size_average=True).cuda()
# True: loss is averaged over each loss element in batch
weightsVal = [0.2, 0.4, 0.6, 1.00, 1.00,
1.00] # default [0.08, 0.37, 0.15, 0.40]
class_weights_Val = torch.FloatTensor(weightsVal).cuda()
criterionVal = nn.CrossEntropyLoss(weight=class_weights_Val,
size_average=True).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)
# #####################################################
# Start training
# #####################################################
global_counter = 0
max_mIoU_test = 0.0
for epoch in range(start_epoch, epochs):
learn_rate_now = adjust_learning_rate(optimizer, global_counter,
BATCH_SIZE,
lr) # Seems not changing, ZZC
iter_loss = 0.0 # Initialisation: loss for one epoch
iterations = 0
cm = ConfusionMatrix(6, class_names=class_names)
cm.clear()
model.train()
train_data_shuffled, train_label_shuffled, _ = shuffle_data(
train_data[:, 0:NUM_POINT, :], train_label)
file_size = train_data_shuffled.shape[
0] # total number of training batches
num_batches = file_size // BATCH_SIZE # number of iterations in one epoch
print('\nnum_batches(training):\t', num_batches)
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]
# print('Here')
# print(feature.shape)
# print(label.shape)
# feature[:, :, 0:2] = 0.0
# feature[:, :, 6:9] = 0.0
# print(feature.shape)
# print(feature[0, 0, 0])
# print(feature[0, 0, 1])
# print(feature[0, 0, 2])
# print(feature[0, 0, 3])
# print(feature[0, 0, 4])
# print(feature[0, 0, 5])
# print(feature[0, 0, 6])
# print(feature[0, 0, 7])
# print(feature[0, 0, 8])
#
feature = np.expand_dims(feature, axis=1)
input = Variable(torch.from_numpy(feature).cuda(),
requires_grad=True)
# print(input.size())
input = torch.transpose(input, 3, 1) # ? ZZC
# print(input.size())
target = Variable(torch.from_numpy(label).cuda(),
requires_grad=False)
# print(target.size())
target = target.view(-1, )
# print(target.size())
output = model(input)
output_reshaped = output.permute(0, 3, 2,
1).contiguous().view(-1, 6)
# exit() # for check, ZZC
_, pred = torch.max(output.data, 1)
pred = pred.view(-1, )
cm.add_batch(target.cpu().numpy(), pred.cpu().numpy()) # detach()
loss = criterionTrain(output_reshaped, target)
iter_loss += loss.item() # Accumulate the loss
iterations += 1
optimizer.zero_grad()
loss.backward()
optimizer.step()
global_counter += 1
if batch_idx % 10 == 0:
print('Epoch: [%3d][%3d]\t Loss: %.4f' %
(epoch, batch_idx, loss)) # Print loss for one bath
# Print training results for 1 epoch
iou0, iou1, iou2, iou3, iou4, iou5, mIoU = cm.class_IoU()
print(
'Epoch: [%3d]\t Train Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%' %
(epoch, iter_loss / iterations, cm.overall_accuracy(),
mIoU)) # Print loss for the epoch
print(
'T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%'
% (iou0, iou1, iou2, iou3, iou4, iou5))
with open(os.path.join(log_dir, 'train_log.txt'), 'a') as f:
f.write(
'Epoch: [%3d]\t Train Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%\n'
% (epoch, iter_loss / iterations, cm.overall_accuracy(), mIoU))
f.write(
'T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%\n\n'
% (iou0, iou1, iou2, iou3, iou4, iou5))
# #####################################################
# Start validation
# #####################################################
model.eval()
iter_loss = 0.0 # Initialisation: loss for one epoch
iterations = 0
cm = ConfusionMatrix(6, class_names=class_names) # ZZC
cm.clear()
file_size = test_data.shape[0]
num_batches = file_size // BATCH_SIZE
print('num_batches(testing):\t', num_batches)
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[:, :, 0:2] = 0.0
# feature[:, :, 6:9] = 0.0
feature = np.expand_dims(feature, axis=1)
input = Variable(torch.from_numpy(feature).cuda(),
requires_grad=True)
input = torch.transpose(input, 3, 1) # ? ZZC
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, 6)
_, pred = torch.max(output.data, 1)
pred = pred.view(-1, )
cm.add_batch(target.cpu().numpy(), pred.cpu().numpy()) # detach()
loss = criterionVal(output_reshaped, target)
iter_loss += loss.item() # Accumulate the loss
iterations += 1
# Print validation results after 1 epoch
iou0, iou1, iou2, iou3, iou4, iou5, mIoU = cm.class_IoU()
print('Epoch: [%3d]\t Test Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%' %
(epoch, iter_loss / iterations, cm.overall_accuracy(),
mIoU)) # Print loss for the epoch
print(
'T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%'
% (iou0, iou1, iou2, iou3, iou4, iou5))
with open(os.path.join(log_dir, 'test_log.txt'), 'a') as f:
f.write(
'Epoch: [%3d]\t Test Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%\n'
% (epoch, iter_loss / iterations, cm.overall_accuracy(), mIoU))
f.write(
'T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%\n\n'
% (iou0, iou1, iou2, iou3, iou4, iou5))
# Check whether best model, -> Save model
if (mIoU > max_mIoU_test or epoch == epochs - 1):
max_mIoU_test = mIoU
print(
'-> Best performance (test mIoU) achieved or This is final epoch.'
)
print('Max_mIoU in testing: %3.2f%%\n' % (max_mIoU_test))
torch.save(
{
'epoch': epoch + 1,
'args': args,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
},
os.path.join(
checkpoint_dir, 'checkpoint_' + str(epoch) +
'_max_mIoU_test_' + str(mIoU) + '.pth.tar'))
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 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=50, help='epochs') # default=100
parser.add_argument('--batchsize', default=4, help='epochs') # default=32
parser.add_argument('--weight_file', default='', help='weights to load')
parser.add_argument('--test_area', type=int, default=2, help='Which area to use for test, option: 1-2 [default: 2]')
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 into a big data_batch & a big label_batch
data_batch_list = []
label_batch_list = []
print(ALL_FILES)
for h5_filename in ALL_FILES:
# print(h5_filename)
# data_batch, label_batch = loadDataFile(h5_filename)
h5_dir = os.path.join('/home/chenkun/pointnet_pytorch-master/indoor3d_sem_seg_hdf5_data', h5_filename)
f = h5py.File(h5_dir)
data_batch = f['data'][:]
label_batch = f['label'][:]
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, ...]
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)
# class_names = ["ground", "vegetation", "building", "clutter"] # ZZC
class_names = ["T2T", "B2B", "BH", "BL", "V2V", "OT"]
# Add weights to the loss function
# weightsTrain = [0.04, 0.20, 0.12, 0.64] # default
# weightsTrain = [0.25, 0.25, 0.25, 0.25]
# weightsTrain = [0.20, 0.50, 0.30, 0.50]
weightsTrain = [0.2, 0.4, 0.6, 1.00, 1.00, 1.00]
class_weights_Train = torch.FloatTensor(weightsTrain).cuda()
criterionTrain = nn.CrossEntropyLoss(weight=class_weights_Train,
size_average=True).cuda()
# True: loss is averaged over each loss element in batch
weightsVal = [0.2, 0.4, 0.6, 1.00, 1.00, 1.00] # default [0.08, 0.37, 0.15, 0.40]
class_weights_Val = torch.FloatTensor(weightsVal).cuda()
criterionVal = nn.CrossEntropyLoss(weight=class_weights_Val,
size_average=True).cuda()
#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
max_mIoU_test = 0.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
iter_loss = 0.0 # Initialisation: loss for one epoch
iterations = 0
cm = ConfusionMatrix(6, class_names=class_names)
cm.clear()
model.train()
train_data_shuffled, train_label_shuffled, _ = shuffle_data(train_data[:, 0:NUM_POINT, :], train_label)
file_size = train_data_shuffled.shape[0] # total number of training batches
num_batches = file_size // BATCH_SIZE # number of iterations in one epoch
print('\nnum_batches(training):\t', num_batches)
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, 6)
_, pred = torch.max(output.data, 1)
pred = pred.view(-1, )
cm.add_batch(target.cpu().numpy(), pred.cpu().numpy()) # detach()
loss = criterionTrain(output_reshaped, target)
iter_loss += loss.item() # Accumulate the loss
iterations += 1
loss = criterion(output_reshaped, target)
prec1 = accuracy(output_reshaped.data, target.data, topk=(1,))
prec1[0] = prec1[0].cpu().numpy()
losses.update(loss.item(), BATCH_SIZE)
top1.update(prec1[0], BATCH_SIZE)
optimizer.zero_grad()
loss.backward()
optimizer.step()
global_counter += 1
if batch_idx%10==0:
print('Epoch: [%3d][%3d]\t Loss: %.4f'%(epoch,batch_idx,loss)) # Print loss for one bath
# 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))
# Print training results for 1 epoch
iou0, iou1, iou2, iou3, iou4, iou5, mIoU = cm.class_IoU()
print('Epoch: [%3d]\t Train Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%' % (epoch, iter_loss / iterations, cm.overall_accuracy(), mIoU)) # Print loss for the epoch
print('T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%' % (iou0, iou1, iou2, iou3, iou4, iou5))
with open(os.path.join(log_dir, 'train_log.txt'), 'a') as f:
f.write('Epoch: [%3d]\t Train Loss: %.4f\t OA: %3.2f%%\t mIoU : %3.2f%%\n' % (epoch, iter_loss / iterations, cm.overall_accuracy(), mIoU))
f.write('T2T: %3.2f%%, B2B: %3.2f%%, BH: %3.2f%%, BL: %3.2f%%, V2V: %3.2f%%, OT: %3.2f%%\n\n' % (iou0, iou1, iou2, iou3, iou4, iou5))
#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()
losses.update(loss.item(), BATCH_SIZE)
top1.update(prec1[0], 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') )
writer.close()