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)