def main():
global args, best_result, output_directory, train_csv, test_csv
args = parser.parse_args()
if args.modality == 'rgb' and args.num_samples != 0:
print("number of samples is forced to be 0 when input modality is rgb")
args.num_samples = 0
if args.modality == 'rgb' and args.max_depth != 0.0:
print("max depth is forced to be 0.0 when input modality is rgb/rgbd")
args.max_depth = 0.0
sparsifier = None
max_depth = args.max_depth if args.max_depth >= 0.0 else np.inf
if args.sparsifier == UniformSampling.name:
sparsifier = UniformSampling(num_samples=args.num_samples,
max_depth=max_depth)
elif args.sparsifier == SimulatedStereo.name:
sparsifier = SimulatedStereo(num_samples=args.num_samples,
max_depth=max_depth)
# create results folder, if not already exists
output_directory = os.path.join(
'/media/kuowei/8EB89C8DB89C7585/results_CS',
'{}'.format(args.outputdir),
'{}.sparsifier={}.modality={}.arch={}.decoder={}.criterion={}.lr={}.bs={}'
.format(args.data, sparsifier, args.modality, args.arch, args.decoder,
args.criterion, args.lr, args.batch_size))
if not os.path.exists(output_directory):
os.makedirs(output_directory)
train_csv = os.path.join(output_directory, 'train.csv')
test_csv = os.path.join(output_directory, 'test.csv')
best_txt = os.path.join(output_directory, 'best.txt')
# define loss function (criterion) and optimizer
if args.criterion == 'l2':
criterion = criteria.MaskedMSELoss().cuda()
elif args.criterion == 'l1':
criterion = criteria.MaskedL1Loss().cuda()
out_channels = 1
# Data loading code
print("=> creating data loaders ...")
traindir = os.path.join('/media/kuowei/8EB89C8DB89C7585/data', args.data,
'train')
valdir = os.path.join('/media/kuowei/8EB89C8DB89C7585/data', args.data,
'val')
# traindir = os.path.join('data', args.data, 'train')
# valdir = os.path.join('data', args.data, 'val')
# if args.data == 'kitti':
# pass
# rgb_dir = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/train/rgb/'
# sparse_depth_dir = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/train/sd/'
# continuous_depth_dir = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/train/cd/'
# ground_dir = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/train/ground/'
# train_dataset = RgbdDataset(rgb_dir, sparse_depth_dir, continuous_depth_dir, ground_dir)
# train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None))
# # set batch size to be 1 for validation
# rgb_dir_val = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/validate/rgb/'
# sparse_depth_dir_val = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/validate/sd/'
# continuous_depth_dir_val = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/validate/cd/'
# ground_dir_val = '/media/kuowei/c9cb78ce-3109-4880-adad-b628c4261d82/rgb/validate/ground/'
# val_dataset = RgbdDataset(rgb_dir_val, sparse_depth_dir_val, continuous_depth_dir_val, ground_dir_val)
# val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True)
# elif args.data == 'nyudepthv2':
train_dataset = NYUDataset(traindir,
type='train',
modality=args.modality,
sparsifier=sparsifier)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None)
# set batch size to be 1 for validation
val_dataset = NYUDataset(valdir,
type='val',
modality=args.modality,
sparsifier=sparsifier)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print("=> data loaders created.")
# evaluation mode
if args.evaluate:
best_model_filename = os.path.join(output_directory,
'model_best.pth.tar')
if os.path.isfile(best_model_filename):
print("=> loading best model '{}'".format(best_model_filename))
checkpoint = torch.load(best_model_filename)
args.start_epoch = checkpoint['epoch']
best_result = checkpoint['best_result']
model = checkpoint['model']
print("=> loaded best model (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no best model found at '{}'".format(best_model_filename))
validate(val_loader, model, checkpoint['epoch'], write_to_file=False)
return
# optionally resume from a checkpoint
elif args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
best_result = checkpoint['best_result']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
return
# create new model
else:
# define model
print("=> creating Model ({}-{}) ...".format(args.arch, args.decoder))
in_channels = len(args.modality)
if args.arch == 'resnet50':
model = ResNet(layers=50,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
elif args.arch == 'resnet18':
model = ResNet(layers=18,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
elif args.arch == 'resnet152':
model = ResNet(layers=152,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
print("=> model created.")
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# create new csv files with only header
with open(train_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
with open(test_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
print(model)
print("=> model transferred to GPU.")
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
result, img_merge = validate(val_loader, model, epoch)
# remember best rmse and save checkpoint
is_best = result.rmse < best_result.rmse
if is_best:
best_result = result
with open(best_txt, 'w') as txtfile:
txtfile.write(
"epoch={}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n"
.format(epoch, result.mse, result.rmse, result.absrel,
result.lg10, result.mae, result.delta1,
result.gpu_time))
if img_merge is not None:
img_filename = output_directory + '/comparison_best.png'
utils.save_image(img_merge, img_filename)
save_checkpoint(
{
'epoch': epoch,
'arch': args.arch,
'model': model,
'best_result': best_result,
'optimizer': optimizer,
}, is_best, epoch)
def main():
global args, best_result, output_directory, train_csv, test_csv
sparsifier = None
max_depth = args.max_depth if args.max_depth >= 0.0 else np.inf
if args.sparsifier == UniformSampling.name:
sparsifier = UniformSampling(num_samples=args.num_samples,
max_depth=max_depth)
elif args.sparsifier == SimulatedStereo.name:
sparsifier = SimulatedStereo(num_samples=args.num_samples,
max_depth=max_depth)
# create results folder, if not already exists
output_directory = utils.get_output_directory(args)
if not os.path.exists(output_directory):
os.makedirs(output_directory)
train_csv = os.path.join(output_directory, 'train.csv')
test_csv = os.path.join(output_directory, 'test.csv')
best_txt = os.path.join(output_directory, 'best.txt')
# define loss function (criterion) and optimizer
if args.criterion == 'l2':
criterion = criteria.MaskedMSELoss().cuda()
elif args.criterion == 'l1':
criterion = criteria.MaskedL1Loss().cuda()
out_channels = 1
# Data loading code
print("=> creating data loaders ...")
traindir = os.path.join('data', args.data, 'train')
valdir = os.path.join('data', args.data, 'val')
train_dataset = NYUDataset(traindir,
type='train',
modality=args.modality,
sparsifier=sparsifier)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None)
# set batch size to be 1 for validation
val_dataset = NYUDataset(valdir,
type='val',
modality=args.modality,
sparsifier=sparsifier)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print("=> data loaders created.")
# evaluation mode
if args.evaluate:
best_model_filename = os.path.join(output_directory,
'model_best.pth.tar')
assert os.path.isfile(best_model_filename), \
"=> no best model found at '{}'".format(best_model_filename)
print("=> loading best model '{}'".format(best_model_filename))
checkpoint = torch.load(best_model_filename)
args.start_epoch = checkpoint['epoch']
best_result = checkpoint['best_result']
model = checkpoint['model']
print("=> loaded best model (epoch {})".format(checkpoint['epoch']))
validate(val_loader, model, checkpoint['epoch'], write_to_file=False)
return
# optionally resume from a checkpoint
elif args.resume:
assert os.path.isfile(args.resume), \
"=> no checkpoint found at '{}'".format(args.resume)
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
best_result = checkpoint['best_result']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
# create new model
else:
# define model
print("=> creating Model ({}-{}) ...".format(args.arch, args.decoder))
in_channels = len(args.modality)
if args.arch == 'resnet50':
model = ResNet(layers=50,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
elif args.arch == 'resnet18':
model = ResNet(layers=18,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
print("=> model created.")
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# create new csv files with only header
with open(train_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
with open(test_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
print(model)
print("=> model transferred to GPU.")
for epoch in range(args.start_epoch, args.epochs):
utils.adjust_learning_rate(optimizer, epoch, args.lr)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
result, img_merge = validate(val_loader, model, epoch)
# remember best rmse and save checkpoint
is_best = result.rmse < best_result.rmse
if is_best:
best_result = result
with open(best_txt, 'w') as txtfile:
txtfile.write(
"epoch={}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n"
.format(epoch, result.mse, result.rmse, result.absrel,
result.lg10, result.mae, result.delta1,
result.gpu_time))
if img_merge is not None:
img_filename = output_directory + '/comparison_best.png'
utils.save_image(img_merge, img_filename)
utils.save_checkpoint(
{
'args': args,
'epoch': epoch,
'arch': args.arch,
'model': model,
'best_result': best_result,
'optimizer': optimizer,
}, is_best, epoch, output_directory)
def main():
global args, best_result, output_directory, train_csv, test_csv # 全局变量
args = parser.parse_args() # 获取参数值
args.data = os.path.join('data', args.data)
# os.path.join()函数：将多个路径组合后返回
# 语法：os.path.join(path1[,path2[,......]])
# 注：第一个绝对路径之前的参数将被忽略
# 注意if的语句后面有冒号
# args中modality的参数值。modality之前定义过
if args.modality == 'rgb' and args.num_samples != 0:
print("number of samples is forced to be 0 when input modality is rgb")
args.num_samples = 0
# 若是RGB的sparse-to-dense，则在生成训练数据时将稀疏深度点设为0
# create results folder, if not already exists
output_directory = os.path.join(
'results',
'NYUDataset.modality={}.nsample={}.arch={}.decoder={}.criterion={}.lr={}.bs={}'
.format(args.modality, args.num_samples, args.arch, args.decoder,
args.criterion, args.lr, args.batch_size)) # 输出文件名的格式
# 如果路径不存在
if not os.path.exists(output_directory):
os.makedirs(output_directory)
train_csv = os.path.join(output_directory, 'train.csv')
test_csv = os.path.join(output_directory, 'test.csv')
best_txt = os.path.join(output_directory, 'best.txt')
# define loss function (criterion) and optimizer
if args.criterion == 'l2':
criterion = criteria.MaskedMSELoss().cuda(
) # 调用别的py文件中的内容时，若被调用的是函数，则直接写函数名即可；若被调用的是类，则要按这句话的格式写
out_channels = 1
# elif: else if
elif args.criterion == 'l1':
criterion = criteria.MaskedL1Loss().cuda()
out_channels = 1
# Data loading code
print("=> creating data loaders ...")
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
train_dataset = NYUDataset(traindir,
type='train',
modality=args.modality,
num_samples=args.num_samples)
# DataLoader是导入数据的函数
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None)
# set batch size to be 1 for validation
val_dataset = NYUDataset(valdir,
type='val',
modality=args.modality,
num_samples=args.num_samples)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print("=> data loaders created.")
# evaluation mode
if args.evaluate:
best_model_filename = os.path.join(output_directory,
'model_best.pth.tar')
if os.path.isfile(best_model_filename):
print("=> loading best model '{}'".format(best_model_filename))
checkpoint = torch.load(best_model_filename)
args.start_epoch = checkpoint['epoch']
best_result = checkpoint['best_result']
model = checkpoint['model']
print("=> loaded best model (epoch {})".format(
checkpoint['epoch']))
else: # else也要加:
print("=> no best model found at '{}'".format(best_model_filename))
validate(val_loader, model, checkpoint['epoch'], write_to_file=False)
return
# optionally resume from a checkpoint
elif args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
best_result = checkpoint['best_result']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# create new model
else:
# define model
print("=> creating Model ({}-{}) ...".format(args.arch, args.decoder))
in_channels = len(args.modality) # len()返回对象的长度或项目个数
if args.arch == 'resnet50':
model = ResNet(layers=50,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
elif args.arch == 'resnet18':
model = ResNet(layers=18,
decoder=args.decoder,
in_channels=in_channels,
out_channels=out_channels,
pretrained=args.pretrained)
print("=> model created.")
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# create new csv files with only header
# with open() as xxx: 的用法详见https://www.cnblogs.com/ymjyqsx/p/6554817.html
with open(train_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
with open(test_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
print(model)
print("=> model transferred to GPU.")
# for循环也要有:
# 一般情况下，循环次数未知采用while循环，循环次数已知采用for
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
result, img_merge = validate(val_loader, model, epoch)
# Python的return可以返回多个值
# remember best rmse and save checkpoint
is_best = result.rmse < best_result.rmse
if is_best:
best_result = result
with open(best_txt, 'w') as txtfile:
# 字符串格式化输出
# :3f中，3表示输出宽度，f表示浮点型。若输出位数小于此宽度，则默认右对齐，左边补空格。
# 若输出位数大于宽度，则按实际位数输出。
# :.3f中，.3表示指定除小数点外的输出位数，f表示浮点型。
txtfile.write(
"epoch={}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n"
.format(epoch, result.mse, result.rmse, result.absrel,
result.lg10, result.mae, result.delta1,
result.gpu_time))
# None表示该值是一个空对象，空值是Python里一个特殊的值，用None表示。None不能理解为0，因为0是有意义的，而None是一个特殊的空值。
# 你可以将None赋值给任何变量，也可以将任何变量赋值给一个None值的对象
# None在判断的时候是False
# NULL是空字符，和None不一样
if img_merge is not None:
img_filename = output_directory + '/comparison_best.png'
utils.save_image(img_merge, img_filename)
# Python中，万物皆对象，所有的操作都是针对对象的。一个对象包括两方面的特征：
# 属性：去描述它的特征
# 方法：它所具有的行为
# 所以，对象=属性+方法 （其实方法也是一种属性，一种区别于数据属性的可调用属性）
save_checkpoint(
{
'epoch': epoch,
'arch': args.arch,
'model': model,
'best_result': best_result,
'optimizer': optimizer,
}, is_best, epoch)