def main():
nclass = 7
batch_size = 4
drop_last = True
train_loader, val_loader = make_dataloader(batch_size=batch_size,
drop_last=drop_last)
model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False)
model = load_pretrained_model(model)
evaluator = Evaluator(nclass)
model = model.cuda()
model.eval()
evaluator.reset()
for i, sample in enumerate(val_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
evaluator.add_batch(target, pred)
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
def evaluate_single(gt,pred,num_of_class):
evaluator = Evaluator(num_of_class)
evaluator.reset()
evaluator.add_batch(gt,pred)
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def evaluate_batch(gt_list,pred_list,num_of_class):
evaluator = Evaluator(num_of_class)
evaluator.reset()
for i in range(len(gt_list)):
evaluator.add_batch(gt_list[i],pred_list[i])
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def evaluate(image_file_paths,
gt_file_paths,
model,
nn_input_size,
num_classes,
window_sizes=None,
step_sizes=None,
device=None):
"""
Evaluate model performance
:param image_file_paths: paths to images that will be predicted
:param gt_file_paths: paths to ground truth masks
:param model: model used for prediction
:param nn_input_size: size that the images will be scaled to before feeding them into the nn
:param num_classes: number of classes
:param window_sizes: list of sizes that determine how the image will be cut (for sliding window). Image will be cut
into squares
:param step_sizes: list of step sizes for sliding window
:param device: PyTorch device (cpu or gpu)
:return: list of prediction masks if res_fcn is None, else Nothing
"""
model.eval()
evaluator = Evaluator(num_classes)
def ev(prediction, index):
pred_class_id_mask = np.argmax(prediction, axis=-1)
gt_colour_mask = cv2.imread(gt_file_paths[index])[:, :, 0]
gt_class_id_mask = colour_mask_to_class_id_mask(gt_colour_mask)
# Add sample into evaluator
evaluator.add_batch(gt_class_id_mask, pred_class_id_mask)
with torch.no_grad():
predict.predict(image_file_paths=image_file_paths,
model=model,
nn_input_size=nn_input_size,
res_fcn=ev,
window_sizes=window_sizes,
step_sizes=step_sizes,
device=device)
with np.errstate(divide='ignore', invalid='ignore'):
acc = evaluator.Pixel_Accuracy()
acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
fWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return acc, acc_class, mIoU, fWIoU
def main():
# 配置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 数据路径
train_dir = './data_road/training'
# 超参数
num_epochs = 60
learning_rate = 0.001
img_size = (640, 192)
num_class = 2
SAVE_INTERVAL = 5
evaluator = Evaluator(num_class=num_class)
# 构建Dataset实例
train_data = LaneDataset(data_dir=train_dir, img_size=img_size)
train_loader = DataLoader(train_data, batch_size=2, shuffle=True)
# 构建Model实例
model = FCN8s(n_class=2)
vgg16_path = "models/vgg16_from_caffe.pth"
vgg16 = load_vgg16(model_file=vgg16_path)
model.copy_params_from_vgg16(vgg16)
model = model.to(device)
print("模型加载成功！！！")
# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)
Acc_list = []
Acc_class_list = []
mIoU_list = []
FWIoU_list = []
loss_list = []
# 训练
for epoch in range(num_epochs):
train_loss = 0.0
evaluator.reset()
for i, (image, label) in enumerate(train_loader):
image = image.to(device)
label = label.to(device)
optimizer.zero_grad()
# 前向传播
output = model(image)
output = torch.sigmoid(output)
loss = criterion(output, label)
output = torch.argmax(output, dim=1).cpu().numpy()
label = torch.argmax(label, dim=1).cpu().numpy()
# 后向传播及优化
loss.backward()
train_loss += loss.item()
optimizer.step()
evaluator.add_batch(label, output) # 添加output和label，用于后续评估
# 计算像素准确率、平均像素准确率、平均交并比、频权交并比
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
Acc_list.append(Acc)
Acc_class_list.append(Acc_class)
mIoU_list.append(mIoU)
FWIoU_list.append(FWIoU)
loss_list.append(train_loss/len(train_loader))
# 保存模型
if epoch % SAVE_INTERVAL == 0 or (epoch+1) == num_epochs:
torch.save(model.state_dict(), './models/fcn8s_{}.pth'.format(epoch))
print("Epoch_{}: train_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}"
.format(epoch+1, train_loss/len(train_loader), Acc, Acc_class, mIoU, FWIoU))
# 绘制曲线
draw_figures(loss_list, title='train_loss')
draw_figures(Acc_list, title='Acc')
draw_figures(Acc_class_list, title='Acc_class')
draw_figures(mIoU_list, title='mIoU')
draw_figures(FWIoU_list, title='FWIoU')
print("完成曲线绘制！！！")
class Trainer:
def __init__(self, data_train, data_valid, image_base_dir, instructions):
"""
:param data_train:
:param data_valid:
:param image_base_dir:
:param instructions:
"""
self.image_base_dir = image_base_dir
self.data_valid = data_valid
self.instructions = instructions
# specify model save dir
self.model_name = instructions[STR.MODEL_NAME]
# now = time.localtime()
# start_time = "{}-{}-{}T{}:{}:{}".format(now.tm_year, now.tm_mon, now.tm_mday, now.tm_hour, now.tm_min,
# now.tm_sec)
experiment_folder_path = os.path.join(paths.MODELS_FOLDER_PATH,
self.model_name)
if os.path.exists(experiment_folder_path):
Warning(
"Experiment folder exists already. Files might be overwritten")
os.makedirs(experiment_folder_path, exist_ok=True)
# define saver and save instructions
self.saver = Saver(folder_path=experiment_folder_path,
instructions=instructions)
self.saver.save_instructions()
# define Tensorboard Summary
self.writer = SummaryWriter(log_dir=experiment_folder_path)
nn_input_size = instructions[STR.NN_INPUT_SIZE]
state_dict_file_path = instructions.get(STR.STATE_DICT_FILE_PATH, None)
self.colour_mapping = mapping.get_colour_mapping()
# define transformers for training
crops_per_image = instructions.get(STR.CROPS_PER_IMAGE, 10)
apply_random_cropping = (STR.CROPS_PER_IMAGE in instructions.keys()) and \
(STR.IMAGES_PER_BATCH in instructions.keys())
print("{}applying random cropping".format(
"" if apply_random_cropping else "_NOT_ "))
t = [Normalize()]
if apply_random_cropping:
t.append(
RandomCrop(min_size=instructions.get(STR.CROP_SIZE_MIN, 400),
max_size=instructions.get(STR.CROP_SIZE_MAX, 1000),
crop_count=crops_per_image))
t += [
Resize(nn_input_size),
Flip(p_vertical=0.2, p_horizontal=0.5),
ToTensor()
]
transformations_train = transforms.Compose(t)
# define transformers for validation
transformations_valid = transforms.Compose(
[Normalize(), Resize(nn_input_size),
ToTensor()])
# set up data loaders
dataset_train = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_train,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_train)
# define batch sizes
self.batch_size = instructions[STR.BATCH_SIZE]
if apply_random_cropping:
self.data_loader_train = DataLoader(
dataset=dataset_train,
batch_size=instructions[STR.IMAGES_PER_BATCH],
shuffle=True,
collate_fn=custom_collate)
else:
self.data_loader_train = DataLoader(dataset=dataset_train,
batch_size=self.batch_size,
shuffle=True,
collate_fn=custom_collate)
dataset_valid = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_valid,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_valid)
self.data_loader_valid = DataLoader(dataset=dataset_valid,
batch_size=self.batch_size,
shuffle=False,
collate_fn=custom_collate)
self.num_classes = dataset_train.num_classes()
# define model
print("Building model")
self.model = DeepLab(num_classes=self.num_classes,
backbone=instructions.get(STR.BACKBONE, "resnet"),
output_stride=instructions.get(
STR.DEEPLAB_OUTPUT_STRIDE, 16))
# load weights
if state_dict_file_path is not None:
print("loading state_dict from:")
print(state_dict_file_path)
load_state_dict(self.model, state_dict_file_path)
learning_rate = instructions.get(STR.LEARNING_RATE, 1e-5)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': learning_rate
}, {
'params': self.model.get_10x_lr_params(),
'lr': learning_rate
}]
# choose gpu or cpu
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
if instructions.get(STR.MULTI_GPU, False):
if torch.cuda.device_count() > 1:
print("Using ", torch.cuda.device_count(), " GPUs!")
self.model = nn.DataParallel(self.model)
self.model.to(self.device)
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=0.9,
weight_decay=5e-4,
nesterov=False)
# calculate class weights
if instructions.get(STR.CLASS_STATS_FILE_PATH, None):
class_weights = calculate_class_weights(
instructions[STR.CLASS_STATS_FILE_PATH],
self.colour_mapping,
modifier=instructions.get(STR.LOSS_WEIGHT_MODIFIER, 1.01))
class_weights = torch.from_numpy(class_weights.astype(np.float32))
else:
class_weights = None
self.criterion = SegmentationLosses(
weight=class_weights, cuda=self.device.type != "cpu").build_loss()
# Define Evaluator
self.evaluator = Evaluator(self.num_classes)
# Define lr scheduler
self.scheduler = None
if instructions.get(STR.USE_LR_SCHEDULER, True):
self.scheduler = LR_Scheduler(mode="cos",
base_lr=learning_rate,
num_epochs=instructions[STR.EPOCHS],
iters_per_epoch=len(
self.data_loader_train))
# print information before training start
print("-" * 60)
print("instructions")
pprint(instructions)
model_parameters = sum([p.nelement() for p in self.model.parameters()])
print("Model parameters: {:.2E}".format(model_parameters))
self.best_prediction = 0.0
def train(self, epoch):
self.model.train()
train_loss = 0.0
# create a progress bar
pbar = tqdm(self.data_loader_train)
num_batches_train = len(self.data_loader_train)
# go through each item in the training data
for i, sample in enumerate(pbar):
# set input and target
nn_input = sample[STR.NN_INPUT].to(self.device)
nn_target = sample[STR.NN_TARGET].to(self.device,
dtype=torch.float)
if self.scheduler:
self.scheduler(self.optimizer, i, epoch, self.best_prediction)
# run model
output = self.model(nn_input)
# calc losses
loss = self.criterion(output, nn_target)
# # save step losses
# combined_loss_steps.append(float(loss))
# regression_loss_steps.append(float(regression_loss))
# classification_loss_steps.append(float(classification_loss))
train_loss += loss.item()
pbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_batches_train * epoch)
# calculate gradient and update model weights
loss.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
self.optimizer.step()
self.optimizer.zero_grad()
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print("[Epoch: {}, num images/crops: {}]".format(
epoch, num_batches_train * self.batch_size))
print("Loss: {:.2f}".format(train_loss))
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
pbar = tqdm(self.data_loader_valid, desc='\r')
num_batches_val = len(self.data_loader_valid)
for i, sample in enumerate(pbar):
# set input and target
nn_input = sample[STR.NN_INPUT].to(self.device)
nn_target = sample[STR.NN_TARGET].to(self.device,
dtype=torch.float)
with torch.no_grad():
output = self.model(nn_input)
loss = self.criterion(output, nn_target)
test_loss += loss.item()
pbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
nn_target = nn_target.cpu().numpy()
# Add batch sample into evaluator
self.evaluator.add_batch(nn_target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print("[Epoch: {}, num crops: {}]".format(
epoch, num_batches_val * self.batch_size))
print(
"Acc:{:.2f}, Acc_class:{:.2f}, mIoU:{:.2f}, fwIoU: {:.2f}".format(
Acc, Acc_class, mIoU, FWIoU))
print("Loss: {:.2f}".format(test_loss))
new_pred = mIoU
is_best = new_pred > self.best_prediction
if is_best:
self.best_prediction = new_pred
self.saver.save_checkpoint(self.model, is_best, epoch)
def main():
nclass = 7
lr = 0.01
num_epochs = 400
batch_size = 4
best_pred = 0.0
drop_last = True
train_loader, val_loader = make_dataloader(batch_size=batch_size,
drop_last=drop_last)
model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': lr
}, {
'params': model.get_10x_lr_params(),
'lr': lr * 10
}]
optimizer = torch.optim.SGD(train_params,
momentum=0.9,
weight_decay=5e-4,
nesterov=False)
criterion = SegmentationLosses(weight=None,
cuda=True).build_loss(mode='ce')
evaluator = Evaluator(nclass)
scheduler = LR_Scheduler(mode='poly',
base_lr=lr,
num_epochs=num_epochs,
iters_per_epoch=len(train_loader))
model = model.cuda()
for epoch in range(num_epochs):
train_loss = 0.0
count = 0
model.train()
for i, sample in enumerate(train_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
scheduler(optimizer, i, epoch, best_pred)
optimizer.zero_grad()
output = model(image)
loss = criterion(output, target)
loss.backward()
optimizer.step()
train_loss += loss.item()
count += 1
train_loss = train_loss / count
print(
'Training [Epoch: %d, best_pred: %.4f, numImages: %5d, Loss: %.3f]'
% (epoch, best_pred, i * batch_size + image.data.shape[0],
train_loss))
writer.add_scalar('scalar/loss_seg_train', train_loss, epoch)
model.eval()
evaluator.reset()
test_loss = 0.0
count = 0
for i, sample in enumerate(val_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
loss = criterion(output, target)
test_loss += loss.item()
count += 1
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
evaluator.add_batch(target, pred)
# Fast test during the training
test_loss = test_loss / count
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print('Validation [Epoch: %d, numImages: %5d, Loss: %.3f]' %
(epoch, i * batch_size + image.data.shape[0], test_loss))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
writer.add_scalar('scalar/loss_seg_val', test_loss, epoch)
writer.add_scalar('scalar/Acc_val', Acc, epoch)
writer.add_scalar('scalar/Acc_class_val', Acc_class, epoch)
writer.add_scalar('scalar/mIou_val', mIoU, epoch)
writer.add_scalar('scalar/FWIou_val', FWIoU, epoch)
path = '/home/user/'
if mIoU > best_pred:
best_pred = mIoU
torch.save(model.state_dict(), path + 'model_params.pkl')
print('save the segmentation ' + str(epoch) +
'model, replace the previous parameters')
#%%
import numpy as np
pred = logit.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
print(target.shape) #(2, 513, 513)
print(pred.shape) #(2, 513, 513)
gt_image, pre_image = target, target
evaluator = Evaluator(18)
evaluator.add_batch(gt_image, pre_image.astype('int'))
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
class Trainer(object):
def __init__(self, args):
self.epochs = args.epochs
self.save_interval = args.save_interval
self.train_data = AerialDataset(args, mode='train')
self.train_loader = DataLoader(self.train_data,
batch_size=args.train_batch_size,
shuffle=True,
num_workers=1)
self.model = models.deeplabv3_resnet50(num_classes=args.num_of_class)
#self.model = models.fcn_resnet50(num_classes=args.num_of_class)
self.loss = args.loss
if self.loss == 'CE':
self.criterion = nn.CrossEntropyLoss()
else: #self.loss == 'LS'
self.criterion = LovaszSoftmax()
self.optimizer = torch.optim.AdamW(self.model.parameters())
self.eval_interval = args.eval_interval
self.eval_data = AerialDataset(args, mode='eval')
self.eval_loader = DataLoader(self.eval_data,
batch_size=args.eval_batch_size,
shuffle=False,
num_workers=1)
self.evaluator = Evaluator(args.num_of_class)
self.cuda = args.cuda
if self.cuda is True:
self.model = self.model.cuda()
self.resume = args.resume
if self.resume != None:
if self.cuda:
checkpoint = torch.load(args.resume)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
self.model.load_state_dict(checkpoint['model'])
self.optimizer.load_state_dict(checkpoint['opt'])
self.start_epoch = checkpoint['epoch'] + 1
#start from next epoch
else:
self.start_epoch = 1
#Note: self.start_epoch and self.epochs are only used in run() to schedule training & validation
def run(self):
self.eval(0)
end_epoch = self.start_epoch + self.epochs
for epoch in range(self.start_epoch, end_epoch):
self.train(epoch)
if (epoch - self.start_epoch + 1) % self.save_interval == 0:
saved_dict = {
'epoch': epoch,
'model': self.model.state_dict(),
'opt': self.optimizer.state_dict()
}
torch.save(saved_dict,
'./model_epoch' + str(epoch) + '.pth.tar')
if (epoch - self.start_epoch + 1) % self.eval_interval == 0:
self.eval(epoch)
def train(self, epoch):
self.model.train()
print(f"----------epoch {epoch}----------")
for i, [img, gt] in enumerate(self.train_loader):
print("epoch:", epoch, " batch:", i + 1)
print("img:", img.shape)
print("gt:", gt.shape)
self.optimizer.zero_grad()
if self.cuda:
img, gt = img.cuda(), gt.cuda()
pred = self.model(img)['out']
print("pred:", pred.shape)
loss = self.criterion(pred, gt.long())
print("loss:", loss)
loss.backward()
self.optimizer.step()
def eval(self, epoch):
self.model.eval()
self.evaluator.reset()
if os.path.exists("epoch" + str(epoch)) is False:
os.mkdir("epoch" + str(epoch))
print(f"-----eval epoch {epoch}-----")
for i, [ori, img, gt] in enumerate(self.eval_loader):
print("batch:", i + 1)
print("img:", img.shape)
print("gt:", gt.shape)
if self.cuda:
img = img.cuda()
out = self.model(img)['out']
pred = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
print("pred:", pred.shape)
gt = gt.numpy().squeeze()
#Both gt and pred are numpy array now
self.evaluator.add_batch(gt, pred)
#TODO: eval_batch_size > 1
#colorise
mask = ret2mask(pred)
gt_color = ret2mask(gt)
ori = ori.numpy().squeeze().transpose(1, 2, 0)
cat = np.concatenate((gt_color, ori, mask), axis=1)
cat = Image.fromarray(np.uint8(cat))
cat.save("epoch" + str(epoch) + "/batch" + str(i + 1) + ".png")
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print("Acc:", Acc)
print("Acc_class:", Acc_class)
print("mIoU:", mIoU)
print("FWIoU:", FWIoU)
def test(self):
my_tester = Tester(self)
my_tester.test()
raise NotImplementedError
def main():
# 配置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 数据路径
train_dir = './data_road/training'
valid_dir = './data_road/validing'
# 超参数
num_epochs = 20
learning_rate = 0.001
img_size = (640, 192)
num_class = 2
SAVE_INTERVAL = 5
evaluator = Evaluator(num_class=num_class)
# 构建Dataset实例
train_data = LaneDataset(data_dir=train_dir, img_size=img_size)
train_loader = DataLoader(train_data, batch_size=2, shuffle=True)
valid_data = LaneDataset(data_dir=valid_dir, img_size=img_size)
valid_loader = DataLoader(valid_data, batch_size=2, shuffle=False)
# 构建Model实例
# model = FCN8s(n_class=2)
# vgg16_path = "models/vgg16_from_caffe.pth"
# vgg16 = load_vgg16(model_file=vgg16_path)
# model.copy_params_from_vgg16(vgg16)
# resnet101-fcn
# resnet101_path = './resnet101.pth'
# resnet101 = torchvision.models.resnet101(pretrained=False)
# state_dict = torch.load(resnet101_path)
# resnet101.load_state_dict(state_dict)
# model = ResnetFCN(resnet101, num_classes=2, expansion=4)
# resnet50-fcn
# resnet50_path = './resnet50.pth'
# resnet50 = torchvision.models.resnet50(pretrained=False)
# state_dict = torch.load(resnet50_path)
# resnet50.load_state_dict(state_dict)
# model = ResnetFCN(resnet50, num_classes=2, expansion=4)
# resnet34-fcn
resnet34_path = './resnet34.pth'
resnet34 = torchvision.models.resnet34(pretrained=False)
state_dict = torch.load(resnet34_path)
resnet34.load_state_dict(state_dict)
model = ResnetFCN(resnet34, num_classes=2, expansion=1)
model = model.to(device)
print("模型加载成功！！！")
# 定义损失函数和优化器
# criterion = nn.BCELoss()
criterion = Focal_Loss()
#criterion = nn.BCELoss(weight=torch.tensor([0.3, 0.7])).to(device)
#criterion = nn.CrossEntropyLoss(weight=torch.tensor([0.25, 0.75])).to(device)
#criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(),
lr=learning_rate,
momentum=0.9,
weight_decay=5e-4)
train_Acc_list = []
train_Acc_class_list = []
train_mIoU_list = []
train_FWIoU_list = []
train_loss_list = []
valid_Acc_list = []
valid_Acc_class_list = []
valid_mIoU_list = []
valid_FWIoU_list = []
valid_loss_list = []
# 训练
for epoch in range(num_epochs):
train_loss = 0.0
evaluator.reset()
for i, (image, label) in enumerate(train_loader):
image = image.to(device)
label = label.to(device)
#print(label[:, 1, :, :])
optimizer.zero_grad()
# 前向传播
output = model(image)
loss = criterion(output, label)
#output = torch.sigmoid(output)
output = torch.softmax(output, dim=1)
#loss = criterion(output.transpose(1,3), label.transpose(1,3))
output = torch.argmax(output, dim=1).cpu().numpy()
label = torch.argmax(label, dim=1).cpu().numpy()
#print(output)
# 后向传播及优化
loss.backward()
train_loss += loss.item()
optimizer.step()
evaluator.add_batch(label, output) # 添加output和label，用于后续评估
# 计算像素准确率、平均像素准确率、平均交并比、频权交并比
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
train_Acc_list.append(Acc)
train_Acc_class_list.append(Acc_class)
train_mIoU_list.append(mIoU)
train_FWIoU_list.append(FWIoU)
train_loss_list.append(train_loss / len(train_loader))
evaluator.reset()
# 保存模型
if (epoch + 1) == num_epochs: #epoch % SAVE_INTERVAL == 0 or
torch.save(model.state_dict(), './models/fcn8s_BFL.pth')
print(
"Epoch_{}: train_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}"
.format(epoch + 1, train_loss / len(train_loader), Acc, Acc_class,
mIoU, FWIoU))
# 验证阶段
model.eval()
valid_loss = 0.0
with torch.no_grad():
for i, (image, label) in enumerate(valid_loader):
image = image.to(device)
label = label.to(device)
output = model(image)
loss = criterion(output, label)
#output = torch.sigmoid(output)
output = torch.softmax(output, dim=1)
# loss = criterion(output.transpose(1, 3), label.transpose(1, 3))
output = torch.argmax(output, dim=1).cpu().numpy()
label = torch.argmax(label, dim=1).cpu().numpy()
#print(output)
#print(label)
valid_loss += loss.item()
evaluator.add_batch(label, output) # 添加output和label，用于后续评估
# 计算像素准确率、平均像素准确率、平均交并比、频权交并比
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print(
"Epoch_{}: valid_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}"
.format(epoch + 1, valid_loss / len(valid_loader), Acc, Acc_class,
mIoU, FWIoU))
valid_Acc_list.append(Acc)
valid_Acc_class_list.append(Acc_class)
valid_mIoU_list.append(mIoU)
valid_FWIoU_list.append(FWIoU)
valid_loss_list.append(valid_loss / len(valid_loader))
# 绘制曲线
draw_figures(train_loss_list,
valid_loss_list,
title='loss',
y_label='loss')
draw_figures(train_Acc_list, valid_Acc_list, title='Acc', y_label='Acc')
draw_figures(train_Acc_class_list,
valid_Acc_class_list,
title='Acc_class',
y_label='Acc_class')
draw_figures(train_mIoU_list,
valid_mIoU_list,
title='mIoU',
y_label='mIoU')
draw_figures(train_FWIoU_list,
valid_FWIoU_list,
title='FWIoU',
y_label='FWIoU')
print("完成曲线绘制！！！")
def train(data_set_type, num_classes, batch_size, epochs, use_gpu,
learning_rate, w_decay):
model = get_fcn_model(num_classes, use_gpu)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params=model.parameters(),
lr=learning_rate,
weight_decay=w_decay)
scheduler = lr_scheduler.StepLR(
optimizer, step_size=step_size,
gamma=gamma) # decay LR by a factor of 0.5 every 5 epochs
data_set, data_loader = get_dataset_dataloader(data_set_type, batch_size)
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
epoch_loss = np.zeros((2, epochs))
epoch_acc = np.zeros((2, epochs))
epoch_iou = np.zeros((2, epochs, num_classes))
epoch_mean_iou = np.zeros((2, epochs))
evaluator = Evaluator(num_classes)
for epoch in range(epochs):
logger.info('Epoch {}/{}'.format(epoch + 1, epochs))
logger.info('-' * 28)
for phase_ind, phase in enumerate(['train', 'val']):
if phase == 'train':
model.train()
logger.info(phase)
else:
model.eval()
logger.info(phase)
evaluator.reset()
running_loss = 0.0
running_acc = 0.0
num_of_batches = math.ceil(len(data_set[phase]) / batch_size)
running_iou = np.zeros((num_of_batches, num_classes))
for batch_ind, batch in enumerate(data_loader[phase]):
imgs, targets = batch
imgs = Variable(imgs).float()
imgs = imgs.to(device)
targets = Variable(targets).type(torch.LongTensor)
targets = targets.to(device)
# zero the learnable parameters gradients
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'):
outputs = model(imgs)
loss = criterion(outputs, targets)
if phase == 'train':
loss.backward()
optimizer.step()
# computes loss and acc for current iteration
preds = torch.argmax(outputs, dim=1)
ious = iou(preds, targets, num_classes)
running_loss += loss * imgs.size(0)
running_acc += pixelwise_acc(preds, targets) * imgs.size(0)
running_iou[batch_ind, :] = ious
logger.debug('Batch {} running loss: {}'.format(
batch_ind, running_loss))
# test the iou and pixelwise accuracy using evaluator
preds = preds.cpu().numpy()
targets = targets.cpu().numpy()
evaluator.add_batch(targets, preds)
epoch_loss[phase_ind, epoch] = running_loss / len(data_set[phase])
epoch_acc[phase_ind, epoch] = running_acc / len(data_set[phase])
epoch_iou[phase_ind, epoch] = np.nanmean(running_iou, axis=0)
epoch_mean_iou[phase_ind, epoch] = np.nanmean(epoch_iou[phase_ind,
epoch])
logger.info('{} loss: {:.4f}, acc: {:.4f}, mean iou: {:.6f}'.format(phase,\
epoch_loss[phase_ind, epoch], epoch_acc[phase_ind, epoch],\
epoch_mean_iou[phase_ind, epoch]))
eva_pixel_acc = evaluator.Pixel_Accuracy()
eva_pixel_acc_class = evaluator.Pixel_Accuracy_Class()
eva_mIOU = evaluator.Mean_Intersection_over_Union()
logger.info('{} - Evaluator - acc: {:.4f}, acc class: {:.4f}, mean iou: {:.6f}'.format(phase,\
eva_pixel_acc, eva_pixel_acc_class, eva_mIOU))
if phase == 'val' and epoch_acc[phase_ind, epoch] > best_acc:
best_acc = epoch_acc[phase_ind, epoch]
best_model_wts = copy.deepcopy(model.state_dict())
print()
time_elapsed = time.time() - since
logger.info('Training completed in {}m {}s'.format(int(time_elapsed / 60),\
int(time_elapsed) % 60))
# load best model weights
model.load_state_dict(best_model_wts)
# save numpy results
np.save(os.path.join(score_dir, 'epoch_accuracy'), epoch_acc)
np.save(os.path.join(score_dir, 'epoch_mean_iou'), epoch_mean_iou)
np.save(os.path.join(score_dir, 'epoch_iou'), epoch_iou)
return model