class ENetDemo(ImageInferEngine):
def __init__(self, f, model_path):
super(ENetDemo, self).__init__(f=f)
self.target_size = (512, 1024)
self.model_path = model_path
self.num_classes = 20
self.image_transform = transforms.Compose(
[transforms.Resize(self.target_size),
transforms.ToTensor()])
self._init_model()
def _init_model(self):
self.model = ENet(self.num_classes).to(device)
checkpoint = torch.load(self.model_path)
self.model.load_state_dict(checkpoint['state_dict'])
print('Model loaded!')
def solve_a_image(self, img):
images = Variable(
self.image_transform(Image.fromarray(img)).to(device).unsqueeze(0))
predictions = self.model(images)
_, predictions = torch.max(predictions.data, 1)
prediction = predictions.cpu().numpy()[0] - 1
return prediction
def vis_result(self, img, net_out):
mask_color = np.asarray(label_to_color_image(net_out, 'cityscapes'),
dtype=np.uint8)
frame = cv2.resize(img, (self.target_size[1], self.target_size[0]))
# mask_color = cv2.resize(mask_color, (frame.shape[1], frame.shape[0]))
res = cv2.addWeighted(frame, 0.5, mask_color, 0.7, 1)
return res
def train(train_loader, val_loader, class_weights):
model = ENet(num_classes)
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = optim.Adam(model.parameters(), lr=5e-4, weight_decay=2e-4)
lr_updater = lr_scheduler.StepLR(
optimizer, 10, 1e-7) # Large dataset, decaying every 10 epochs..
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
model = model.cuda()
criterion = criterion.cuda()
# model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
# model, optimizer, args.save_dir, args.name)
# print("Resuming from model: Start epoch = {0} "
# "| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
start_epoch = 0
best_miou = 0
train = Train(model,
train_loader,
optimizer,
criterion,
metric,
use_cuda=True)
val = Test(model, val_loader, criterion, metric, use_cuda=True)
n_epochs = 200
for epoch in range(start_epoch, n_epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step()
epoch_loss, (iou, miou) = train.run_epoch(iteration_loss=True)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
if (epoch + 1) % 10 == 0 or epoch + 1 == n_epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(iteration_loss=True)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
# Print per class IoU on last epoch or if best iou
if epoch + 1 == n_epochs or miou > best_miou:
for class_iou in iou:
print(class_iou)
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
torch.save(
model.state_dict(),
'/mnt/disks/data/d4dl/snapshots/snapshot_' + str(epoch) +
'.pt')
return model
def main_script(args):
# Fail fast if the dataset directory doesn't exist
assert os.path.isdir(
args.dataset_dir), "The directory \"{0}\" doesn't exist.".format(
args.dataset_dir)
# Fail fast if the saving directory doesn't exist
assert os.path.isdir(
args.save_dir), "The directory \"{0}\" doesn't exist.".format(
args.save_dir)
# Import the requested dataset
if args.dataset.lower() == 'camvid':
from data import CamVid as dataset
elif args.dataset.lower() == 'cityscapes':
from data import Cityscapes as dataset
else:
# Should never happen...but just in case it does
raise RuntimeError("\"{0}\" is not a supported dataset.".format(
args.dataset))
loaders, w_class, class_encoding = load_dataset(dataset, args.color_space,
args.hue_value)
train_loader, val_loader, test_loader = loaders
if args.mode.lower() in {'train', 'full'}:
model = train(train_loader, val_loader, w_class, class_encoding)
if args.mode.lower() == 'full':
test(model, test_loader, w_class, class_encoding)
elif args.mode.lower() == 'test':
# Intialize a new ENet model
num_classes = len(class_encoding)
model = ENet(num_classes)
if use_cuda:
model = model.cuda()
# Here we register forward hooks for each layer.
# model.initial_block.register_forward_hook(save_activations)
# model.downsample1_0.register_forward_hook(save_activations)
# model.regular1_1.register_forward_hook(save_activations)
# model.downsample2_0.register_forward_hook(save_activations)
# Initialize a optimizer just so we can retrieve the model from the
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
test(model, test_loader, w_class, class_encoding)
else:
# Should never happen...but just in case it does
raise RuntimeError(
"\"{0}\" is not a valid choice for execution mode.".format(
args.mode))
def predict():
image_transform = transforms.Compose(
[transforms.Resize(target_size),
transforms.ToTensor()])
label_transform = transforms.Compose(
[transforms.Resize(target_size),
ext_transforms.PILToLongTensor()])
# Get selected dataset
# Load the training set as tensors
train_set = Cityscapes(data_dir,
mode='test',
transform=image_transform,
label_transform=label_transform)
class_encoding = train_set.color_encoding
num_classes = len(class_encoding)
model = ENet(num_classes).to(device)
# Initialize a optimizer just so we can retrieve the model from the
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, 'save',
'ENet_cityscapes_mine.pth')[0]
# print(model)
image = Image.open('images/mainz_000000_008001_leftImg8bit.png')
images = Variable(image_transform(image).to(device).unsqueeze(0))
image = np.array(image)
# Make predictions!
predictions = model(images)
_, predictions = torch.max(predictions.data, 1)
# 0~18
prediction = predictions.cpu().numpy()[0] - 1
mask_color = np.asarray(label_to_color_image(prediction, 'cityscapes'),
dtype=np.uint8)
mask_color = cv2.resize(mask_color, (image.shape[1], image.shape[0]))
print(image.shape)
print(mask_color.shape)
res = cv2.addWeighted(image, 0.3, mask_color, 0.7, 0.6)
# cv2.imshow('rr', mask_color)
cv2.imshow('combined', res)
cv2.waitKey(0)
def main():
"""Main function."""
loaders, class_weights, class_encoding = load_dataset(dataset)
train_loader, val_loader, test_loader = loaders
num_classes = len(class_encoding)
critic = DiscriminativeNet()
generator = ENet(num_classes)
dataloader = load_real_data(real_dataset)
optimizer_D = optim.Adam(critic.parameters(),
lr=0.0001,
betas=(0.5, 0.999))
optimizer_G = optim.Adam(generator.parameters(),
lr=0.0001,
betas=(0.5, 0.999))
gan = WGANGP(generator,
critic,
dataloader,
train_loader,
test_loader,
class_weights,
class_encoding,
ngpu=ngpu,
device=device,
nr_epochs=500,
print_every=10,
save_every=400,
optimizer_D=optimizer_D,
optimizer_G=optimizer_G)
gan.train()
samples_l, D_losses, G_losses = gan.get_training_results()
def create_model(args,n_classes):
"""
Creates a model according to:
- args.model_type: ENet or BiSeNet
- args.dropout: if different from 0 or None, a dropout is done
Dropout is done by adding a forward hook to the batch-norm layers. Dropped-
out pixels are replaced by the batch norm bias for their channel. It is
designed for MC dropout (i.e. even at inference stage) and is always active
even after calling model.eval() or model.train(False).
Returns a Torch model
"""
if not hasattr(args,"model_type") or args.model_type.lower()=="enet":
model = ENet(n_classes)
elif args.model_type.lower()=="bisenet":
model = BiSeNetv2(3, n_classes, ratio=8)
else:
raise ValueError("Model type: expected ENet or BiSeNet")
if hasattr(args,"dropout") and args.dropout not in [0, None]:
utils.add_dropout(model, args.dropout)
return model
import torch
from models.enet import ENet
import os
from configs import config_factory
if __name__ == '__main__':
save_pth = os.path.join(config_factory['resnet_cityscapes'].respth,
'model_final.pth')
model = ENet(nb_classes=19)
model.load_state_dict(torch.load(save_pth))
model.eval()
example = torch.rand(2, 3, 1024, 1024).cpu()
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(
os.path.join(config_factory['resnet_cityscapes'].respth,
"model_dfanet_1024.pt"))
def __init__(self, s_exp_name, t_exp_name):
cfg_path = os.path.join(os.getcwd(), 'config/tusimple_config.yaml')
self.s_exp_name = s_exp_name
self.t_exp_name = t_exp_name
self.writer = SummaryWriter('tensorboard/' + self.s_exp_name)
self.metric = SegmentationMetric(7)
with open(cfg_path) as cfg:
config = yaml.load(cfg, Loader=yaml.FullLoader)
self.device = torch.device(config['DEVICE'])
self.max_epochs = config['TRAIN']['MAX_EPOCHS']
self.dataset_path = config['DATASET']['PATH']
self.mean = config['DATASET']['MEAN']
self.std = config['DATASET']['STD']
'''
self.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(self.mean, self.std),
])
'''
self.train_transform = Compose(Resize(size=(645, 373)),
RandomCrop(size=(640, 368)),
RandomFlip(0.5), Rotation(2),
ToTensor(),
Normalize(mean=self.mean, std=self.std))
self.val_transform = Compose(Resize(size=(640, 368)), ToTensor(),
Normalize(mean=self.mean, std=self.std))
self.train_dataset = tuSimple(path=config['DATASET']['PATH'],
image_set='train',
transforms=self.train_transform)
self.val_dataset = tuSimple(
path=config['DATASET']['PATH'],
image_set='val',
transforms=self.val_transform,
)
self.train_loader = data.DataLoader(
dataset=self.train_dataset,
batch_size=config['TRAIN']['BATCH_SIZE'],
shuffle=True,
num_workers=0,
pin_memory=True,
drop_last=True,
)
self.val_loader = data.DataLoader(
dataset=self.val_dataset,
batch_size=config['TRAIN']['BATCH_SIZE'],
shuffle=False,
num_workers=0,
pin_memory=True,
drop_last=False,
)
self.iters_per_epoch = len(
self.train_dataset) // config['TRAIN']['BATCH_SIZE']
self.max_iters = self.max_epochs * self.iters_per_epoch
# ------------network------------
self.s_model = ENet(num_classes=7).to(self.device)
self.t_model = ENet(num_classes=7).to(self.device)
self.optimizer = optim.SGD(
self.s_model.parameters(),
lr=config['OPTIM']['LR'],
weight_decay=config['OPTIM']['DECAY'],
momentum=0.9,
)
self.lr_scheduler = get_scheduler(
self.optimizer,
max_iters=self.max_iters,
iters_per_epoch=self.iters_per_epoch,
)
self.ce = nn.CrossEntropyLoss(weight=torch.tensor(
[0.4, 1, 1, 1, 1, 1, 1])).cuda() #background weight 0.4
self.bce = nn.BCELoss().cuda()
self.kl = nn.KLDivLoss().cuda() #reduction='batchmean' gives NaN
self.mse = nn.MSELoss().cuda()
class Trainer(object):
def __init__(self, s_exp_name, t_exp_name):
cfg_path = os.path.join(os.getcwd(), 'config/tusimple_config.yaml')
self.s_exp_name = s_exp_name
self.t_exp_name = t_exp_name
self.writer = SummaryWriter('tensorboard/' + self.s_exp_name)
self.metric = SegmentationMetric(7)
with open(cfg_path) as cfg:
config = yaml.load(cfg, Loader=yaml.FullLoader)
self.device = torch.device(config['DEVICE'])
self.max_epochs = config['TRAIN']['MAX_EPOCHS']
self.dataset_path = config['DATASET']['PATH']
self.mean = config['DATASET']['MEAN']
self.std = config['DATASET']['STD']
'''
self.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(self.mean, self.std),
])
'''
self.train_transform = Compose(Resize(size=(645, 373)),
RandomCrop(size=(640, 368)),
RandomFlip(0.5), Rotation(2),
ToTensor(),
Normalize(mean=self.mean, std=self.std))
self.val_transform = Compose(Resize(size=(640, 368)), ToTensor(),
Normalize(mean=self.mean, std=self.std))
self.train_dataset = tuSimple(path=config['DATASET']['PATH'],
image_set='train',
transforms=self.train_transform)
self.val_dataset = tuSimple(
path=config['DATASET']['PATH'],
image_set='val',
transforms=self.val_transform,
)
self.train_loader = data.DataLoader(
dataset=self.train_dataset,
batch_size=config['TRAIN']['BATCH_SIZE'],
shuffle=True,
num_workers=0,
pin_memory=True,
drop_last=True,
)
self.val_loader = data.DataLoader(
dataset=self.val_dataset,
batch_size=config['TRAIN']['BATCH_SIZE'],
shuffle=False,
num_workers=0,
pin_memory=True,
drop_last=False,
)
self.iters_per_epoch = len(
self.train_dataset) // config['TRAIN']['BATCH_SIZE']
self.max_iters = self.max_epochs * self.iters_per_epoch
# ------------network------------
self.s_model = ENet(num_classes=7).to(self.device)
self.t_model = ENet(num_classes=7).to(self.device)
self.optimizer = optim.SGD(
self.s_model.parameters(),
lr=config['OPTIM']['LR'],
weight_decay=config['OPTIM']['DECAY'],
momentum=0.9,
)
self.lr_scheduler = get_scheduler(
self.optimizer,
max_iters=self.max_iters,
iters_per_epoch=self.iters_per_epoch,
)
self.ce = nn.CrossEntropyLoss(weight=torch.tensor(
[0.4, 1, 1, 1, 1, 1, 1])).cuda() #background weight 0.4
self.bce = nn.BCELoss().cuda()
self.kl = nn.KLDivLoss().cuda() #reduction='batchmean' gives NaN
self.mse = nn.MSELoss().cuda()
def train(self, epoch, start_time):
running_loss = 0.0
is_better = True
prev_loss = float('inf')
logging.info(
'Start training, Total Epochs: {:d}, Total Iterations: {:d}'.
format(self.max_epochs, self.max_iters))
print("Train Epoch: {}".format(epoch))
self.s_model.train()
self.t_model.eval()
epoch_loss = 0
iteration = epoch * self.iters_per_epoch if epoch > 0 else 0
start_time = start_time
for batch_idx, sample in enumerate(self.train_loader):
iteration += 1
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
with torch.no_grad():
t_outputs, t_sig = self.t_model(img)
s_outputs, s_sig = self.s_model(img)
ce = self.ce(s_outputs, segLabel)
bce = self.bce(s_sig, exist)
kl = self.kl(
F.log_softmax(s_outputs, dim=1),
F.softmax(t_outputs, dim=1),
)
mse = self.mse(s_outputs, t_outputs) #/ s_outputs.size(0)
loss = ce + (0.1 * bce) + kl + (0.5 * mse)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
epoch_loss += loss.item()
running_loss += loss.item()
eta_seconds = ((time.time() - start_time) /
iteration) * (self.max_iters - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
iter_idx = epoch * len(self.train_loader) + batch_idx
if iteration % 10 == 0:
logging.info(
"Epoch: {:d}/{:d} || Iters: {:d}/{:d} || Lr: {:6f} || Loss: {:.4f} || Cost Time: {} || Estimated Time: {}"
.format(
epoch,
self.max_epochs,
iteration % self.iters_per_epoch,
self.iters_per_epoch,
self.optimizer.param_groups[0]['lr'],
loss.item(),
str(
datetime.timedelta(seconds=int(time.time() -
start_time))),
eta_string,
))
if batch_idx % 10 == 9:
self.writer.add_scalar(
'train_loss', running_loss / 10,
epoch * len(self.train_loader) + batch_idx + 1)
running_loss = 0.0
if epoch % 1 == 0:
save_dict = {
"epoch": epoch,
"model": self.s_model.state_dict(),
"optim": self.optimizer.state_dict(),
"best_mIoU": best_mIoU,
"best_val_loss": best_val_loss,
}
save_name = os.path.join(os.getcwd(), 'results', self.s_exp_name,
'run.pth')
torch.save(save_dict, save_name)
print("Model is saved: {}".format(save_name))
def val(self, epoch):
self.metric.reset()
global best_val_loss
global best_mIoU
print("Val Epoch: {}".format(epoch))
self.s_model.eval()
val_loss = 0
with torch.no_grad():
for batch_idx, sample in enumerate(self.val_loader):
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
outputs, sig = self.s_model(img)
ce = self.ce(outputs, segLabel)
bce = self.bce(sig, exist)
loss = ce + (0.1 * bce)
self.metric.update(outputs, segLabel)
pixAcc, mIoU = self.metric.get()
logging.info(
"Sample: {:d}, pixAcc: {:.3f}, mIoU: {:.3f}".format(
batch_idx + 1, pixAcc * 100, mIoU * 100))
pixAcc, mIoU, category_iou = self.metric.get(return_category_iou=True)
print(category_iou)
logging.info("Final pixAcc: {:.3f}, mIoU: {:.3f}".format(
pixAcc * 100,
mIoU * 100,
))
iter_idx = (epoch + 1) * len(self.train_loader)
if (mIoU * 100) > best_mIoU:
best_mIoU = mIoU * 100
save_dict = {
"epoch": epoch,
"model": self.s_model.state_dict(),
"optim": self.optimizer.state_dict(),
"best_val_loss": best_val_loss,
"best_mIoU": best_mIoU,
}
save_name = os.path.join(os.getcwd(), 'results', self.s_exp_name,
'best_mIoU.pth')
torch.save(save_dict, save_name)
print("mIoU is higher than best mIoU! Model saved to {}".format(
save_name))
class Trainer(object):
def __init__(self, exp, exp2):
cfg_path = os.path.join(os.getcwd(), 'config/tusimple_config.yaml')
self.exp_name = exp
self.exp_name2 = exp2
self.writer = SummaryWriter('tensorboard/' + self.exp_name)
with open(cfg_path) as file:
cfg = yaml.load(file, Loader=yaml.FullLoader)
self.device = torch.device(cfg['DEVICE'])
self.max_epochs = cfg['TRAIN']['MAX_EPOCHS']
self.dataset_path = cfg['DATASET']['PATH']
# TODO remove this and refactor PROPERLY
self.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cfg['DATASET']['MEAN'],
cfg['DATASET']['STD']),
])
mean = cfg['DATASET']['MEAN']
std = cfg['DATASET']['STD']
self.train_transform = Compose(Resize(size=(645, 373)),
RandomCrop(size=(640, 368)),
RandomFlip(0.5), Rotation(2),
ToTensor(), Normalize(mean=mean,
std=std))
self.val_transform = Compose(Resize(size=(640, 368)), ToTensor(),
Normalize(mean=mean, std=std))
data_kwargs = {
'transform': self.input_transform,
'size': cfg['DATASET']['SIZE'],
}
self.train_dataset = tuSimple(path=cfg['DATASET']['PATH'],
image_set='train',
transforms=self.train_transform)
self.val_dataset = tuSimple(
path=cfg['DATASET']['PATH'],
image_set='val',
transforms=self.val_transform,
)
self.train_loader = data.DataLoader(
dataset=self.train_dataset,
batch_size=cfg['TRAIN']['BATCH_SIZE'],
shuffle=True,
num_workers=0,
pin_memory=True,
drop_last=True,
)
self.val_loader = data.DataLoader(
dataset=self.val_dataset,
batch_size=cfg['TRAIN']['BATCH_SIZE'],
shuffle=False,
num_workers=0,
pin_memory=True,
drop_last=False,
)
# -------- network --------
weight = [0.4, 1, 1, 1, 1, 1, 1]
self.model = ENet(num_classes=7).to(self.device)
self.model2 = ENet(num_classes=7).to(self.device)
self.optimizer = optim.SGD(
self.model.parameters(),
lr=cfg['OPTIM']['LR'],
weight_decay=cfg['OPTIM']['DECAY'],
momentum=0.9,
)
#self.optimizer = optim.Adam(
# self.model.parameters(),
# lr = cfg['OPTIM']['LR'],
# weight_decay=0,
# )
self.criterion = nn.CrossEntropyLoss(
weight=torch.tensor([0.4, 1, 1, 1, 1, 1, 1])).cuda()
self.bce = nn.BCELoss().cuda()
def train(self, epoch):
running_loss = 0.0
is_better = True
prev_loss = float('inf')
print("Train Epoch: {}".format(epoch))
self.model.train()
epoch_loss = 0
progressbar = tqdm(range(len(self.train_loader)))
for batch_idx, sample in enumerate(self.train_loader):
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
# outputs is crossentropy, sig is binary cross entropy
outputs, sig = self.model(img)
ce = self.criterion(outputs, segLabel)
bce = self.bce(sig, exist)
loss = ce + (0.1 * bce)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
epoch_loss += loss.item()
running_loss += loss.item()
iter_idx = epoch * len(self.train_loader) + batch_idx
progressbar.set_description("Batch loss: {:.3f}".format(
loss.item()))
progressbar.update(1)
# Tensorboard
if batch_idx % 10 == 9:
self.writer.add_scalar(
'train loss', running_loss / 10,
epoch * len(self.train_loader) + batch_idx + 1)
running_loss = 0.0
progressbar.close()
if epoch % 1 == 0:
save_dict = {
"epoch": epoch,
"model": self.model.state_dict(),
"optim": self.optimizer.state_dict(),
"best_val_loss": best_val_loss,
}
os.makedirs(os.path.join(os.getcwd(), 'results', self.exp_name),
exist_ok=True)
save_name = os.path.join(os.getcwd(), 'results', self.exp_name,
'run.pth')
torch.save(save_dict, save_name)
print("Model is saved: {}".format(save_name))
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
return epoch_loss / len(self.train_loader)
def val(self, epoch, train_loss):
global best_val_loss
print("Val Epoch: {}".format(epoch))
self.model.eval()
val_loss = 0
progressbar = tqdm(range(len(self.val_loader)))
with torch.no_grad():
for batch_idx, sample in enumerate(self.val_loader):
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
outputs, sig = self.model(img)
ce = self.criterion(outputs, segLabel)
bce = self.bce(sig, exist)
loss = ce + (0.1 * bce)
val_loss += loss.item()
progressbar.set_description("Batch loss: {:3f}".format(
loss.item()))
progressbar.update(1)
# Tensorboard
if batch_idx + 1 == len(self.val_loader):
self.writer.add_scalar(
'train - val loss',
train_loss - (val_loss / len(self.val_loader)), epoch)
progressbar.close()
iter_idx = (epoch + 1) * len(self.train_loader)
print("Validation loss: {}".format(val_loss))
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
if val_loss < best_val_loss:
best_val_loss = val_loss
save_name = os.path.join(os.getcwd(), 'results', self.exp_name,
'run.pth')
copy_name = os.path.join(os.getcwd(), 'results', self.exp_name,
'run_best.pth')
print("val loss is lower than best val loss! Model saved to {}".
format(copy_name))
shutil.copyfile(save_name, copy_name)
def eval(self):
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
print("Evaluating.. ")
self.model.eval()
self.model2.eval()
val_loss = 0
dump_to_json = []
test_dataset = tuSimple(path=self.dataset_path,
image_set='test',
transforms=self.val_transform)
test_loader = data.DataLoader(
dataset=test_dataset,
batch_size=12,
shuffle=False,
num_workers=0,
pin_memory=True,
drop_last=False,
)
progressbar = tqdm(range(len(test_loader)))
with torch.no_grad():
with open('exist_out.txt', 'w') as f:
for batch_idx, sample in enumerate(test_loader):
img = sample['img'].to(self.device)
img_name = sample['img_name']
#segLabel = sample['segLabel'].to(self.device)
outputs, sig = self.model(img)
outputs2, sig2 = self.model2(img)
#added_sig = sig2.add(sig)
#div_sig = torch.div(added_sig, 2.0)
#added_out = outputs.add(outputs2)
#div_out = torch.div(added_out, 2.0)
seg_pred1 = F.softmax(outputs, dim=1)
seg_pred2 = F.softmax(outputs2, dim=1)
seg_pred = seg_pred1.add(seg_pred2)
seg_pred = torch.div(seg_pred, 2.0)
seg_pred = seg_pred.detach().cpu().numpy()
sig_pred = sig.add(sig2)
exist_pred = sig_pred.detach().cpu().numpy()
count = 0
for img_idx in range(len(seg_pred)):
seg = seg_pred[img_idx]
exist = [
1 if exist_pred[img_idx, i] > 0.8 else 0
for i in range(6)
]
lane_coords = getLane.prob2lines_tusimple(
seg,
exist,
resize_shape=(720, 1280),
y_px_gap=10,
pts=56)
for i in range(len(lane_coords)):
# sort lane coords
lane_coords[i] = sorted(lane_coords[i],
key=lambda pair: pair[1])
#print(len(lane_coords))
# Visualisation
savename = "{}/{}_{}_vis.png".format(
os.path.join(os.getcwd(), 'vis'), batch_idx, count)
count += 1
raw_file_name = img_name[img_idx]
pred_json = {}
pred_json['lanes'] = []
pred_json['h_samples'] = []
# truncate everything before 'clips' to be consistent with test_label.json gt
pred_json['raw_file'] = raw_file_name[raw_file_name.
find('clips'):]
pred_json['run_time'] = 0
for l in lane_coords:
empty = all(lane[0] == -2 for lane in l)
if len(l) == 0:
continue
if empty:
continue
pred_json['lanes'].append([])
for (x, y) in l:
pred_json['lanes'][-1].append(int(x))
for (x, y) in lane_coords[0]:
pred_json['h_samples'].append(int(y))
dump_to_json.append(json.dumps(pred_json))
progressbar.update(1)
progressbar.close()
with open(
os.path.join(os.getcwd(), "results", self.exp_name,
"pred_json.json"), "w") as f:
for line in dump_to_json:
print(line, end="\n", file=f)
print("Saved pred_json.json to {}".format(
os.path.join(os.getcwd(), "results", self.exp_name,
"pred_json.json")))
'''
raw_img = img[b].cpu().detach().numpy()
raw_img = raw_img.transpose(1, 2, 0)
# Normalize both to 0..1
min_val, max_val = np.min(raw_img), np.max(raw_img)
raw_img = (raw_img - min_val) / (max_val - min_val)
#rgb = rgb / 255.
#stack = np.hstack((raw_img, rgb))
background = Image.fromarray(np.uint8(raw_img*255))
overlay = Image.fromarray(rgb)
new_img = Image.blend(background, overlay, 0.4)
new_img.save(savename, "PNG")
'''
'''
from data import Cityscapes as dataset
else:
# Should never happen...but just in case it does
raise RuntimeError("\"{0}\" is not a supported dataset.".format(
args.dataset))
loaders, w_class, class_encoding = load_dataset(dataset)
train_loader, val_loader, test_loader = loaders
if args.mode.lower() in {'train', 'full'}:
model = train(train_loader, val_loader, w_class, class_encoding)
if args.mode.lower() in {'test', 'full'}:
if args.mode.lower() == 'test':
# Intialize a new ENet model
num_classes = len(class_encoding)
model = ENet(num_classes).to(device)
# Initialize a optimizer just so we can retrieve the model from the
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
if args.mode.lower() == 'test':
print(model)
test(model, test_loader, w_class, class_encoding)
class Trainer(object):
def __init__(self, exp):
# IoU and pixAcc Metric calculator
self.metric = SegmentationMetric(7)
cfg_path = os.path.join(os.getcwd(), 'config/tusimple_config.yaml')
self.exp_name = exp
self.writer = SummaryWriter('tensorboard/' + self.exp_name)
with open(cfg_path) as file:
cfg = yaml.load(file, Loader=yaml.FullLoader)
self.device = torch.device(cfg['DEVICE'])
self.max_epochs = cfg['TRAIN']['MAX_EPOCHS']
self.dataset_path = cfg['DATASET']['PATH']
# TODO remove this and refactor PROPERLY
self.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cfg['DATASET']['MEAN'], cfg['DATASET']['STD']),
])
mean = cfg['DATASET']['MEAN']
std = cfg['DATASET']['STD']
self.train_transform = Compose(Resize(size=(645,373)), RandomCrop(size=(640,368)), RandomFlip(0.5), Rotation(2), ToTensor(), Normalize(mean=mean, std=std))
self.val_transform = Compose(Resize(size=(640,368)), ToTensor(), Normalize(mean=mean, std=std))
data_kwargs = {
'transform': self.input_transform,
'size': cfg['DATASET']['SIZE'],
}
self.train_dataset = tuSimple(
path=cfg['DATASET']['PATH'],
image_set='train',
transforms=self.train_transform
)
self.val_dataset = tuSimple(
path = cfg['DATASET']['PATH'],
image_set = 'val',
transforms =self.val_transform,
)
self.train_loader = data.DataLoader(
dataset = self.train_dataset,
batch_size = cfg['TRAIN']['BATCH_SIZE'],
shuffle = True,
num_workers = 0,
pin_memory = True,
drop_last = True,
)
self.val_loader = data.DataLoader(
dataset = self.val_dataset,
batch_size = cfg['TRAIN']['BATCH_SIZE'],
shuffle = False,
num_workers = 0,
pin_memory = True,
drop_last = False,
)
self.iters_per_epoch = len(self.train_dataset) // (cfg['TRAIN']['BATCH_SIZE'])
self.max_iters = cfg['TRAIN']['MAX_EPOCHS'] * self.iters_per_epoch
# -------- network --------
weight = [0.4, 1, 1, 1, 1, 1, 1]
self.model = ENet(num_classes=7).to(self.device)
self.optimizer = optim.SGD(
self.model.parameters(),
lr=cfg['OPTIM']['LR'],
weight_decay=cfg['OPTIM']['DECAY'],
momentum=0.9,
)
self.lr_scheduler = get_scheduler(self.optimizer, max_iters=self.max_iters, iters_per_epoch=self.iters_per_epoch)
#self.optimizer = optim.Adam(
# self.model.parameters(),
# lr = cfg['OPTIM']['LR'],
# weight_decay=0,
# )
self.criterion = nn.CrossEntropyLoss(weight=torch.tensor([0.4, 1, 1, 1, 1, 1, 1])).cuda()
self.bce = nn.BCELoss().cuda()
def train(self, epoch, start_time):
running_loss = 0.0
is_better = True
prev_loss = float('inf')
logging.info('Start training, Total Epochs: {:d}, Total Iterations: {:d}'.format(self.max_epochs, self.max_iters))
print("Train Epoch: {}".format(epoch))
self.model.train()
epoch_loss = 0
#progressbar = tqdm(range(len(self.train_loader)))
iteration = epoch * self.iters_per_epoch if epoch > 0 else 0
start_time = start_time
for batch_idx, sample in enumerate(self.train_loader):
iteration += 1
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
# outputs is crossentropy, sig is binary cross entropy
outputs, sig = self.model(img)
ce = self.criterion(outputs,segLabel)
bce = self.bce(sig, exist)
loss = ce + (0.1 * bce)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
#print("LR", self.optimizer.param_groups[0]['lr'])
epoch_loss += loss.item()
running_loss += loss.item()
eta_seconds = ((time.time() - start_time) / iteration) * (self.max_iters - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
iter_idx = epoch * len(self.train_loader) + batch_idx
#progressbar.set_description("Batch loss: {:.3f}".format(loss.item()))
#progressbar.update(1)
# Tensorboard
if iteration % 10 == 0:
logging.info(
"Epoch: {:d}/{:d} || Iters: {:d}/{:d} || Lr: {:6f} || "
"Loss: {:.4f} || Cost Time: {} || Estimated Time: {}".format(
epoch, self.max_epochs, iteration % self.iters_per_epoch, self.iters_per_epoch,
self.optimizer.param_groups[0]['lr'], loss.item(), str(datetime.timedelta(seconds=int(time.time() - start_time))), eta_string))
if batch_idx % 10 == 9:
self.writer.add_scalar('train loss',
running_loss / 10,
epoch * len(self.train_loader) + batch_idx + 1)
running_loss = 0.0
#progressbar.close()
if epoch % 1 == 0:
save_dict = {
"epoch": epoch,
"model": self.model.state_dict(),
"optim": self.optimizer.state_dict(),
"best_val_loss": best_val_loss,
}
save_name = os.path.join(os.getcwd(), 'results', self.exp_name, 'run.pth')
save_name_epoch = os.path.join(os.getcwd(), 'results', self.exp_name, '{}.pth'.format(epoch))
torch.save(save_dict, save_name)
torch.save(save_dict, save_name_epoch)
print("Model is saved: {}".format(save_name))
print("Model is saved: {}".format(save_name_epoch))
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
return epoch_loss/len(self.train_loader)
def val(self, epoch, train_loss):
self.metric.reset()
global best_val_loss
global best_mIoU
print("Val Epoch: {}".format(epoch))
self.model.eval()
val_loss = 0
#progressbar = tqdm(range(len(self.val_loader)))
with torch.no_grad():
for batch_idx, sample in enumerate(self.val_loader):
img = sample['img'].to(self.device)
segLabel = sample['segLabel'].to(self.device)
exist = sample['exist'].to(self.device)
outputs, sig = self.model(img)
ce = self.criterion(outputs, segLabel)
bce = self.bce(sig, exist)
loss = ce + (0.1*bce)
val_loss += loss.item()
self.metric.update(outputs, segLabel)
pixAcc, mIoU = self.metric.get()
logging.info("Sample: {:d}, validation pixAcc: {:.3f}, mIoU: {:.3f}".format(
batch_idx + 1, pixAcc * 100, mIoU * 100))
#progressbar.set_description("Batch loss: {:3f}".format(loss.item()))
#progressbar.update(1)
# Tensorboard
#if batch_idx + 1 == len(self.val_loader):
# self.writer.add_scalar('train - val loss',
# train_loss - (val_loss / len(self.val_loader)),
# epoch)
#progressbar.close()
pixAcc, mIoU, category_iou = self.metric.get(return_category_iou = True)
print(category_iou)
logging.info('End validation pixAcc: {:.3f}, mIoU: {:.3f}'.format(
pixAcc * 100, mIoU * 100))
iter_idx = (epoch + 1) * len(self.train_loader)
with open('val_out.txt', 'a') as out:
sys.stdout = out
print(self.exp_name, 'Epoch:', epoch, 'pixAcc: {:.3f}, mIoU: {:.3f}'.format(pixAcc*100, mIoU*100))
sys.stdout = original_stdout
print("Validation loss: {}".format(val_loss))
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
if (mIoU * 100) > best_mIoU:
best_mIoU = mIoU*100
save_dict = {
"epoch": epoch,
"model": self.model.state_dict(),
"optim": self.optimizer.state_dict(),
"best_val_loss": best_val_loss,
"best_mIoU": best_mIoU,
}
save_name = os.path.join(os.getcwd(), 'results', self.exp_name, 'best_mIoU.pth')
torch.save(save_dict, save_name)
print("mIoU is higher than best mIoU! Model saved to {}".format(save_name))
#if val_loss < best_val_loss:
# best_val_loss = val_loss
# save_name = os.path.join(os.getcwd(), 'results', self.exp_name, 'run.pth')
# copy_name = os.path.join(os.getcwd(), 'results', self.exp_name, 'run_best.pth')
# print("val loss is lower than best val loss! Model saved to {}".format(copy_name))
# shutil.copyfile(save_name, copy_name)
def eval(self):
print("+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*+*")
print("Evaluating.. ")
self.model.eval()
val_loss = 0
dump_to_json = []
test_dataset = tuSimple(
path=self.dataset_path,
image_set='test',
transforms=self.val_transform
)
test_loader = data.DataLoader(
dataset = test_dataset,
batch_size = 12,
shuffle = False,
num_workers = 0,
pin_memory = True,
drop_last = False,
)
progressbar = tqdm(range(len(test_loader)))
with torch.no_grad():
with open('exist_out.txt','w') as f:
for batch_idx, sample in enumerate(test_loader):
img = sample['img'].to(self.device)
img_name = sample['img_name']
#segLabel = sample['segLabel'].to(self.device)
outputs, sig = self.model(img)
seg_pred = F.softmax(outputs, dim=1)
seg_pred = seg_pred.detach().cpu().numpy()
exist_pred = sig.detach().cpu().numpy()
count = 0
for img_idx in range(len(seg_pred)):
seg = seg_pred[img_idx]
exist = [1 if exist_pred[img_idx ,i] > 0.5 else 0 for i in range(6)]
lane_coords = getLane.prob2lines_tusimple(seg, exist, resize_shape=(720,1280), y_px_gap=10, pts=56)
for i in range(len(lane_coords)):
# sort lane coords
lane_coords[i] = sorted(lane_coords[i], key=lambda pair:pair[1])
#print(len(lane_coords))
# Visualisation
savename = "{}/{}_{}_vis.png".format(os.path.join(os.getcwd(), 'vis'), batch_idx, count)
count += 1
raw_file_name = img_name[img_idx]
pred_json = {}
pred_json['lanes'] = []
pred_json['h_samples'] = []
# truncate everything before 'clips' to be consistent with test_label.json gt
pred_json['raw_file'] = raw_file_name[raw_file_name.find('clips'):]
pred_json['run_time'] = 0
for l in lane_coords:
empty = all(lane[0] == -2 for lane in l)
if len(l)==0:
continue
if empty:
continue
pred_json['lanes'].append([])
for (x,y) in l:
pred_json['lanes'][-1].append(int(x))
for (x, y) in lane_coords[0]:
pred_json['h_samples'].append(int(y))
dump_to_json.append(json.dumps(pred_json))
progressbar.update(1)
progressbar.close()
with open(os.path.join(os.getcwd(), "results", self.exp_name, "pred_json.json"), "w") as f:
for line in dump_to_json:
print(line, end="\n", file=f)
print("Saved pred_json.json to {}".format(os.path.join(os.getcwd(), 'results', self.exp_name, "pred_json.json")))
'''
raw_img = img[b].cpu().detach().numpy()
raw_img = raw_img.transpose(1, 2, 0)
# Normalize both to 0..1
min_val, max_val = np.min(raw_img), np.max(raw_img)
raw_img = (raw_img - min_val) / (max_val - min_val)
#rgb = rgb / 255.
#stack = np.hstack((raw_img, rgb))
background = Image.fromarray(np.uint8(raw_img*255))
overlay = Image.fromarray(rgb)
new_img = Image.blend(background, overlay, 0.4)
new_img.save(savename, "PNG")
'''
'''
def create_model(num_classes=20, device='cuda'):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ENet(num_classes).to(device)
return model
else:
# Should never happen...but just in case it does
raise RuntimeError("\"{0}\" is not a supported dataset.".format(
args.dataset))
loaders, w_class, class_encoding = load_dataset(dataset)
train_loader, val_loader, test_loader = loaders
if args.mode.lower() in {'train', 'full'}:
model = train(train_loader, val_loader, w_class, class_encoding)
if args.mode.lower() == 'full':
test(model, test_loader, w_class, class_encoding)
elif args.mode.lower() == 'test':
# Intialize a new ENet model
num_classes = len(class_encoding)
model = ENet(num_classes)
if use_cuda:
model = model.cuda()
# Initialize a optimizer just so we can retrieve the model from the
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
print(model)
test(model, test_loader, w_class, class_encoding)
else:
# Should never happen...but just in case it does
raise RuntimeError(
def train(train_loader, val_loader, class_weights, class_encoding):
print("\nTraining...\n")
vis_calling_times = 0
num_classes = len(class_encoding)
# Intialize ENet
model = ENet(num_classes).to(device)
# Check if the network architecture is correct
if torch.cuda.device_count() > 1:
print(">>>Use mult GPU for trainning>>>")
gpu_num = torch.cuda.device_count()
gpu_list = list(range(gpu_num))
model = nn.DataParallel(model, device_ids=gpu_list)
print(model)
# We are going to use the CrossEntropyLoss loss function as it's most
# frequentely used in classification problems with multiple classes which
# fits the problem. This criterion combines LogSoftMax and NLLLoss.
criterion = nn.CrossEntropyLoss(weight=class_weights)
# ENet authors used Adam as the optimizer
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
# Optionally resume from a checkpoint
if args.resume:
model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# Start Training
print()
train = Train(model, train_loader, optimizer, criterion, metric, device)
val = Test(model, val_loader, criterion, metric, device)
for epoch in range(start_epoch, args.epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step()
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
if (epoch + 1) % 10 == 0 or epoch + 1 == args.epochs:
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
args)
if vis_calling_times == 0:
# set to false
vis_calling_times = 1
win = viz.line(X=np.column_stack(
(np.array(epoch), np.array(epoch))),
Y=np.column_stack(
(np.array(epoch_loss), np.array(loss))),
opts=dict(legend=['training loss', 'eval loss'],
title='loss'))
else:
viz.line(
X=np.column_stack((np.array(epoch), np.array(epoch))),
Y=np.column_stack((np.array(epoch_loss), np.array(loss))),
win=win, #win要保持一致
update='append')
# if vis_first_create:
# vis_first_create = false
# win = viz.line( X=np.column_stack((np.array(epoch),np.array(epoch))),
# Y=np.column_stack((np.array(epoch_loss),np.array(loss))),
# name=
# opts=dict(title='loss'))
# else:
# viz.line( X=np.column_stack((np.array(epoch),np.array(epoch))),
# Y=np.column_stack((np.array(epoch_loss),np.array(loss))),
# win=win,#win要保持一致
# update='append')
return model
def _init_model(self):
self.model = ENet(self.num_classes).to(device)
checkpoint = torch.load(self.model_path)
self.model.load_state_dict(checkpoint['state_dict'])
print('Model loaded!')
def initialize(self):
args = self.args
if args.architecture == 'deeplab':
print('Using Deeplab')
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
elif args.architecture == 'enet':
print('Using ENet')
model = ENet(num_classes=self.nclass,
encoder_relu=True,
decoder_relu=True)
train_params = [{'params': model.parameters(), 'lr': args.lr}]
elif args.architecture == 'fastscnn':
print('Using FastSCNN')
model = FastSCNN(3, self.nclass)
train_params = [{'params': model.parameters(), 'lr': args.lr}]
if args.optimizer == 'SGD':
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
elif args.optimizer == 'Adam':
optimizer = torch.optim.Adam(train_params,
weight_decay=args.weight_decay)
else:
raise NotImplementedError
if args.use_balanced_weights:
weight = calculate_weights_labels(args.dataset, self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
self.evaluator = Evaluator(self.nclass)
if args.use_lr_scheduler:
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.train_loader))
else:
self.scheduler = None
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
self.best_pred = 0.0
def __init__(self, args):
self.args = args
self.saver = PassiveSaver(args)
self.saver.save_experiment_config()
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
kwargs = {'pin_memory': False, 'memory_hog': args.memory_hog}
self.train_set, self.train_loader, self.val_loader, self.test_loader, self.nclass = make_dataloader(
args.dataset, args.base_size, args.crop_size, args.batch_size,
args.workers, args.overfit, **kwargs)
self.train_set.make_dataset_multiple_of_batchsize(args.batch_size)
if args.architecture == 'deeplab':
print('Using Deeplab')
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
elif args.architecture == 'enet':
print('Using ENet')
model = ENet(num_classes=self.nclass,
encoder_relu=True,
decoder_relu=True)
train_params = [{'params': model.parameters(), 'lr': args.lr}]
elif args.architecture == 'fastscnn':
print('Using FastSCNN')
model = FastSCNN(3, self.nclass)
train_params = [{'params': model.parameters(), 'lr': args.lr}]
if args.optimizer == 'SGD':
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
elif args.optimizer == 'Adam':
optimizer = torch.optim.Adam(train_params,
weight_decay=args.weight_decay)
else:
raise NotImplementedError
if args.use_balanced_weights:
weight = calculate_weights_labels(args.dataset, self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
self.evaluator = Evaluator(self.nclass)
if args.use_lr_scheduler:
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.train_loader))
else:
self.scheduler = None
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError(f"=> no checkpoint found at {args.resume}")
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print(
f'=> loaded checkpoint {args.resume} (epoch {checkpoint["epoch"]})'
)
def __init__(self, exp):
# IoU and pixAcc Metric calculator
self.metric = SegmentationMetric(7)
cfg_path = os.path.join(os.getcwd(), 'config/tusimple_config.yaml')
self.exp_name = exp
self.writer = SummaryWriter('tensorboard/' + self.exp_name)
with open(cfg_path) as file:
cfg = yaml.load(file, Loader=yaml.FullLoader)
self.device = torch.device(cfg['DEVICE'])
self.max_epochs = cfg['TRAIN']['MAX_EPOCHS']
self.dataset_path = cfg['DATASET']['PATH']
# TODO remove this and refactor PROPERLY
self.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cfg['DATASET']['MEAN'], cfg['DATASET']['STD']),
])
mean = cfg['DATASET']['MEAN']
std = cfg['DATASET']['STD']
self.train_transform = Compose(Resize(size=(645,373)), RandomCrop(size=(640,368)), RandomFlip(0.5), Rotation(2), ToTensor(), Normalize(mean=mean, std=std))
self.val_transform = Compose(Resize(size=(640,368)), ToTensor(), Normalize(mean=mean, std=std))
data_kwargs = {
'transform': self.input_transform,
'size': cfg['DATASET']['SIZE'],
}
self.train_dataset = tuSimple(
path=cfg['DATASET']['PATH'],
image_set='train',
transforms=self.train_transform
)
self.val_dataset = tuSimple(
path = cfg['DATASET']['PATH'],
image_set = 'val',
transforms =self.val_transform,
)
self.train_loader = data.DataLoader(
dataset = self.train_dataset,
batch_size = cfg['TRAIN']['BATCH_SIZE'],
shuffle = True,
num_workers = 0,
pin_memory = True,
drop_last = True,
)
self.val_loader = data.DataLoader(
dataset = self.val_dataset,
batch_size = cfg['TRAIN']['BATCH_SIZE'],
shuffle = False,
num_workers = 0,
pin_memory = True,
drop_last = False,
)
self.iters_per_epoch = len(self.train_dataset) // (cfg['TRAIN']['BATCH_SIZE'])
self.max_iters = cfg['TRAIN']['MAX_EPOCHS'] * self.iters_per_epoch
# -------- network --------
weight = [0.4, 1, 1, 1, 1, 1, 1]
self.model = ENet(num_classes=7).to(self.device)
self.optimizer = optim.SGD(
self.model.parameters(),
lr=cfg['OPTIM']['LR'],
weight_decay=cfg['OPTIM']['DECAY'],
momentum=0.9,
)
self.lr_scheduler = get_scheduler(self.optimizer, max_iters=self.max_iters, iters_per_epoch=self.iters_per_epoch)
#self.optimizer = optim.Adam(
# self.model.parameters(),
# lr = cfg['OPTIM']['LR'],
# weight_decay=0,
# )
self.criterion = nn.CrossEntropyLoss(weight=torch.tensor([0.4, 1, 1, 1, 1, 1, 1])).cuda()
self.bce = nn.BCELoss().cuda()
print("Preprocessed sample image dimension:", sample_image.shape)
# Load the required parameters for inference
color_encoding = OrderedDict([
('unlabeled', (0, 0, 0)), ('road', (128, 64, 128)),
('sidewalk', (244, 35, 232)), ('building', (70, 70, 70)),
('wall', (102, 102, 156)), ('fence', (190, 153, 153)),
('pole', (153, 153, 153)), ('traffic_light', (250, 170, 30)),
('traffic_sign', (220, 220, 0)), ('vegetation', (107, 142, 35)),
('terrain', (152, 251, 152)), ('sky', (70, 130, 180)),
('person', (220, 20, 60)), ('rider', (255, 0, 0)), ('car', (0, 0, 142)),
('truck', (0, 0, 70)), ('bus', (0, 60, 100)), ('train', (0, 80, 100)),
('motorcycle', (0, 0, 230)), ('bicycle', (119, 11, 32))
])
num_classes = len(color_encoding)
model = ENet(num_classes).to(device)
# Load the pre-trained weights
model_path = "./save/ENet_Cityscapes/ENet"
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['state_dict'])
print('Model loaded successfully!')
# Run the inference
# If args.test, then showcase how this model works
if not args.test:
model.eval()
sample_image = torch.unsqueeze(sample_image, 0)
with torch.no_grad():
output = model(sample_image)
print("Model output dimension:", output.shape)
def train(train_loader, val_loader, class_weights, class_encoding):
print("\nTraining...\n")
num_classes = len(class_encoding)
# Intialize ENet
model = ENet(num_classes).to(device)
# Check if the network architecture is correct
print(model)
# We are going to use the CrossEntropyLoss loss function as it's most
# frequentely used in classification problems with multiple classes which
# fits the problem. This criterion combines LogSoftMax and NLLLoss.
criterion = nn.CrossEntropyLoss(weight=class_weights)
# ENet authors used Adam as the optimizer
optimizer = optim.Adam(
model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
# Optionally resume from a checkpoint
if args.resume:
model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# Start Training
print()
train = Train(model, train_loader, optimizer, criterion, metric, device)
val = Test(model, val_loader, criterion, metric, device)
for epoch in range(start_epoch, args.epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
lr_updater.step()
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
if (epoch + 1) % 10 == 0 or epoch + 1 == args.epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
args)
return model
def train(train_loader, val_loader, class_weights, class_encoding):
print("\nTraining...\n")
num_classes = len(class_encoding)
# Intialize ENet
model = ENet(num_classes).to(device)
# Check if the network architecture is correct
print(model)
# We are going to use the CrossEntropyLoss loss function as it's most
# frequentely used in classification problems with multiple classes which
# fits the problem. This criterion combines LogSoftMax and NLLLoss.
criterion = nn.CrossEntropyLoss(weight=class_weights)
# ENet authors used Adam as the optimizer
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
# Optionally resume from a checkpoint
if args.resume:
model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# Start Training
print()
train = Train(model, train_loader, optimizer, criterion, metric, device)
val = Test(model, val_loader, criterion, metric, device)
for epoch in range(start_epoch, args.epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step()
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
# Visualization by TensorBoardX
writer.add_scalar('data/train/loss', epoch_loss, epoch)
writer.add_scalar('data/train/mean_IoU', miou, epoch)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
if (epoch + 1) % 1 == 0 or epoch + 1 == args.epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
# Visualization by TensorBoardX
writer.add_scalar('data/val/loss', loss, epoch)
writer.add_scalar('data/val/mean_IoU', miou, epoch)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
args)
# Visualization of the predicted batch in TensorBoard
for i, batch in enumerate(val_loader):
if i == 1:
break
# Get the inputs and labels
inputs = batch[0].to(device)
labels = batch[1].to(device)
# Forward propagation
with torch.no_grad():
predictions = model(inputs)
# Predictions is one-hot encoded with "num_classes" channels.
# Convert it to a single int using the indices where the maximum (1) occurs
_, predictions = torch.max(predictions.data, 1)
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_predictions = utils.batch_transform(
predictions.cpu(), label_to_rgb)
in_training_visualization(model, inputs, labels,
class_encoding, writer, epoch, 'val')
return model
