def main(opt):
writer = SummaryWriter()
log_dir = writer.get_logdir()
os.makedirs(os.path.join(log_dir, "images"), exist_ok=True)
os.makedirs(os.path.join(log_dir, "test"), exist_ok=True)
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
# Initialize generator and discriminator
generator = UNet(opt.sample_num, opt.channels, opt.batch_size, opt.alpha)
discriminator = Discriminator(opt.batch_size, opt.alpha)
generator.to(device=device)
discriminator.to(device=device)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(),
lr=opt.lr_d,
betas=(opt.b1, opt.b2))
if opt.mode == 'train':
generator = train(writer, log_dir, device, generator, discriminator,
optimizer_G, optimizer_D, opt)
test(opt, log_dir, generator=generator)
if opt.mode == 'test':
test(opt, log_dir)
test_moving(opt, log_dir)
def test_moving(opt, log_dir, generator=None):
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
if generator == None:
generator = UNet(opt.sample_num, opt.channels, opt.batch_size,
opt.alpha)
checkpoint = torch.load(opt.load_model, map_location=device)
generator.load_state_dict(checkpoint['g_state_dict'])
del checkpoint
torch.cuda.empty_cache()
generator.to(device)
generator.eval()
dataloader = torch.utils.data.DataLoader(
MyDataset_test_moving(opt),
opt.batch_size,
shuffle=True,
num_workers=opt.num_workers_dataloader)
for i, (imgs, filename) in enumerate(dataloader):
with torch.no_grad():
filename = filename[0].split('/')[-1]
for k in range(len(imgs)):
test_img = generator(imgs[k].to(device))
folder_path = os.path.join(log_dir, "test/%s" % filename)
os.makedirs(folder_path, exist_ok=True)
filename_ = filename + '_' + str(k) + '.png'
test_img = convert_im(test_img,
os.path.join(folder_path, filename_),
nrow=5,
normalize=True,
save_im=True)
def test(opt, log_dir, generator=None):
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
if generator == None:
generator = UNet(opt.sample_num, opt.channels, opt.batch_size,
opt.alpha)
checkpoint = torch.load(opt.load_model, map_location=device)
generator.load_state_dict(checkpoint['g_state_dict'])
del checkpoint
torch.cuda.empty_cache()
generator.to(device)
generator.eval()
dataloader = torch.utils.data.DataLoader(MyDataset_test(opt),
opt.batch_size,
shuffle=True,
num_workers=0)
for i, (imgs, filename) in enumerate(dataloader):
with torch.no_grad():
test_img = generator(imgs.to(device))
filename = filename[0].split('/')[-1]
filename = "test/" + filename + '.png'
test_img = convert_im(test_img,
os.path.join(log_dir, filename),
nrow=5,
normalize=True,
save_im=True)
def main():
device = torch.device("cuda")
inputPath = "D:/VolumeSuperResolution-InputData/sparse-rendering/cleveland60/sample00004_sparse.npy"
input = torch.from_numpy(np.load(inputPath)).to(device=device,
dtype=torch.float32)
print("input:", input.shape)
B, C, H, W = input.shape
mask = torch.abs(input[:, 0:1, :, :])
print("Filled pixels:", end='')
for b in range(B):
print(" %d" % int(torch.sum(mask[b])), end='')
print(" of %d pixels" % (H * W))
print(
"Now test different u-net configuations and report the filled pixels for the first batch"
)
wf = 2
with torch.no_grad():
for depth in range(1, 10):
unet = UNet(in_channels=C,
depth=depth,
wf=wf,
padding='partial',
return_masks=True)
unet.to(device)
output, masks = unet(input, mask)
print("Depth:", depth)
print(" output shape:", output.shape)
print(" Layer: shape + filled pixels")
for i, m in enumerate(masks):
print(
" ", m.shape, " -> %d of %d pixels" %
(int(torch.sum(m[0])), m.shape[2] * m.shape[3]))
print("Save mask")
img = mask[0, 0].cpu().numpy()
plt.imsave("unet-test-mask.png", img, cmap='Greys')
def test(weights_path):
# Get all images in train set
image_names = os.listdir('dataset/train/images/')
image_names = [name for name in image_names if name.endswith(('.jpg', '.JPG', '.png'))]
# Initialize model and transfer to device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = UNet()
model = model.to(device)
model.eval()
# Load weights
model.load_state_dict(torch.load(weights_path, map_location=device))
# Misc info
img_size = 512
# Predict on images
for image_name in tqdm(image_names):
# Load image, prepare for inference
img = cv2.imread(os.path.join('dataset/train/images/', image_name))
img_torch = prepare_image(img, img_size)
with torch.no_grad():
# Get predictions for image
pred_egg_mask, pred_pan_mask = model(img_torch)
# Threshold by 0.5
pred_egg_mask = (torch.sigmoid(pred_egg_mask) >= 0.5).type(pred_egg_mask.dtype)
pred_pan_mask = (torch.sigmoid(pred_pan_mask) >= 0.5).type(pred_pan_mask.dtype)
pred_egg_mask, pred_pan_mask = pred_egg_mask.cpu().detach().numpy(), pred_pan_mask.cpu().detach().numpy()
# Resize masks back to original shape
pred_egg_mask, pred_pan_mask = pred_egg_mask[0][0] * 256, pred_pan_mask[0][0] * 256
pred_egg_mask, pred_pan_mask = postprocess_masks(img, pred_egg_mask, pred_pan_mask)
cv2.imwrite('test_vis/' + image_name[:-4] + '_egg' + image_name[-4:], pred_egg_mask)
cv2.imwrite('test_vis/' + image_name[:-4] + '_pan' + image_name[-4:], pred_pan_mask)
cv2.imwrite('test_vis/' + image_name, img)
def inference(img_dir='data/samples',
img_size=256,
output_dir='outputs',
weights='weights/best_miou.pt',
unet=False):
os.makedirs(output_dir, exist_ok=True)
if unet:
model = UNet(30)
else:
model = DeepLabV3Plus(30)
model = model.to(device)
state_dict = torch.load(weights, map_location=device)
model.load_state_dict(state_dict['model'])
model.eval()
names = [
n for n in os.listdir(img_dir)
if os.path.splitext(n)[1] in ['.jpg', '.jpeg', '.png', '.tiff']
]
with torch.no_grad():
for name in tqdm(names):
path = os.path.join(img_dir, name)
img = cv2.imread(path)
img_shape = img.shape
h = (img.shape[0] / max(img.shape[:2]) * img_size) // 32
w = (img.shape[1] / max(img.shape[:2]) * img_size) // 32
img = cv2.resize(img, (int(w * 32), int(h * 32)))
img = img[:, :, ::-1]
img = img.transpose(2, 0, 1)
img = torch.FloatTensor([img], device=device) / 255.
output = model(img)[0].cpu().numpy().transpose(1, 2, 0)
output = cv2.resize(output, (img_shape[1], img_shape[0]))
output = output.argmax(2)
seg = np.zeros(img_shape, dtype=np.uint8)
for ci, color in enumerate(VOC_COLORMAP):
seg[output == ci] = color
cv2.imwrite(os.path.join(output_dir, name), seg)
model = deeplabv3_resnet101(pretrained=False, progress=True, num_classes=num_classes)
#model = models.segmentation.deeplabv3_resnet101(pretrained=False, progress=True, num_classes=num_classes)
elif args.model == 'FCN':
model = fcn_resnet101(pretrained=False, progress=True, num_classes=num_classes)
else:
model = UNet11(num_classes=num_classes, input_channels=input_channels)
if torch.cuda.is_available():
if args.device_ids:#
device_ids = list(map(int, args.device_ids.split(',')))
else:
device_ids = None
#model = nn.DataParallel(model, device_ids=device_ids).cuda()
model = nn.DataParallel(model, device_ids = device_ids)
model.to(f'cuda:{model.device_ids[0]}')
print('device model',device_ids)
cudnn.benchmark = True
####################Change the files_names ######################################
out_path = Path(('logs/mapping/{}').format(args.dataset_file))
name_file = '_'+ str(int(args.percent*100))+'_percent_' + args.dataset_file
print(args.dataset_file,name_file)
data_all=args.data_all ##file with all the data
data_path = Path(args.dataset_path)
print("data_path:",data_path)
#################################################################################
def train_val(config):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_loader = get_dataloader(img_dir=config.train_img_dir,
mask_dir=config.train_mask_dir,
mode="train",
batch_size=config.batch_size,
num_workers=config.num_workers,
smooth=config.smooth)
val_loader = get_dataloader(img_dir=config.val_img_dir,
mask_dir=config.val_mask_dir,
mode="val",
batch_size=4,
num_workers=config.num_workers)
writer = SummaryWriter(
comment="LR_%f_BS_%d_MODEL_%s_DATA_%s" %
(config.lr, config.batch_size, config.model_type, config.data_type))
if config.model_type == "UNet":
model = UNet()
elif config.model_type == "UNet++":
model = UNetPP()
elif config.model_type == "SEDANet":
model = SEDANet()
elif config.model_type == "RefineNet":
model = rf101()
elif config.model_type == "BASNet":
model = BASNet(n_classes=8)
elif config.model_type == "DANet":
model = DANet(backbone='resnet101',
nclass=config.output_ch,
pretrained=True,
norm_layer=nn.BatchNorm2d)
elif config.model_type == "Deeplabv3+":
model = deeplabv3_plus.DeepLabv3_plus(in_channels=3,
num_classes=8,
backend='resnet101',
os=16,
pretrained=True,
norm_layer=nn.BatchNorm2d)
elif config.model_type == "HRNet_OCR":
model = seg_hrnet_ocr.get_seg_model()
elif config.model_type == "scSEUNet":
model = scSEUNet(pretrained=True, norm_layer=nn.BatchNorm2d)
else:
model = UNet()
if config.iscontinue:
model = torch.load("./exp/24_Deeplabv3+_0.7825757691389714.pth").module
for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
labels = [100, 200, 300, 400, 500, 600, 700, 800]
objects = ['水体', '交通建筑', '建筑', '耕地', '草地', '林地', '裸土', '其他']
if config.optimizer == "sgd":
optimizer = SGD(model.parameters(),
lr=config.lr,
weight_decay=1e-4,
momentum=0.9)
elif config.optimizer == "adamw":
optimizer = adamw.AdamW(model.parameters(), lr=config.lr)
else:
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
# weight = torch.tensor([1, 1.5, 1, 2, 1.5, 2, 2, 1.2]).to(device)
# criterion = nn.CrossEntropyLoss(weight=weight)
criterion = BasLoss()
# scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[25, 30, 35, 40], gamma=0.5)
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.1, patience=5, verbose=True)
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
T_0=15,
eta_min=1e-4)
global_step = 0
max_fwiou = 0
frequency = np.array(
[0.1051, 0.0607, 0.1842, 0.1715, 0.0869, 0.1572, 0.0512, 0.1832])
for epoch in range(config.num_epochs):
epoch_loss = 0.0
cm = np.zeros([8, 8])
print(optimizer.param_groups[0]['lr'])
with tqdm(total=config.num_train,
desc="Epoch %d / %d" % (epoch + 1, config.num_epochs),
unit='img',
ncols=100) as train_pbar:
model.train()
for image, mask in train_loader:
image = image.to(device, dtype=torch.float32)
mask = mask.to(device, dtype=torch.float16)
pred = model(image)
loss = criterion(pred, mask)
epoch_loss += loss.item()
writer.add_scalar('Loss/train', loss.item(), global_step)
train_pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_pbar.update(image.shape[0])
global_step += 1
# if global_step > 10:
# break
# scheduler.step()
print("\ntraining epoch loss: " +
str(epoch_loss / (float(config.num_train) /
(float(config.batch_size)))))
torch.cuda.empty_cache()
val_loss = 0
with torch.no_grad():
with tqdm(total=config.num_val,
desc="Epoch %d / %d validation round" %
(epoch + 1, config.num_epochs),
unit='img',
ncols=100) as val_pbar:
model.eval()
locker = 0
for image, mask in val_loader:
image = image.to(device, dtype=torch.float32)
target = mask.to(device, dtype=torch.long).argmax(dim=1)
mask = mask.cpu().numpy()
pred, _, _, _, _, _, _, _ = model(image)
val_loss += F.cross_entropy(pred, target).item()
pred = pred.cpu().detach().numpy()
mask = semantic_to_mask(mask, labels)
pred = semantic_to_mask(pred, labels)
cm += get_confusion_matrix(mask, pred, labels)
val_pbar.update(image.shape[0])
if locker == 25:
writer.add_images('mask_a/true',
mask[2, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_a/pred',
pred[2, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_b/true',
mask[3, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_b/pred',
pred[3, :, :],
epoch + 1,
dataformats='HW')
locker += 1
# break
miou = get_miou(cm)
fw_miou = (miou * frequency).sum()
scheduler.step()
if fw_miou > max_fwiou:
if torch.__version__ == "1.6.0":
torch.save(model,
config.result_path + "/%d_%s_%.4f.pth" %
(epoch + 1, config.model_type, fw_miou),
_use_new_zipfile_serialization=False)
else:
torch.save(
model, config.result_path + "/%d_%s_%.4f.pth" %
(epoch + 1, config.model_type, fw_miou))
max_fwiou = fw_miou
print("\n")
print(miou)
print("testing epoch loss: " + str(val_loss),
"FWmIoU = %.4f" % fw_miou)
writer.add_scalar('mIoU/val', miou.mean(), epoch + 1)
writer.add_scalar('FWIoU/val', fw_miou, epoch + 1)
writer.add_scalar('loss/val', val_loss, epoch + 1)
for idx, name in enumerate(objects):
writer.add_scalar('iou/val' + name, miou[idx], epoch + 1)
torch.cuda.empty_cache()
writer.close()
print("Training finished")
def train_val(config):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_loader = get_dataloader(img_dir=config.train_img_dir,
mask_dir=config.train_mask_dir,
mode="train",
batch_size=config.batch_size,
num_workers=config.num_workers,
smooth=config.smooth)
val_loader = get_dataloader(img_dir=config.val_img_dir,
mask_dir=config.val_mask_dir,
mode="val",
batch_size=config.batch_size,
num_workers=config.num_workers)
writer = SummaryWriter(
comment="LR_%f_BS_%d_MODEL_%s_DATA_%s" %
(config.lr, config.batch_size, config.model_type, config.data_type))
if config.model_type == "UNet":
model = UNet()
elif config.model_type == "UNet++":
model = UNetPP()
elif config.model_type == "SEDANet":
model = SEDANet()
elif config.model_type == "RefineNet":
model = rf101()
elif config.model_type == "DANet":
# src = "./pretrained/60_DANet_0.8086.pth"
# pretrained_dict = torch.load(src, map_location='cpu').module.state_dict()
# print("load pretrained params from stage 1: " + src)
# pretrained_dict.pop('seg1.1.weight')
# pretrained_dict.pop('seg1.1.bias')
model = DANet(backbone='resnext101',
nclass=config.output_ch,
pretrained=True,
norm_layer=nn.BatchNorm2d)
# model_dict = model.state_dict()
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
elif config.model_type == "Deeplabv3+":
# src = "./pretrained/Deeplabv3+.pth"
# pretrained_dict = torch.load(src, map_location='cpu').module.state_dict()
# print("load pretrained params from stage 1: " + src)
# # print(pretrained_dict.keys())
# for key in list(pretrained_dict.keys()):
# if key.split('.')[0] == "cbr_last":
# pretrained_dict.pop(key)
model = deeplabv3_plus.DeepLabv3_plus(in_channels=3,
num_classes=config.output_ch,
backend='resnet101',
os=16,
pretrained=True,
norm_layer=nn.BatchNorm2d)
# model_dict = model.state_dict()
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
elif config.model_type == "HRNet_OCR":
model = seg_hrnet_ocr.get_seg_model()
elif config.model_type == "scSEUNet":
model = scSEUNet(pretrained=True, norm_layer=nn.BatchNorm2d)
else:
model = UNet()
if config.iscontinue:
model = torch.load("./exp/13_Deeplabv3+_0.7619.pth",
map_location='cpu').module
for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
labels = [1, 2, 3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
objects = [
'水体', '道路', '建筑物', '机场', '停车场', '操场', '普通耕地', '农业大棚', '自然草地', '绿地绿化',
'自然林', '人工林', '自然裸土', '人为裸土', '其它'
]
frequency = np.array([
0.0279, 0.0797, 0.1241, 0.00001, 0.0616, 0.0029, 0.2298, 0.0107,
0.1207, 0.0249, 0.1470, 0.0777, 0.0617, 0.0118, 0.0187
])
if config.optimizer == "sgd":
optimizer = SGD(model.parameters(),
lr=config.lr,
weight_decay=1e-4,
momentum=0.9)
elif config.optimizer == "adamw":
optimizer = adamw.AdamW(model.parameters(), lr=config.lr)
else:
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
# weight = torch.tensor([1, 1.5, 1, 2, 1.5, 2, 2, 1.2]).to(device)
# criterion = nn.CrossEntropyLoss(weight=weight)
if config.smooth == "all":
criterion = LabelSmoothSoftmaxCE()
elif config.smooth == "edge":
criterion = LabelSmoothCE()
else:
criterion = nn.CrossEntropyLoss()
# scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[25, 30, 35, 40], gamma=0.5)
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.1, patience=5, verbose=True)
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
T_0=15,
eta_min=1e-4)
global_step = 0
max_fwiou = 0
for epoch in range(config.num_epochs):
epoch_loss = 0.0
seed = np.random.randint(0, 2, 1)
seed = 0
print("seed is ", seed)
if seed == 1:
train_loader = get_dataloader(img_dir=config.train_img_dir,
mask_dir=config.train_mask_dir,
mode="train",
batch_size=config.batch_size // 2,
num_workers=config.num_workers,
smooth=config.smooth)
val_loader = get_dataloader(img_dir=config.val_img_dir,
mask_dir=config.val_mask_dir,
mode="val",
batch_size=config.batch_size // 2,
num_workers=config.num_workers)
else:
train_loader = get_dataloader(img_dir=config.train_img_dir,
mask_dir=config.train_mask_dir,
mode="train",
batch_size=config.batch_size,
num_workers=config.num_workers,
smooth=config.smooth)
val_loader = get_dataloader(img_dir=config.val_img_dir,
mask_dir=config.val_mask_dir,
mode="val",
batch_size=config.batch_size,
num_workers=config.num_workers)
cm = np.zeros([15, 15])
print(optimizer.param_groups[0]['lr'])
with tqdm(total=config.num_train,
desc="Epoch %d / %d" % (epoch + 1, config.num_epochs),
unit='img',
ncols=100) as train_pbar:
model.train()
for image, mask in train_loader:
image = image.to(device, dtype=torch.float32)
if seed == 0:
pass
elif seed == 1:
image = F.interpolate(image,
size=(384, 384),
mode='bilinear',
align_corners=True)
mask = F.interpolate(mask.float(),
size=(384, 384),
mode='nearest')
if config.smooth == "edge":
mask = mask.to(device, dtype=torch.float32)
else:
mask = mask.to(device, dtype=torch.long).argmax(dim=1)
aux_out, out = model(image)
aux_loss = criterion(aux_out, mask)
seg_loss = criterion(out, mask)
loss = aux_loss + seg_loss
# pred = model(image)
# loss = criterion(pred, mask)
epoch_loss += loss.item()
writer.add_scalar('Loss/train', loss.item(), global_step)
train_pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_pbar.update(image.shape[0])
global_step += 1
# if global_step > 10:
# break
# scheduler.step()
print("\ntraining epoch loss: " +
str(epoch_loss / (float(config.num_train) /
(float(config.batch_size)))))
torch.cuda.empty_cache()
val_loss = 0
with torch.no_grad():
with tqdm(total=config.num_val,
desc="Epoch %d / %d validation round" %
(epoch + 1, config.num_epochs),
unit='img',
ncols=100) as val_pbar:
model.eval()
locker = 0
for image, mask in val_loader:
image = image.to(device, dtype=torch.float32)
target = mask.to(device, dtype=torch.long).argmax(dim=1)
mask = mask.cpu().numpy()
_, pred = model(image)
val_loss += F.cross_entropy(pred, target).item()
pred = pred.cpu().detach().numpy()
mask = semantic_to_mask(mask, labels)
pred = semantic_to_mask(pred, labels)
cm += get_confusion_matrix(mask, pred, labels)
val_pbar.update(image.shape[0])
if locker == 5:
writer.add_images('mask_a/true',
mask[2, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_a/pred',
pred[2, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_b/true',
mask[3, :, :],
epoch + 1,
dataformats='HW')
writer.add_images('mask_b/pred',
pred[3, :, :],
epoch + 1,
dataformats='HW')
locker += 1
# break
miou = get_miou(cm)
fw_miou = (miou * frequency).sum()
scheduler.step()
if True:
if torch.__version__ == "1.6.0":
torch.save(model,
config.result_path + "/%d_%s_%.4f.pth" %
(epoch + 1, config.model_type, fw_miou),
_use_new_zipfile_serialization=False)
else:
torch.save(
model, config.result_path + "/%d_%s_%.4f.pth" %
(epoch + 1, config.model_type, fw_miou))
max_fwiou = fw_miou
print("\n")
print(miou)
print("testing epoch loss: " + str(val_loss),
"FWmIoU = %.4f" % fw_miou)
writer.add_scalar('FWIoU/val', fw_miou, epoch + 1)
writer.add_scalar('loss/val', val_loss, epoch + 1)
for idx, name in enumerate(objects):
writer.add_scalar('iou/val' + name, miou[idx], epoch + 1)
torch.cuda.empty_cache()
writer.close()
print("Training finished")
logging.info(
'Model: UNet + Loss: DynamicContextLoss {}'.format(object_type))
train_file_names = glob.glob(train_path)
random.shuffle(train_file_names)
val_file_names = glob.glob(val_path)
device = torch.device(CUDA_SELECT if torch.cuda.is_available() else "cpu")
model = UNet(num_classes=2, input_channels=1)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
# To handle epoch start number and pretrained weight
epoch_start = '0'
if (use_pretrained):
print("Loading Model {}".format(
os.path.basename(pretrained_model_path)))
model.load_state_dict(torch.load(pretrained_model_path))
epoch_start = os.path.basename(pretrained_model_path).split('.')[0]
print(epoch_start)
trainLoader = DataLoader(DatasetImageMaskLocal(train_file_names,
object_type,
mode='train'),
batch_size=batch_size)
devLoader = DataLoader(
valid_ds = LandmarkDataset(valid_fnames, annotations_path, args.GAUSS_SIGMA,
args.GAUSS_AMPLITUDE)
valid_dl = DataLoader(valid_ds,
args.BATCH_SIZE,
shuffle=False,
num_workers=num_workers)
# Net model
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device(
'cpu')
if args.MODEL == 'unet':
net = UNet(in_ch=1,
out_ch=N_LANDMARKS,
down_drop=args.DOWN_DROP,
up_drop=args.UP_DROP)
net.to(device)
# Optimizer + loss
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(),
lr=args.LEARN_RATE,
weight_decay=args.WEIGHT_DECAY)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=args.OPTIM_PATIENCE,
verbose=True)
def train():
train_loss, train_mre, train_sdr_4mm = 0, 0, 0
train_examples = 0
net.train()
img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
return img
#-----------------------------------------------------------------------------------------------------------#
#-----------------------------------------Model loading-----------------------------------------------------#
# Load segmentation net model
print("\nLoading Segmentation Network")
model = UNet(backbone="resnet18", num_classes=2)
if torch.cuda.is_available():
trained_dict = torch.load(SEGMENTATION_NET_CHECKPOINT)['state_dict']
else:
trained_dict = torch.load(SEGMENTATION_NET_CHECKPOINT, map_location="cpu")['state_dict']
model.load_state_dict(trained_dict, strict=False)
model.to(device)
model.eval()
# Create segmentation object
segmentObj = VideoInference(
model=model,
video_path=0,
input_size=320,
height=OUT_HEIGHT,
width=OUT_WIDTH,
use_cuda=torch.cuda.is_available(),
draw_mode='matting',
)
print("Done Loading Segmentation Network\n")
# Load style transfer net model
def train_val(config):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_loader = get_dataloader(img_dir=config.train_img_dir,
mask_dir=config.train_mask_dir,
mode="train",
batch_size=config.batch_size,
num_workers=config.num_workers)
val_loader = get_dataloader(img_dir=config.val_img_dir,
mask_dir=config.val_mask_dir,
mode="val",
batch_size=config.batch_size,
num_workers=config.num_workers)
writer = SummaryWriter(
comment="LR_%f_BS_%d_MODEL_%s_DATA_%s" %
(config.lr, config.batch_size, config.model_type, config.data_type))
if config.model_type not in [
'UNet', 'R2UNet', 'AUNet', 'R2AUNet', 'SEUNet', 'SEUNet++',
'UNet++', 'DAUNet', 'DANet', 'AUNetR', 'RendDANet', "BASNet"
]:
print('ERROR!! model_type should be selected in supported models')
print('Choose model %s' % config.model_type)
return
if config.model_type == "UNet":
model = UNet()
elif config.model_type == "AUNet":
model = AUNet()
elif config.model_type == "R2UNet":
model = R2UNet()
elif config.model_type == "SEUNet":
model = SEUNet(useCSE=False, useSSE=False, useCSSE=True)
elif config.model_type == "UNet++":
model = UNetPP()
elif config.model_type == "DANet":
model = DANet(backbone='resnet101', nclass=1)
elif config.model_type == "AUNetR":
model = AUNet_R16(n_classes=1, learned_bilinear=True)
elif config.model_type == "RendDANet":
model = RendDANet(backbone='resnet101', nclass=1)
elif config.model_type == "BASNet":
model = BASNet(n_channels=3, n_classes=1)
else:
model = UNet()
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device, dtype=torch.float)
if config.optimizer == "sgd":
optimizer = SGD(model.parameters(),
lr=config.lr,
weight_decay=1e-6,
momentum=0.9)
else:
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.loss == "dice":
criterion = DiceLoss()
elif config.loss == "bce":
criterion = nn.BCELoss()
elif config.loss == "bas":
criterion = BasLoss()
else:
criterion = MixLoss()
scheduler = lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1)
global_step = 0
best_dice = 0.0
for epoch in range(config.num_epochs):
epoch_loss = 0.0
with tqdm(total=config.num_train,
desc="Epoch %d / %d" % (epoch + 1, config.num_epochs),
unit='img') as train_pbar:
model.train()
for image, mask in train_loader:
image = image.to(device, dtype=torch.float)
mask = mask.to(device, dtype=torch.float)
d0, d1, d2, d3, d4, d5, d6, d7 = model(image)
loss = criterion(d0, d1, d2, d3, d4, d5, d6, d7, mask)
epoch_loss += loss.item()
writer.add_scalar('Loss/train', loss.item(), global_step)
train_pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_pbar.update(image.shape[0])
global_step += 1
# if global_step % 100 == 0:
# writer.add_images('masks/true', mask, global_step)
# writer.add_images('masks/pred', d0 > 0.5, global_step)
scheduler.step()
epoch_dice = 0.0
epoch_acc = 0.0
epoch_sen = 0.0
epoch_spe = 0.0
epoch_pre = 0.0
current_num = 0
with tqdm(total=config.num_val,
desc="Epoch %d / %d validation round" %
(epoch + 1, config.num_epochs),
unit='img') as val_pbar:
model.eval()
locker = 0
for image, mask in val_loader:
current_num += image.shape[0]
image = image.to(device, dtype=torch.float)
mask = mask.to(device, dtype=torch.float)
d0, d1, d2, d3, d4, d5, d6, d7 = model(image)
batch_dice = dice_coeff(mask, d0).item()
epoch_dice += batch_dice * image.shape[0]
epoch_acc += get_accuracy(pred=d0, true=mask) * image.shape[0]
epoch_sen += get_sensitivity(pred=d0,
true=mask) * image.shape[0]
epoch_spe += get_specificity(pred=d0,
true=mask) * image.shape[0]
epoch_pre += get_precision(pred=d0, true=mask) * image.shape[0]
if locker == 200:
writer.add_images('masks/true', mask, epoch + 1)
writer.add_images('masks/pred', d0 > 0.5, epoch + 1)
val_pbar.set_postfix(**{'dice (batch)': batch_dice})
val_pbar.update(image.shape[0])
locker += 1
epoch_dice /= float(current_num)
epoch_acc /= float(current_num)
epoch_sen /= float(current_num)
epoch_spe /= float(current_num)
epoch_pre /= float(current_num)
epoch_f1 = get_F1(SE=epoch_sen, PR=epoch_pre)
if epoch_dice > best_dice:
best_dice = epoch_dice
writer.add_scalar('Best Dice/test', best_dice, epoch + 1)
torch.save(
model, config.result_path + "/%s_%s_%d.pth" %
(config.model_type, str(epoch_dice), epoch + 1))
logging.info('Validation Dice Coeff: {}'.format(epoch_dice))
print("epoch dice: " + str(epoch_dice))
writer.add_scalar('Dice/test', epoch_dice, epoch + 1)
writer.add_scalar('Acc/test', epoch_acc, epoch + 1)
writer.add_scalar('Sen/test', epoch_sen, epoch + 1)
writer.add_scalar('Spe/test', epoch_spe, epoch + 1)
writer.add_scalar('Pre/test', epoch_pre, epoch + 1)
writer.add_scalar('F1/test', epoch_f1, epoch + 1)
writer.close()
print("Training finished")
def train_unet(epoch=100):
# Get all images in train set
image_names = os.listdir('dataset/train/images/')
image_names = [name for name in image_names if name.endswith(('.jpg', '.JPG', '.png'))]
# Split into train and validation sets
np.random.shuffle(image_names)
split = int(len(image_names) * 0.9)
train_image_names = image_names[:split]
val_image_names = image_names[split:]
# Create a dataset
train_dataset = EggsPansDataset('dataset/train', train_image_names, mode='train')
val_dataset = EggsPansDataset('dataset/train', val_image_names, mode='val')
# Create a dataloader
train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=False, num_workers=0)
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0)
# Initialize model and transfer to device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = UNet()
model = model.to(device)
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, mode='max', verbose=True)
loss_obj = EggsPansLoss()
metrics_obj = EggsPansMetricIoU()
# Keep best IoU and checkpoint
best_iou = 0.0
# Train epochs
for epoch_idx in range(epoch):
print('Epoch: {:2}/{}'.format(epoch_idx + 1, epoch))
# Reset metrics and loss
loss_obj.reset_loss()
metrics_obj.reset_iou()
# Train phase
model.train()
# Train epoch
pbar = tqdm(train_dataloader)
for imgs, egg_masks, pan_masks in pbar:
# Convert to device
imgs = imgs.to(device)
gt_egg_masks = egg_masks.to(device)
gt_pan_masks = pan_masks.to(device)
# Zero gradients
optim.zero_grad()
# Forward through net, and get the loss
pred_egg_masks, pred_pan_masks = model(imgs)
loss = loss_obj([gt_egg_masks, gt_pan_masks], [pred_egg_masks, pred_pan_masks])
iou = metrics_obj([gt_egg_masks, gt_pan_masks], [pred_egg_masks, pred_pan_masks])
# Compute gradients and compute them
loss.backward()
optim.step()
# Update metrics
pbar.set_description('Loss: {:5.6f}, IoU: {:5.6f}'.format(loss_obj.get_running_loss(),
metrics_obj.get_running_iou()))
print('Validation: ')
# Reset metrics and loss
loss_obj.reset_loss()
metrics_obj.reset_iou()
# Val phase
model.eval()
# Val epoch
pbar = tqdm(val_dataloader)
for imgs, egg_masks, pan_masks in pbar:
# Convert to device
imgs = imgs.to(device)
gt_egg_masks = egg_masks.to(device)
gt_pan_masks = pan_masks.to(device)
with torch.no_grad():
# Forward through net, and get the loss
pred_egg_masks, pred_pan_masks = model(imgs)
loss = loss_obj([gt_egg_masks, gt_pan_masks], [pred_egg_masks, pred_pan_masks])
iou = metrics_obj([gt_egg_masks, gt_pan_masks], [pred_egg_masks, pred_pan_masks])
pbar.set_description('Val Loss: {:5.6f}, IoU: {:5.6f}'.format(loss_obj.get_running_loss(),
metrics_obj.get_running_iou()))
# Save best model
if best_iou < metrics_obj.get_running_iou():
best_iou = metrics_obj.get_running_iou()
torch.save(model.state_dict(), os.path.join('checkpoints/', 'epoch_{}_{:.4f}.pth'.format(
epoch_idx + 1, metrics_obj.get_running_iou())))
# Reduce learning rate on plateau
lr_scheduler.step(metrics_obj.get_running_iou())
print('\n')
print('-'*100)
img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
return img
#-----------------------------------------------------------------------------------------------------------#
#-----------------------------------------Model loading-----------------------------------------------------#
# Load segmentation net model
print("\nLoading Segmentation Network")
segmentation_model = UNet(backbone="resnet18", num_classes=2)
if torch.cuda.is_available():
trained_dict = torch.load(SEGMENTATION_NET_CHECKPOINT)['state_dict']
else:
trained_dict = torch.load(SEGMENTATION_NET_CHECKPOINT, map_location="cpu")['state_dict']
segmentation_model.load_state_dict(trained_dict, strict=False)
segmentation_model.to(device)
segmentation_model.eval()
# Create segmentation object
segmentObj = VideoInference(
model=segmentation_model,
video_path=0,
input_size=320,
height=OUT_HEIGHT,
width=OUT_WIDTH,
use_cuda=torch.cuda.is_available(),
draw_mode='matting',
)
print("Done Loading Segmentation Network\n")
# Load style transfer net model
def main():
parser = argparse.ArgumentParser()
arg = parser.add_argument
arg('--jaccard-weight', type=float, default=1)
arg('--root', type=str, default='runs/debug', help='checkpoint root')
arg('--image-path', type=str, default='data', help='image path')
arg('--batch-size', type=int, default=2)
arg('--n-epochs', type=int, default=100)
arg('--optimizer', type=str, default='Adam', help='Adam or SGD')
arg('--lr', type=float, default=0.001)
arg('--workers', type=int, default=10)
arg('--model',
type=str,
default='UNet16',
choices=[
'UNet', 'UNet11', 'UNet16', 'LinkNet34', 'FCDenseNet57',
'FCDenseNet67', 'FCDenseNet103'
])
arg('--model-weight', type=str, default=None)
arg('--resume-path', type=str, default=None)
arg('--attribute',
type=str,
default='all',
choices=[
'pigment_network', 'negative_network', 'streaks',
'milia_like_cyst', 'globules', 'all'
])
args = parser.parse_args()
## folder for checkpoint
root = Path(args.root)
root.mkdir(exist_ok=True, parents=True)
image_path = args.image_path
#print(args)
if args.attribute == 'all':
num_classes = 5
else:
num_classes = 1
args.num_classes = num_classes
### save initial parameters
print('--' * 10)
print(args)
print('--' * 10)
root.joinpath('params.json').write_text(
json.dumps(vars(args), indent=True, sort_keys=True))
## load pretrained model
if args.model == 'UNet':
model = UNet(num_classes=num_classes)
elif args.model == 'UNet11':
model = UNet11(num_classes=num_classes, pretrained='vgg')
elif args.model == 'UNet16':
model = UNet16(num_classes=num_classes, pretrained='vgg')
elif args.model == 'LinkNet34':
model = LinkNet34(num_classes=num_classes, pretrained=True)
elif args.model == 'FCDenseNet103':
model = FCDenseNet103(num_classes=num_classes)
else:
model = UNet(num_classes=num_classes, input_channels=3)
## multiple GPUs
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.to(device)
## load pretrained model
if args.model_weight is not None:
state = torch.load(args.model_weight)
#epoch = state['epoch']
#step = state['step']
model.load_state_dict(state['model'])
print('--' * 10)
print('Load pretrained model', args.model_weight)
#print('Restored model, epoch {}, step {:,}'.format(epoch, step))
print('--' * 10)
## replace the last layer
## although the model and pre-trained weight have differernt size (the last layer is different)
## pytorch can still load the weight
## I found that the weight for one layer just duplicated for all layers
## therefore, the following code is not necessary
# if args.attribute == 'all':
# model = list(model.children())[0]
# num_filters = 32
# model.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)
# print('--' * 10)
# print('Load pretrained model and replace the last layer', args.model_weight, num_classes)
# print('--' * 10)
# if torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# model.to(device)
## model summary
print_model_summay(model)
## define loss
loss_fn = LossBinary(jaccard_weight=args.jaccard_weight)
## It enables benchmark mode in cudnn.
## benchmark mode is good whenever your input sizes for your network do not vary. This way, cudnn will look for the
## optimal set of algorithms for that particular configuration (which takes some time). This usually leads to faster runtime.
## But if your input sizes changes at each iteration, then cudnn will benchmark every time a new size appears,
## possibly leading to worse runtime performances.
cudnn.benchmark = True
## get train_test_id
train_test_id = get_split()
## train vs. val
print('--' * 10)
print('num train = {}, num_val = {}'.format(
(train_test_id['Split'] == 'train').sum(),
(train_test_id['Split'] != 'train').sum()))
print('--' * 10)
train_transform = DualCompose(
[HorizontalFlip(),
VerticalFlip(),
ImageOnly(Normalize())])
val_transform = DualCompose([ImageOnly(Normalize())])
## define data loader
train_loader = make_loader(train_test_id,
image_path,
args,
train=True,
shuffle=True,
transform=train_transform)
valid_loader = make_loader(train_test_id,
image_path,
args,
train=False,
shuffle=True,
transform=val_transform)
if True:
print('--' * 10)
print('check data')
train_image, train_mask, train_mask_ind = next(iter(train_loader))
print('train_image.shape', train_image.shape)
print('train_mask.shape', train_mask.shape)
print('train_mask_ind.shape', train_mask_ind.shape)
print('train_image.min', train_image.min().item())
print('train_image.max', train_image.max().item())
print('train_mask.min', train_mask.min().item())
print('train_mask.max', train_mask.max().item())
print('train_mask_ind.min', train_mask_ind.min().item())
print('train_mask_ind.max', train_mask_ind.max().item())
print('--' * 10)
valid_fn = validation_binary
###########
## optimizer
if args.optimizer == 'Adam':
optimizer = Adam(model.parameters(), lr=args.lr)
elif args.optimizer == 'SGD':
optimizer = SGD(model.parameters(), lr=args.lr, momentum=0.9)
## loss
criterion = loss_fn
## change LR
scheduler = ReduceLROnPlateau(optimizer,
'min',
factor=0.8,
patience=5,
verbose=True)
##########
## load previous model status
previous_valid_loss = 10
model_path = root / 'model.pt'
if args.resume_path is not None and model_path.exists():
state = torch.load(str(model_path))
epoch = state['epoch']
step = state['step']
model.load_state_dict(state['model'])
epoch = 1
step = 0
try:
previous_valid_loss = state['valid_loss']
except:
previous_valid_loss = 10
print('--' * 10)
print('Restored previous model, epoch {}, step {:,}'.format(
epoch, step))
print('--' * 10)
else:
epoch = 1
step = 0
#########
## start training
log = root.joinpath('train.log').open('at', encoding='utf8')
writer = SummaryWriter()
meter = AllInOneMeter()
#if previous_valid_loss = 10000
print('Start training')
print_model_summay(model)
previous_valid_jaccard = 0
for epoch in range(epoch, args.n_epochs + 1):
model.train()
random.seed()
#jaccard = []
start_time = time.time()
meter.reset()
w1 = 1.0
w2 = 0.5
w3 = 0.5
try:
train_loss = 0
valid_loss = 0
# if epoch == 1:
# freeze_layer_names = get_freeze_layer_names(part='encoder')
# set_freeze_layers(model, freeze_layer_names=freeze_layer_names)
# #set_train_layers(model, train_layer_names=['module.final.weight','module.final.bias'])
# print_model_summay(model)
# elif epoch == 5:
# w1 = 1.0
# w2 = 0.0
# w3 = 0.5
# freeze_layer_names = get_freeze_layer_names(part='encoder')
# set_freeze_layers(model, freeze_layer_names=freeze_layer_names)
# # set_train_layers(model, train_layer_names=['module.final.weight','module.final.bias'])
# print_model_summay(model)
#elif epoch == 3:
# set_train_layers(model, train_layer_names=['module.dec5.block.0.conv.weight','module.dec5.block.0.conv.bias',
# 'module.dec5.block.1.weight','module.dec5.block.1.bias',
# 'module.dec4.block.0.conv.weight','module.dec4.block.0.conv.bias',
# 'module.dec4.block.1.weight','module.dec4.block.1.bias',
# 'module.dec3.block.0.conv.weight','module.dec3.block.0.conv.bias',
# 'module.dec3.block.1.weight','module.dec3.block.1.bias',
# 'module.dec2.block.0.conv.weight','module.dec2.block.0.conv.bias',
# 'module.dec2.block.1.weight','module.dec2.block.1.bias',
# 'module.dec1.conv.weight','module.dec1.conv.bias',
# 'module.final.weight','module.final.bias'])
# print_model_summa zvgf t5y(model)
# elif epoch == 50:
# set_freeze_layers(model, freeze_layer_names=None)
# print_model_summay(model)
for i, (train_image, train_mask,
train_mask_ind) in enumerate(train_loader):
# inputs, targets = variable(inputs), variable(targets)
train_image = train_image.permute(0, 3, 1, 2)
train_mask = train_mask.permute(0, 3, 1, 2)
train_image = train_image.to(device)
train_mask = train_mask.to(device).type(torch.cuda.FloatTensor)
train_mask_ind = train_mask_ind.to(device).type(
torch.cuda.FloatTensor)
# if args.problem_type == 'binary':
# train_mask = train_mask.to(device).type(torch.cuda.FloatTensor)
# else:
# #train_mask = train_mask.to(device).type(torch.cuda.LongTensor)
# train_mask = train_mask.to(device).type(torch.cuda.FloatTensor)
outputs, outputs_mask_ind1, outputs_mask_ind2 = model(
train_image)
#print(outputs.size())
#print(outputs_mask_ind1.size())
#print(outputs_mask_ind2.size())
### note that the last layer in the model is defined differently
# if args.problem_type == 'binary':
# train_prob = F.sigmoid(outputs)
# loss = criterion(outputs, train_mask)
# else:
# #train_prob = outputs
# train_prob = F.sigmoid(outputs)
# loss = torch.tensor(0).type(train_mask.type())
# for feat_inx in range(train_mask.shape[1]):
# loss += criterion(outputs, train_mask)
train_prob = F.sigmoid(outputs)
train_mask_ind_prob1 = F.sigmoid(outputs_mask_ind1)
train_mask_ind_prob2 = F.sigmoid(outputs_mask_ind2)
loss1 = criterion(outputs, train_mask)
#loss1 = F.binary_cross_entropy_with_logits(outputs, train_mask)
#loss2 = nn.BCEWithLogitsLoss()(outputs_mask_ind1, train_mask_ind)
#print(train_mask_ind.size())
#weight = torch.ones_like(train_mask_ind)
#weight[:, 0] = weight[:, 0] * 1
#weight[:, 1] = weight[:, 1] * 14
#weight[:, 2] = weight[:, 2] * 14
#weight[:, 3] = weight[:, 3] * 4
#weight[:, 4] = weight[:, 4] * 4
#weight = weight * train_mask_ind + 1
#weight = weight.to(device).type(torch.cuda.FloatTensor)
loss2 = F.binary_cross_entropy_with_logits(
outputs_mask_ind1, train_mask_ind)
loss3 = F.binary_cross_entropy_with_logits(
outputs_mask_ind2, train_mask_ind)
#loss3 = criterion(outputs_mask_ind2, train_mask_ind)
loss = loss1 * w1 + loss2 * w2 + loss3 * w3
#print(loss1.item(), loss2.item(), loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
step += 1
#jaccard += [get_jaccard(train_mask, (train_prob > 0).float()).item()]
meter.add(train_prob, train_mask, train_mask_ind_prob1,
train_mask_ind_prob2, train_mask_ind, loss1.item(),
loss2.item(), loss3.item(), loss.item())
# print(train_mask.data.shape)
# print(train_mask.data.sum(dim=-2).shape)
# print(train_mask.data.sum(dim=-2).sum(dim=-1).shape)
# print(train_mask.data.sum(dim=-2).sum(dim=-1).sum(dim=0).shape)
# intersection = train_mask.data.sum(dim=-2).sum(dim=-1)
# print(intersection.shape)
# print(intersection.dtype)
# print(train_mask.data.shape[0])
#torch.zeros([2, 4], dtype=torch.float32)
#########################
## at the end of each epoch, evualte the metrics
epoch_time = time.time() - start_time
train_metrics = meter.value()
train_metrics['epoch_time'] = epoch_time
train_metrics['image'] = train_image.data
train_metrics['mask'] = train_mask.data
train_metrics['prob'] = train_prob.data
#train_jaccard = np.mean(jaccard)
#train_auc = str(round(mtr1.value()[0],2))+' '+str(round(mtr2.value()[0],2))+' '+str(round(mtr3.value()[0],2))+' '+str(round(mtr4.value()[0],2))+' '+str(round(mtr5.value()[0],2))
valid_metrics = valid_fn(model, criterion, valid_loader, device,
num_classes)
##############
## write events
write_event(log,
step,
epoch=epoch,
train_metrics=train_metrics,
valid_metrics=valid_metrics)
#save_weights(model, model_path, epoch + 1, step)
#########################
## tensorboard
write_tensorboard(writer,
model,
epoch,
train_metrics=train_metrics,
valid_metrics=valid_metrics)
#########################
## save the best model
valid_loss = valid_metrics['loss1']
valid_jaccard = valid_metrics['jaccard']
if valid_loss < previous_valid_loss:
save_weights(model, model_path, epoch + 1, step, train_metrics,
valid_metrics)
previous_valid_loss = valid_loss
print('Save best model by loss')
if valid_jaccard > previous_valid_jaccard:
save_weights(model, model_path, epoch + 1, step, train_metrics,
valid_metrics)
previous_valid_jaccard = valid_jaccard
print('Save best model by jaccard')
#########################
## change learning rate
scheduler.step(valid_metrics['loss1'])
except KeyboardInterrupt:
# print('--' * 10)
# print('Ctrl+C, saving snapshot')
# save_weights(model, model_path, epoch, step)
# print('done.')
# print('--' * 10)
writer.close()
#return
writer.close()
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Segmentation Example')
parser.add_argument('--test-batch-size', type=int, default=16, metavar='N',
help='input batch size for testing (default: 64)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--mesh_folder', type=str, default="../../mesh_files",
help='path to mesh folder (default: ../../mesh_files)')
parser.add_argument('--ckpt', type=str, default="log/log_f32_cv1_l5_lw/checkpoint_latest.pth.tar_UNet_best.pth.tar")
parser.add_argument('--data_folder', type=str, default="data_small",
help='path to data folder (default: processed_data)')
parser.add_argument('--feat', type=int, default=32, help='filter dimensions')
parser.add_argument('--export_file', type=str, default='samples.npz', help='file name for exporting samples')
parser.add_argument('--in_ch', type=str, default="rgbd", choices=["rgb", "rgbd"], help="input channels")
parser.add_argument('--fold', type=int, choices=[1, 2, 3], default=1, help="choice among 3 fold for cross-validation")
parser.add_argument('--model', type=str, choices=["ResNetDUCHDC", "FCN8s", "UNet"], required=True, help="model of choice")
args = parser.parse_args()
print("%s", repr(args))
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
torch.manual_seed(args.seed)
if args.model == "FCN8s":
model = FCN8s(len(classes), pretrained=False, feat=args.feat, in_ch=len(args.in_ch))
elif args.model == "UNet":
model = UNet(len(classes), len(args.in_ch), feat=args.feat)
model = nn.DataParallel(model)
model.to(device)
# load checkpoint
assert(os.path.isfile(args.ckpt))
print("=> loading checkpoint '{}'".format(args.ckpt))
resume_dict = torch.load(args.ckpt)
start_ep = resume_dict['epoch']
best_miou = resume_dict['best_miou']
def load_my_state_dict(self, state_dict, exclude='none'):
from torch.nn.parameter import Parameter
own_state = self.state_dict()
for name, param in state_dict.items():
if name not in own_state:
continue
if exclude in name:
continue
if isinstance(param, Parameter):
# backwards compatibility for serialized parameters
param = param.data
own_state[name].copy_(param)
load_my_state_dict(model, resume_dict['state_dict'])
print("=> loaded checkpoint '{}' (epoch {} loss {:.03f}) "
.format(args.ckpt, resume_dict['epoch'], best_miou))
testset = SemSegLoader(args.data_folder, "test", fold=args.fold, in_ch=len(args.in_ch))
test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False)
if args.export_file:
export(args, model, test_loader)
with open(args.model+"_v1_files.txt", "w") as f:
f.writelines([l+"\n" for l in testset.rgb_list])
else:
test(args, model, test_loader, epoch, device)