def func_eval(model, criterion, val_dataset, val_loader, post_crf=False):
print("Start validation.\n")
model.eval() # make model evaluation mode
with torch.no_grad():
n_class = 12
total_loss_sum = 0
mIoU_list = []
hist = np.zeros((n_class, n_class)) # confusion matrix
for step, (images, masks, img_info) in enumerate(val_loader):
images = torch.stack(images).to(
CFG.device) # (batch, channel, height, width)
masks = torch.stack(masks).long().to(
CFG.device) # (batch, channel, height, width)
# forward pass (get logits)
logits = model(images)
# loss 계산 (cross entropy loss)
loss = criterion(logits, masks)
total_loss_sum += loss.item() * images.shape[0]
# use softmax to get probability
probs = F.softmax(logits, dim=1)
probs = probs.data.cpu().numpy()
# Postprocessing
if post_crf:
pool = mp.Pool(mp.cpu_count())
images = images.data.cpu().numpy().astype(np.uint8).transpose(
0, 2, 3, 1)
probs = pool.map(dense_crf_wrapper, zip(images, probs))
pool.close()
# get class index which has biggest probability
preds = np.argmax(probs, axis=1)
masks = masks.detach().cpu().numpy()
hist = add_hist(hist, masks, preds, n_class=n_class)
if step == 0:
fig_mask = log_images(masks, preds, img_info)
del images, masks, logits, probs, preds
val_loss = total_loss_sum / len(val_dataset)
acc, acc_cls, mIoU, iu, fwavacc = label_accuracy_score(hist)
recycle = [
'Background', 'UNKNOWN', 'General trash', 'Paper', 'Paper pack',
'Metal', 'Glass', 'Plastic', 'Styrofoam', 'Plastic bag', 'Battery',
'Clothing'
]
mIoU_df = pd.DataFrame({'Recycle Type': recycle, 'IoU': iu})
return val_loss, acc, mIoU, mIoU_df, fig_mask
def train(num_epochs, model, data_loader, val_loader, val_every, device, file_name):
learning_rate = 0.0001
from torch.optim.swa_utils import AveragedModel, SWALR
from torch.optim.lr_scheduler import CosineAnnealingLR
from segmentation_models_pytorch.losses import SoftCrossEntropyLoss, JaccardLoss
from adamp import AdamP
criterion = [SoftCrossEntropyLoss(smooth_factor=0.1), JaccardLoss('multiclass', classes=12)]
optimizer = AdamP(params=model.parameters(), lr=learning_rate, weight_decay=1e-6)
swa_scheduler = SWALR(optimizer, swa_lr=learning_rate)
swa_model = AveragedModel(model)
look = Lookahead(optimizer, la_alpha=0.5)
print('Start training..')
best_miou = 0
for epoch in range(num_epochs):
hist = np.zeros((12, 12))
model.train()
for step, (images, masks, _) in enumerate(data_loader):
loss = 0
images = torch.stack(images) # (batch, channel, height, width)
masks = torch.stack(masks).long() # (batch, channel, height, width)
# gpu 연산을 위해 device 할당
images, masks = images.to(device), masks.to(device)
# inference
outputs = model(images)
for i in criterion:
loss += i(outputs, masks)
# loss 계산 (cross entropy loss)
look.zero_grad()
loss.backward()
look.step()
outputs = torch.argmax(outputs.squeeze(), dim=1).detach().cpu().numpy()
hist = add_hist(hist, masks.detach().cpu().numpy(), outputs, n_class=12)
acc, acc_cls, mIoU, fwavacc = label_accuracy_score(hist)
# step 주기에 따른 loss, mIoU 출력
if (step + 1) % 25 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, mIoU: {:.4f}'.format(
epoch + 1, num_epochs, step + 1, len(data_loader), loss.item(), mIoU))
# validation 주기에 따른 loss 출력 및 best model 저장
if (epoch + 1) % val_every == 0:
avrg_loss, val_miou = validation(epoch + 1, model, val_loader, criterion, device)
if val_miou > best_miou:
print('Best performance at epoch: {}'.format(epoch + 1))
print('Save model in', saved_dir)
best_miou = val_miou
save_model(model, file_name = file_name)
if epoch > 3:
swa_model.update_parameters(model)
swa_scheduler.step()
def train(num_epochs, model, data_loader, val_loader, criterion, optimizer,
saved_dir, val_every, device, file_name, n_class):
print('Start training..')
best_mIoU = 0
for epoch in range(num_epochs):
hist = np.zeros((n_class, n_class))
model.train()
for step, (images, masks, _) in enumerate(data_loader):
# (batch, channel, height, width)
images = torch.stack(images)
# (batch, channel, height, width)
masks = torch.stack(masks).long()
# gpu 연산을 위해 device 할당
images, masks = images.to(device), masks.to(device)
# inference
outputs = model(images)
# loss 계산 (cross entropy loss)
loss = criterion(outputs, masks)
optimizer.zero_grad()
loss.backward()
optimizer.step()
outputs = torch.argmax(outputs.squeeze(),
dim=1).detach().cpu().numpy()
hist = add_hist(hist,
masks.detach().cpu().numpy(),
outputs,
n_class=n_class)
acc, acc_cls, mIoU, fwavacc = label_accuracy_score(hist)
wandb.log({"loss": loss, "mIoU": mIoU}) # wandb 로그출력
# step 주기에 따른 loss 출력
if (step + 1) % 25 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, mIoU:{:.4f}'.
format(epoch + 1, num_epochs, step + 1,
len(train_loader), loss.item(), mIoU))
# validation 주기에 따른 loss 출력 및 best model 저장
# mIoU에 따라 모델 저장
if (epoch + 1) % val_every == 0:
avrg_loss, val_mIoU = validation(epoch + 1, model, val_loader,
criterion, device, n_class)
if val_mIoU > best_mIoU:
print('Best performance at epoch: {}'.format(epoch + 1))
print('Save model in', saved_dir)
best_mIoU = val_mIoU
save_model(model, saved_dir, file_name)
wandb.log({
"val_loss": avrg_loss,
"val_mIoU": val_mIoU,
"best_mIoU": best_mIoU
})
def validation(epoch, model, data_loader, criterion, device, n_class):
print('Start validation #{}'.format(epoch))
model.eval()
with torch.no_grad():
total_loss = 0
cnt = 0
mIoU_list = []
hist = np.zeros((n_class, n_class)) # 중첩을위한 변수
for step, (images, masks, _) in enumerate(data_loader):
# (batch, channel, height, width)
images = torch.stack(images)
# (batch, channel, height, width)
masks = torch.stack(masks).long()
images, masks = images.to(device), masks.to(device)
outputs = model(images)
loss = criterion(outputs, masks)
total_loss += loss
cnt += 1
outputs = torch.argmax(outputs.squeeze(),
dim=1).detach().cpu().numpy()
# 계산을 위한 중첩
hist = add_hist(hist,
masks.detach().cpu().numpy(),
outputs,
n_class=n_class)
# mIoU = label_accuracy_score(
# masks.detach().cpu().numpy(), outputs, n_class=12)[2]
# mIoU_list.append(mIoU)
# mIoU가 전체에대해 계산
acc, acc_cls, mIoU, fwavacc = label_accuracy_score(hist)
avrg_loss = total_loss / cnt
print('Validation #{} Average Loss: {:.4f}, mIoU: {:.4f}'.format(
epoch, avrg_loss, mIoU))
return avrg_loss, mIoU
def psudo_labeling(num_epochs, model, data_loader, val_loader,
unlabeled_loader, criterion, optimizer, device, n_class,
saved_dir, file_name, val_every):
# Instead of using current epoch we use a "step" variable to calculate alpha_weight
# This helps the model converge faster
step = 100
size = 256
transform = A.Compose([A.Resize(256, 256)])
preds_array = np.empty((0, size * size), dtype=np.long)
file_name_list = []
best_mIoU = 0
model.train()
for epoch in range(num_epochs):
hist = np.zeros((n_class, n_class))
for batch_idx, (imgs, image_infos) in enumerate(unlabeled_loader):
# Forward Pass to get the pseudo labels
# --------------------------------------------- test(unlabelse)를 모델에 통과
model.eval()
outs = model(torch.stack(imgs).to(device))
oms = torch.argmax(outs.squeeze(), dim=1).detach().cpu().numpy()
oms = torch.Tensor(oms)
oms = oms.long()
oms = oms.to(device)
# --------------------------------------------- 학습
model.train()
# Now calculate the unlabeled loss using the pseudo label
imgs = torch.stack(imgs)
imgs = imgs.to(device)
# preds_array = preds_array.to(device)
output = model(imgs)
unlabeled_loss = alpha_weight(step) * criterion(output, oms)
# Backpropogate
optimizer.zero_grad()
unlabeled_loss.backward()
optimizer.step()
output = torch.argmax(output.squeeze(),
dim=1).detach().cpu().numpy()
hist = add_hist(hist,
oms.detach().cpu().numpy(),
output,
n_class=n_class)
if (batch_idx + 1) % 25 == 0:
acc, acc_cls, mIoU, fwavacc = label_accuracy_score(hist)
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, mIoU:{:.4f}'.
format(epoch + 1, num_epochs, batch_idx + 1,
len(unlabeled_loader), unlabeled_loss.item(),
mIoU))
# For every 50 batches train one epoch on labeled data
# 50배치마다 라벨데이터를 1 epoch학습
if batch_idx % 50 == 0:
# Normal training procedure
for batch_idx, (images, masks, _) in enumerate(train_loader):
images = torch.stack(images)
# (batch, channel, height, width)
masks = torch.stack(masks).long()
# gpu 연산을 위해 device 할당
images, masks = images.to(device), masks.to(device)
output = model(images)
labeled_loss = criterion(output, masks)
optimizer.zero_grad()
labeled_loss.backward()
optimizer.step()
# Now we increment step by 1
step += 1
if (epoch + 1) % val_every == 0:
avrg_loss, val_mIoU = validation(epoch + 1, model, val_loader,
criterion, device, n_class)
if val_mIoU > best_mIoU:
print('Best performance at epoch: {}'.format(epoch + 1))
print('Save model in', saved_dir)
best_mIoU = val_mIoU
save_model(model, saved_dir, file_name)
wandb.log({
"val_loss": avrg_loss,
"val_mIoU": val_mIoU,
"best_mIoU": best_mIoU
})
model.train()
def pseudo_labeling(num_epochs, model, data_loader, val_loader,
unlabeled_loader, device, val_every, file_name):
# Instead of using current epoch we use a "step" variable to calculate alpha_weight
# This helps the model converge faster
from torch.optim.swa_utils import AveragedModel, SWALR
from segmentation_models_pytorch.losses import SoftCrossEntropyLoss, JaccardLoss
from adamp import AdamP
criterion = [
SoftCrossEntropyLoss(smooth_factor=0.1),
JaccardLoss('multiclass', classes=12)
]
optimizer = AdamP(params=model.parameters(), lr=0.0001, weight_decay=1e-6)
swa_scheduler = SWALR(optimizer, swa_lr=0.0001)
swa_model = AveragedModel(model)
optimizer = Lookahead(optimizer, la_alpha=0.5)
step = 100
size = 256
best_mIoU = 0
model.train()
print('Start Pseudo-Labeling..')
for epoch in range(num_epochs):
hist = np.zeros((12, 12))
for batch_idx, (imgs, image_infos) in enumerate(unlabeled_loader):
# Forward Pass to get the pseudo labels
# --------------------------------------------- test(unlabelse)를 모델에 통과
model.eval()
outs = model(torch.stack(imgs).to(device))
oms = torch.argmax(outs.squeeze(), dim=1).detach().cpu().numpy()
oms = torch.Tensor(oms)
oms = oms.long()
oms = oms.to(device)
# --------------------------------------------- 학습
model.train()
# Now calculate the unlabeled loss using the pseudo label
imgs = torch.stack(imgs)
imgs = imgs.to(device)
# preds_array = preds_array.to(device)
output = model(imgs)
loss = 0
for each in criterion:
loss += each(output, oms)
unlabeled_loss = alpha_weight(step) * loss
# Backpropogate
optimizer.zero_grad()
unlabeled_loss.backward()
optimizer.step()
output = torch.argmax(output.squeeze(),
dim=1).detach().cpu().numpy()
hist = add_hist(hist,
oms.detach().cpu().numpy(),
output,
n_class=12)
if (batch_idx + 1) % 25 == 0:
acc, acc_cls, mIoU, fwavacc = label_accuracy_score(hist)
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, mIoU:{:.4f}'.
format(epoch + 1, num_epochs, batch_idx + 1,
len(unlabeled_loader), unlabeled_loss.item(),
mIoU))
# For every 50 batches train one epoch on labeled data
# 50배치마다 라벨데이터를 1 epoch학습
if batch_idx % 50 == 0:
# Normal training procedure
for batch_idx, (images, masks, _) in enumerate(data_loader):
labeled_loss = 0
images = torch.stack(images)
# (batch, channel, height, width)
masks = torch.stack(masks).long()
# gpu 연산을 위해 device 할당
images, masks = images.to(device), masks.to(device)
output = model(images)
for each in criterion:
labeled_loss += each(output, masks)
optimizer.zero_grad()
labeled_loss.backward()
optimizer.step()
# Now we increment step by 1
step += 1
if (epoch + 1) % val_every == 0:
avrg_loss, val_mIoU = validation(epoch + 1, model, val_loader,
criterion, device)
if val_mIoU > best_mIoU:
print('Best performance at epoch: {}'.format(epoch + 1))
print('Save model in', saved_dir)
best_mIoU = val_mIoU
save_model(model, file_name=file_name)
model.train()
if epoch > 3:
swa_model.update_parameters(model)
swa_scheduler.step()
def evaluate(args, model, criterions, dataloader):
model.eval()
epoch_loss = 0
n_class = 12
example_images = []
with torch.no_grad():
hist = np.zeros((n_class, n_class))
miou_images = []
for images, masks, _ in dataloader:
images = torch.stack(images) # (batch, channel, height, width)
masks = torch.stack(
masks).long() # (batch, channel, height, width)
images, masks = images.to(args.device), masks.to(args.device)
outputs = model(images)
flag = criterions[0]
if flag == "+":
loss = criterions[1](outputs, masks) + criterions[2](outputs,
masks)
elif flag == "-":
loss = criterions[1](outputs, masks) - criterions[2](outputs,
masks)
else:
loss = criterions[1](outputs, masks)
epoch_loss += loss
inputs_np = torch.clone(images).detach().cpu().permute(0, 2, 3,
1).numpy()
inputs_np = denormalize_image(inputs_np,
mean=(0.4611, 0.4403, 0.4193),
std=(0.2107, 0.2074, 0.2157))
example_images.append(
wb_mask(
inputs_np[0],
pred_mask=outputs.argmax(1)[0].detach().cpu().numpy(),
true_mask=masks[0].detach().cpu().numpy(),
))
outputs = torch.argmax(outputs.squeeze(),
dim=1).detach().cpu().numpy()
hist = add_hist(hist,
masks.detach().cpu().numpy(),
outputs,
n_class=n_class)
# 이미지별 miou 저장
miou_list = get_miou(masks.detach().cpu().numpy(),
outputs,
n_class=n_class)
miou_images.extend(miou_list)
# metrics
acc, acc_cls, miou, fwavacc = label_accuracy_score(hist)
# 리더보드 miou
lb_miou = np.nanmean(miou_images)
print(f"acc:{acc:.4f}, acc_cls:{acc_cls:.4f}, fwavacc:{fwavacc:.4f}")
# hist wandb에 저장
summa = hist.sum(1).reshape(-1, 1)
percent = hist / summa
plt.figure(figsize=(10, 10))
sns.heatmap(percent, annot=True, fmt=".2%", annot_kws={"size": 8})
wandb.log({"percent_hist": wandb.Image(plt)}, commit=False)
return (epoch_loss / len(dataloader)), lb_miou, miou, example_images
