def main():
args = parse_args()
save_path = args.save
if not os.path.isdir(save_path):
os.makedirs(save_path)
train_dataset_file = os.path.join(args.dataset, 'train.txt')
# val_dataset_file = os.path.join(args.dataset, 'val.txt')
train_dataset = LaneDataSet(train_dataset_file,
transform=transforms.Compose(
[Rescale((512, 256))]))
# val_dataset = LaneDataSet(val_dataset_file, transform=transforms.Compose([Rescale((512, 256))]))
model = LaneNet()
model.to(DEVICE)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
# val_loader = DataLoader(val_dataset, batch_size=8, shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
print(f"{args.epochs} epochs {len(train_dataset)} training samples\n")
for epoch in range(0, args.epochs):
print(f"Epoch {epoch}")
train_iou = train(train_loader, model, optimizer, epoch)
# val_iou = test(val_loader, model, epoch)
if (epoch + 1) % 5 == 0:
print("should save model")
# save_model(save_path, epoch, model)
# best_iou = max(val_iou, best_iou)
print(f"Best IoU : {train_iou}")
def main():
args = parse_args()
save_path = args.save
if not os.path.isdir(save_path):
os.makedirs(save_path)
image_output_path = args.image
if not os.path.isdir(image_output_path):
os.makedirs(image_output_path)
train_dataset_file = os.path.join(args.dataset, 'train.txt')
val_dataset_file = os.path.join(args.dataset, 'val.txt')
train_dataset = LaneDataSet(train_dataset_file,
transform=transforms.Compose(
[Rescale((512, 256))]))
# MES changes to use less GPU memory - batch size of 10 is max without OOM errors
# train_loader = DataLoader(train_dataset, batch_size=args.bs, shuffle=True)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
if args.val:
val_dataset = LaneDataSet(val_dataset_file,
transform=transforms.Compose(
[Rescale((512, 256))]))
val_loader = DataLoader(val_dataset, batch_size=args.bs, shuffle=True)
model = LaneNet()
model.to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
print(f"{args.epochs} epochs {len(train_dataset)} training samples\n")
for epoch in range(0, args.epochs):
print(f"Epoch {epoch}")
train_iou = train(train_loader, model, optimizer, epoch,
image_output_path)
if args.val:
val_iou = test(val_loader, model, epoch)
if (epoch + 1) % 5 == 0:
save_model(save_path, epoch, model)
print(f"Train IoU : {train_iou}")
if args.val:
print(f"Val IoU : {val_iou}")