net = UNet(n_channels=3, n_classes=1)
net_dict = net.state_dict()
# # (3) apply pre-trained params in new model
net_old_dict = {k: v for k, v in net_old_dict.items() if k in net_dict}
net_dict.update(net_old_dict) # update params using pre-trained model
net.load_state_dict(net_dict) # update the model
# for name, param in net.named_parameters():
# if param.requires_grad:
# print(name)
# -------------------------- fix parameters related area ---------------------------------- #
# whether fix parameters related area
for name, param in net.named_parameters():
if "down" in name:
param.requires_grad = False
if "inc" in name:
param.requires_grad = False
if have_gpu and args.gpu:
print('Using GPU !')
net = net.cuda()
try:
train_net(image_dir=args.imagedir,
label_dir=args.gt,
boundary_dir=args.bd,
checkpoint_dir=args.checkpoint,
def main():
net = UNet(3, n_classes=3)
#net.load_state_dict(torch.load("./MODEL.pth"))
#print("Model loaded.")
if len(args['snapshot']) == 0:
# net.load_state_dict(torch.load(os.path.join(ckpt_path, 'cityscapes (coarse)-psp_net', 'xx.pth')))
curr_epoch = 1
args['best_record'] = {
'epoch': 0,
'iter': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {
'epoch': int(split_snapshot[1]),
'iter': int(split_snapshot[3]),
'val_loss': float(split_snapshot[5]),
'acc': float(split_snapshot[7]),
'acc_cls': float(split_snapshot[9]),
'mean_iu': float(split_snapshot[11]),
'fwavacc': float(split_snapshot[13])
}
net.cuda().train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(args['longer_size']),
joint_transforms.RandomRotate(10),
joint_transforms.RandomHorizontallyFlip()
])
sliding_crop = joint_transforms.SlidingCrop(args['crop_size'],
args['stride_rate'],
ignore_label)
train_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
visualize = standard_transforms.Compose([
standard_transforms.Scale(args['val_img_display_size']),
standard_transforms.ToTensor()
])
train_set = Retinaimages('training',
joint_transform=train_joint_transform,
sliding_crop=sliding_crop,
transform=train_input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=2,
shuffle=True)
val_set = Retinaimages('validate',
transform=val_input_transform,
sliding_crop=sliding_crop,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=2,
shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True).cuda()
optimizer = optim.SGD([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
args['lr'],
'weight_decay':
args['weight_decay']
}],
momentum=args['momentum'],
nesterov=True)
if len(args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, "_1" + '.txt'),
'w').write(str(args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, args,
val_loader, visualize, val_set)
