def loading_data(): mean_std = cfg.DATA.MEAN_STD train_simul_transform = own_transforms.Compose([ own_transforms.Scale(int(cfg.TRAIN.IMG_SIZE[0] / 0.875)), own_transforms.RandomCrop(cfg.TRAIN.IMG_SIZE), own_transforms.RandomHorizontallyFlip() ]) val_simul_transform = own_transforms.Compose([ own_transforms.Scale(int(cfg.TRAIN.IMG_SIZE[0] / 0.875)), own_transforms.CenterCrop(cfg.TRAIN.IMG_SIZE) ]) img_transform = standard_transforms.Compose([ standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) target_transform = standard_transforms.Compose([ own_transforms.MaskToTensor(), own_transforms.ChangeLabel(cfg.DATA.IGNORE_LABEL, cfg.DATA.NUM_CLASSES - 1) ]) restore_transform = standard_transforms.Compose([ own_transforms.DeNormalize(*mean_std), standard_transforms.ToPILImage() ]) train_set = CityScapes('train', simul_transform=train_simul_transform, transform=img_transform, target_transform=target_transform) train_loader = DataLoader(train_set, batch_size=cfg.TRAIN.BATCH_SIZE, num_workers=16, shuffle=True) val_set = CityScapes('val', simul_transform=val_simul_transform, transform=img_transform, target_transform=target_transform) val_loader = DataLoader(val_set, batch_size=cfg.VAL.BATCH_SIZE, num_workers=16, shuffle=False) return train_loader, val_loader, restore_transform
def main(train_args): net = FCN32VGG(num_classes=voc.num_classes) if len(train_args['snapshot']) == 0: curr_epoch = 1 train_args['best_record'] = { 'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0 } else: print('training resumes from ' + train_args['snapshot']) net.load_state_dict( torch.load( os.path.join(ckpt_path, exp_name, train_args['snapshot']))) split_snapshot = train_args['snapshot'].split('_') curr_epoch = int(split_snapshot[1]) + 1 train_args['best_record'] = { 'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]), 'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]), 'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11]) } curr_epoch = 1 train_args['best_record'] = { 'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0 } mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) input_transform = standard_transforms.Compose([ standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) target_transform = extended_transforms.MaskToTensor() restore_transform = standard_transforms.Compose([ extended_transforms.DeNormalize(*mean_std), standard_transforms.ToPILImage(), ]) visualize = standard_transforms.Compose([ standard_transforms.Resize(400), # was .Scale(400) standard_transforms.CenterCrop(400), standard_transforms.ToTensor() ]) train_set = voc.VOC('train', transform=input_transform, target_transform=target_transform) train_loader = DataLoader(train_set, batch_size=1, num_workers=4, shuffle=True) val_set = voc.VOC('val', transform=input_transform, target_transform=target_transform) val_loader = DataLoader(val_set, batch_size=1, num_workers=4, shuffle=False) # was CrossEntropyLoss2d criterion = torch.nn.CrossEntropyLoss( size_average=False, ignore_index=voc.ignore_label).to(device) optimizer = optim.Adam([{ 'params': [ param for name, param in net.named_parameters() if name[-4:] == 'bias' ], 'lr': 2 * train_args['lr'] }, { 'params': [ param for name, param in net.named_parameters() if name[-4:] != 'bias' ], 'lr': train_args['lr'], 'weight_decay': train_args['weight_decay'] }], betas=(train_args['momentum'], 0.999)) scheduler = ReduceLROnPlateau(optimizer, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True) for epoch in range(curr_epoch, train_args['epoch_num'] + 1): train(train_loader, net, criterion, optimizer, epoch, train_args) val_loss = validate(val_loader, net, criterion, optimizer, epoch, train_args, restore_transform, visualize) scheduler.step(val_loss)
def main(args): print('=' * 10, 'Starting', '=' * 10, '\n') print(device) # Set the seed for reproducing the results random.seed(args.manual_seed) np.random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.manual_seed) cudnn.benchmark = True # Set up results folder if not os.path.exists(os.path.join(ROOT, RESULT, 'saved_val_images')): os.makedirs(os.path.join(ROOT, RESULT, 'saved_val_images')) if not os.path.exists(os.path.join(ROOT, RESULT, 'saved_train_images')): os.makedirs(os.path.join(ROOT, RESULT, 'saved_train_images')) # Setup Dataloader data_loader = get_loader(args.dataset) input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) target_transform = extended_transforms.MaskToTensor() traindata = data_loader('train', n_classes=args.n_classes, transform=input_transform, target_transform=target_transform, do_transform=True) trainloader = data.DataLoader(traindata, batch_size=args.batch_size, num_workers=2, shuffle=True) valdata = data_loader('val', n_classes=args.n_classes, transform=input_transform, target_transform=target_transform) valloader = data.DataLoader(valdata, batch_size=args.batch_size, num_workers=2, shuffle=False) n_classes = traindata.n_classes n_trainsamples = len(traindata) n_iters_per_epoch = np.ceil(n_trainsamples / float(args.batch_size * args.iter_size)) # Setup Model model = get_model(name=args.arch, n_classes=n_classes, ignore_index=traindata.ignore_index, output_stride=args.output_stride, pretrained=args.pretrained, momentum_bn=args.momentum_bn, dprob=args.dprob).to(device) epochs_done = 0 X = [] Y1 = [] Y1_test = [] Y2 = [] Y2_test = [] avg_pixel_acc = 0 mean_class_acc = 0 mIoU = 0 avg_pixel_acc_test = 0 mean_class_acc_test = 0 mIoU_test = 0 best_mIoU = 0 best_epoch = 0 if args.model_path: model_name = args.model_path.split('.') checkpoint_name = model_name[0] + '_optimizer.pkl' checkpoint = torch.load(os.path.join(ROOT, RESULT, checkpoint_name)) optm = checkpoint['optimizer'] model.load_state_dict(checkpoint['state_dict']) split_str = model_name[0].split('_') epochs_done = int(split_str[-1]) saved_loss = pickle.load( open(os.path.join(ROOT, RESULT, "saved_loss.p"), "rb")) saved_accuracy = pickle.load( open(os.path.join(ROOT, RESULT, "saved_accuracy.p"), "rb")) X = saved_loss["X"][:epochs_done] Y = saved_loss["Y"][:epochs_done] Y_test = saved_loss["Y_test"][:epochs_done] avg_pixel_acc = saved_accuracy["P"][:epochs_done, :] mean_class_acc = saved_accuracy["M"][:epochs_done, :] mIoU = saved_accuracy["I"][:epochs_done, :] avg_pixel_acc_test = saved_accuracy["P_test"][:epochs_done, :] mean_class_acc_test = saved_accuracy["M_test"][:epochs_done, :] mIoU_test = saved_accuracy["I_test"][:epochs_done, :] if args.best_model_path: best_model_name = args.best_model_path.split('_') best_mIoU = float(best_model_name[-2]) best_epoch = int(best_model_name[-3]) # Learning rates: For new layers (such as final layer), we set lr to be 10x the learning rate of layers already trained bias_10x_params = filter( lambda x: ('bias' in x[0]) and ('final' in x[0]) and ('conv' in x[0]), model.named_parameters()) bias_10x_params = list(map(lambda x: x[1], bias_10x_params)) bias_params = filter(lambda x: ('bias' in x[0]) and ('final' not in x[0]), model.named_parameters()) bias_params = list(map(lambda x: x[1], bias_params)) nonbias_10x_params = filter( lambda x: (('bias' not in x[0]) or ('bn' in x[0])) and ('final' in x[0]), model.named_parameters()) nonbias_10x_params = list(map(lambda x: x[1], nonbias_10x_params)) nonbias_params = filter( lambda x: ('bias' not in x[0]) and ('final' not in x[0]), model.named_parameters()) nonbias_params = list(map(lambda x: x[1], nonbias_params)) optimizer = torch.optim.SGD([ { 'params': bias_params, 'lr': args.lr }, { 'params': bias_10x_params, 'lr': 20 * args.lr if args.pretrained else args.lr }, { 'params': nonbias_10x_params, 'lr': 10 * args.lr if args.pretrained else args.lr }, { 'params': nonbias_params, 'lr': args.lr }, ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=(args.optim == 'Nesterov')) num_param_groups = 4 # optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # Setting up scheduler if args.model_path and args.restore: # Here we restore all states of optimizer optimizer.load_state_dict(optm) total_iters = n_iters_per_epoch * args.epochs lambda1 = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters) scheduler = lr_scheduler.LambdaLR( optimizer, lr_lambda=[lambda1] * num_param_groups, last_epoch=epochs_done * n_iters_per_epoch) # scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1, last_epoch=epochs_done) else: # scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1) # Here we simply restart the training # if args.T0: # total_iters = args.T0 * n_iters_per_epoch # else: total_iters = ((args.epochs - epochs_done) * n_iters_per_epoch) lambda1 = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters) scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=[lambda1] * num_param_groups) global l_avg, totalclasswise_pixel_acc, totalclasswise_gtpixels, totalclasswise_predpixels global l_avg_test, totalclasswise_pixel_acc_test, totalclasswise_gtpixels_test, totalclasswise_predpixels_test global steps, steps_test criterion_sbd = nn.CrossEntropyLoss(size_average=False, ignore_index=traindata.ignore_index) criterion_lip = nn.CrossEntropyLoss(size_average=False, ignore_index=traindata.ignore_index) criterions = [criterion_sbd, criterion_lip] for epoch in range(epochs_done, args.epochs): print('=' * 10, 'Epoch %d' % (epoch + 1), '=' * 10) l_avg = [0, 0] totalclasswise_pixel_acc = [0, 0] totalclasswise_gtpixels = [0, 0] totalclasswise_predpixels = [0, 0] l_avg_test = [0, 0] totalclasswise_pixel_acc_test = [0, 0] totalclasswise_gtpixels_test = [0, 0] totalclasswise_predpixels_test = [0, 0] steps = [0, 0] steps_test = [0, 0] # scheduler.step() train(model, optimizer, criterions, trainloader, epoch, scheduler, traindata) val(model, criterions, valloader, epoch, valdata) # save the model every 5 epochs if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1: if (epoch + 1) > 5: os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}.pkl".format(args.arch, args.dataset, epoch - 4))) os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer.pkl".format( args.arch, args.dataset, epoch - 4))) torch.save( model, os.path.join( ROOT, RESULT, "{}_{}_{}.pkl".format(args.arch, args.dataset, epoch + 1))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer.pkl".format(args.arch, args.dataset, epoch + 1))) # remove old loss & accuracy files if os.path.isfile(os.path.join(ROOT, RESULT, "saved_loss.p")): os.remove(os.path.join(ROOT, RESULT, "saved_loss.p")) if os.path.isfile(os.path.join(ROOT, RESULT, "saved_accuracy.p")): os.remove(os.path.join(ROOT, RESULT, "saved_accuracy.p")) # save train and validation loss X.append(epoch + 1) Y1.append(l_avg[0] / steps[0]) Y1_test.append(l_avg_test[0] / steps_test[0]) Y2.append(l_avg[1] / steps[1]) Y2_test.append(l_avg_test[1] / steps_test[1]) saved_loss = { "X": X, "Y1": Y1, "Y2": Y2, "Y1_test": Y1_test, "Y2_test": Y2_test } pickle.dump(saved_loss, open(os.path.join(ROOT, RESULT, "saved_loss.p"), "wb")) # pixel accuracy totalclasswise_pixel_acc[0] = totalclasswise_pixel_acc[0].reshape( (-1, n_classes[0])).astype(np.float32) totalclasswise_gtpixels[0] = totalclasswise_gtpixels[0].reshape( (-1, n_classes[0])) totalclasswise_predpixels[0] = totalclasswise_predpixels[0].reshape( (-1, n_classes[0])) totalclasswise_pixel_acc_test[0] = totalclasswise_pixel_acc_test[ 0].reshape((-1, n_classes[0])).astype(np.float32) totalclasswise_gtpixels_test[0] = totalclasswise_gtpixels_test[ 0].reshape((-1, n_classes[0])) totalclasswise_predpixels_test[0] = totalclasswise_predpixels_test[ 0].reshape((-1, n_classes[0])) totalclasswise_pixel_acc[1] = totalclasswise_pixel_acc[1].reshape( (-1, n_classes[1])).astype(np.float32) totalclasswise_gtpixels[1] = totalclasswise_gtpixels[1].reshape( (-1, n_classes[1])) totalclasswise_predpixels[1] = totalclasswise_predpixels[1].reshape( (-1, n_classes[1])) totalclasswise_pixel_acc_test[1] = totalclasswise_pixel_acc_test[ 1].reshape((-1, n_classes[1])).astype(np.float32) totalclasswise_gtpixels_test[1] = totalclasswise_gtpixels_test[ 1].reshape((-1, n_classes[1])) totalclasswise_predpixels_test[1] = totalclasswise_predpixels_test[ 1].reshape((-1, n_classes[1])) if isinstance(avg_pixel_acc, list): avg_pixel_acc[0] = np.vstack( (avg_pixel_acc[0], np.sum(totalclasswise_pixel_acc[0], axis=1) / np.sum(totalclasswise_gtpixels[0], axis=1))) mean_class_acc[0] = np.vstack( (mean_class_acc[0], np.mean(totalclasswise_pixel_acc[0] / totalclasswise_gtpixels[0], axis=1))) mIoU[0] = np.vstack( (mIoU[0], np.mean( totalclasswise_pixel_acc[0] / (totalclasswise_gtpixels[0] + totalclasswise_predpixels[0] - totalclasswise_pixel_acc[0]), axis=1))) avg_pixel_acc[1] = np.vstack( (avg_pixel_acc[1], np.sum(totalclasswise_pixel_acc[1], axis=1) / np.sum(totalclasswise_gtpixels[1], axis=1))) mean_class_acc[1] = np.vstack( (mean_class_acc[1], np.mean(totalclasswise_pixel_acc[1] / totalclasswise_gtpixels[1], axis=1))) mIoU[1] = np.vstack( (mIoU[1], np.mean( totalclasswise_pixel_acc[1] / (totalclasswise_gtpixels[1] + totalclasswise_predpixels[1] - totalclasswise_pixel_acc[1]), axis=1))) avg_pixel_acc_test[0] = np.vstack( (avg_pixel_acc_test[0], np.sum(totalclasswise_pixel_acc_test[0], axis=1) / np.sum(totalclasswise_gtpixels_test[0], axis=1))) mean_class_acc_test[0] = np.vstack( (mean_class_acc_test[0], np.mean(totalclasswise_pixel_acc_test[0] / totalclasswise_gtpixels_test[0], axis=1))) mIoU_test[0] = np.vstack( (mIoU_test[0], np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1))) avg_pixel_acc_test[1] = np.vstack( (avg_pixel_acc_test[1], np.sum(totalclasswise_pixel_acc_test[1], axis=1) / np.sum(totalclasswise_gtpixels_test[1], axis=1))) mean_class_acc_test[1] = np.vstack( (mean_class_acc_test[1], np.mean(totalclasswise_pixel_acc_test[1] / totalclasswise_gtpixels_test[1], axis=1))) mIoU_test[1] = np.vstack( (mIoU_test[1], np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1))) else: avg_pixel_acc = [] mean_class_acc = [] mIoU = [] avg_pixel_acc.append( np.sum(totalclasswise_pixel_acc[0], axis=1) / np.sum(totalclasswise_gtpixels[0], axis=1)) mean_class_acc.append( np.mean(totalclasswise_pixel_acc[0] / totalclasswise_gtpixels[0], axis=1)) mIoU.append( np.mean( totalclasswise_pixel_acc[0] / (totalclasswise_gtpixels[0] + totalclasswise_predpixels[0] - totalclasswise_pixel_acc[0]), axis=1)) avg_pixel_acc.append( np.sum(totalclasswise_pixel_acc[1], axis=1) / np.sum(totalclasswise_gtpixels[1], axis=1)) mean_class_acc.append( np.mean(totalclasswise_pixel_acc[1] / totalclasswise_gtpixels[1], axis=1)) mIoU.append( np.mean( totalclasswise_pixel_acc[1] / (totalclasswise_gtpixels[1] + totalclasswise_predpixels[1] - totalclasswise_pixel_acc[1]), axis=1)) avg_pixel_acc_test = [] mean_class_acc_test = [] mIoU_test = [] avg_pixel_acc_test.append( np.sum(totalclasswise_pixel_acc_test[0], axis=1) / np.sum(totalclasswise_gtpixels_test[0], axis=1)) mean_class_acc_test.append( np.mean(totalclasswise_pixel_acc_test[0] / totalclasswise_gtpixels_test[0], axis=1)) mIoU_test.append( np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1)) avg_pixel_acc_test.append( np.sum(totalclasswise_pixel_acc_test[1], axis=1) / np.sum(totalclasswise_gtpixels_test[1], axis=1)) mean_class_acc_test.append( np.mean(totalclasswise_pixel_acc_test[1] / totalclasswise_gtpixels_test[1], axis=1)) mIoU_test.append( np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1)) saved_accuracy = { "X": X, "P1": avg_pixel_acc[0], "P2": avg_pixel_acc[1], "M1": mean_class_acc[0], "M2": mean_class_acc[1], "I1": mIoU[0], "I2": mIoU[1], "P1_test": avg_pixel_acc_test[0], "P2_test": avg_pixel_acc_test[1], "M1_test": mean_class_acc_test[0], "M2_test": mean_class_acc_test[1], "I1_test": mIoU_test[0], "I2_test": mIoU_test[1] } pickle.dump(saved_accuracy, open(os.path.join(ROOT, RESULT, "saved_accuracy.p"), "wb")) # print validation mIoU of both tasks this_mIoU1 = np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1)[0] this_mIoU2 = np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1)[0] print('Val: mIoU_sbd = {}, mIoU_lip = {}'.format( this_mIoU1, this_mIoU2))
import numpy as np from datasets.dataset import yaogan from torch.utils.data import DataLoader import transforms ####计算weight权重 target_transform = transforms.MaskToTensor() input_transfom = transforms.MaskToTensor() if __name__ == "__main__": yaogan = yaogan(mode='train', input_transform=input_transfom, target_transform=target_transform) train_loader = DataLoader(yaogan, batch_size=24, num_workers=8) weight = np.zeros((1, 8)).squeeze(0).astype(np.uint) for index, data in enumerate(train_loader): img, label = data label = np.array(label).reshape(256 * 256 * 24, 1).squeeze(1).astype(np.uint8) weight_ = np.bincount(label, minlength=8) weight = weight_ + weight print(weight) pix_partion = weight / weight.sum() ##计算像素占比 weight_all = 1 / pix_partion ##计算权重 max = max(weight_all) weight_all = weight_all / max ##权重归一 print(weight_all) print(pix_partion) # import os
def main(train_args, model): print(train_args) net = model.cuda() if len(train_args['snapshot']) == 0: curr_epoch = 1 train_args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0} else: print('training resumes from ' + train_args['snapshot']) net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, train_args['snapshot']))) split_snapshot = train_args['snapshot'].split('_') curr_epoch = int(split_snapshot[1]) + 1 train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]), 'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]), 'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])} net.train() mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) input_transform = standard_transforms.Compose([ standard_transforms.Pad(200), standard_transforms.CenterCrop(320), standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) train_transform = standard_transforms.Compose([ standard_transforms.Pad(200), standard_transforms.CenterCrop(320), extended_transforms.MaskToTensor()]) target_transform = extended_transforms.MaskToTensor() restore_transform = standard_transforms.Compose([ extended_transforms.DeNormalize(*mean_std), standard_transforms.ToPILImage(), ]) visualize = standard_transforms.Compose([ standard_transforms.Resize(400), standard_transforms.CenterCrop(400), standard_transforms.ToTensor() ]) train_set = VOCSegmentation(root='./', image_set='train', transform=input_transform, target_transform=train_transform) train_loader = DataLoader(train_set, batch_size=1, num_workers=4, shuffle=True) val_set = VOCSegmentation(root='./', image_set='val', transform=input_transform, target_transform=train_transform) val_loader = DataLoader(val_set, batch_size=1, num_workers=4, shuffle=False) #criterion = CrossEntropyLoss().cuda()#2d(size_average=False, ignore_index=voc.ignore_label).cuda() criterion = CrossEntropyLoss(size_average=False, ignore_index=255).cuda() optimizer = optim.SGD([ {'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'], 'lr': 2 * train_args['lr']}, {'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'], 'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']} ], momentum=train_args['momentum']) """optimizer = optim.Adam([ {'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'], 'lr': 2 * train_args['lr']}, {'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'], 'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']} ], betas=(train_args['momentum'], 0.999))""" if len(train_args['snapshot']) > 0: optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + train_args['snapshot']))) optimizer.param_groups[0]['lr'] = 2 * train_args['lr'] optimizer.param_groups[1]['lr'] = train_args['lr'] """check_mkdir(ckpt_path) check_mkdir(os.path.join(ckpt_path, exp_name)) open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')""" scheduler = ReduceLROnPlateau(optimizer, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True) for epoch in range(curr_epoch, train_args['epoch_num'] + 1): train(train_loader, net, criterion, optimizer, epoch, train_args) val_loss, imges = validate(val_loader, net, criterion, optimizer, epoch, train_args, restore_transform, visualize) #imges.show() scheduler.step(val_loss) return imges
def main(train_args): check_mkdir(os.path.join(train_args['ckpt_path'], args['exp'])) check_mkdir( os.path.join(train_args['ckpt_path'], args['exp'], train_args['exp_name'])) model = DeepLabV3('1') # print(model) device = torch.device("cuda") num_gpu = list(range(torch.cuda.device_count())) """###############------use gpu--------###############""" if args['use_gpu']: ts = time.time() print(torch.cuda.current_device()) print(torch.cuda.get_device_name(0)) model = nn.DataParallel(model, device_ids=num_gpu) model = model.to(device) print("Finish cuda loading ,time elapsed {}", format(time.time() - ts)) else: print("please check your gpu device,start training on cpu") """###############-------中间开始训练--------###############""" if len(train_args['snapshot']) == 0: curr_epoch = 1 train_args['best_record'] = { 'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0 } # model.apply(weights_init) else: print("train resume from " + train_args['snapshot']) state_dict = torch.load( os.path.join(train_args['ckpt_path'], args['exp'], train_args['exp_name'], train_args['snapshot'])) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] new_state_dict[name] = v model.load_state_dict(new_state_dict) # model.load_state_dict( # torch.load(os.path.join(train_args['ckpt_path'],args['exp'],train_args['exp_name'], train_args['snapshot']))) split_snapshot = train_args['snapshot'].split('_') curr_epoch = int(split_snapshot[1]) + 1 train_args['best_record'] = { 'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]), 'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]), 'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11]) } model.train() mean_std = ([0.485, 0.456, 0.406, 0.450], [0.229, 0.224, 0.225, 0.225]) """#################---数据增强和数据变换等操作------########""" input_transform = standard_transforms.Compose([ standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) ##Nomorlized target_transform = extended_transforms.MaskToTensor() # target to tensor joint_transform = joint_transforms.Compose([ joint_transforms.RandomHorizontallyFlip(), joint_transforms.RandomCrop((256, 256), padding=0), joint_transforms.Rotate(degree=90) ]) ###data_augment restore = standard_transforms.Compose([ extended_transforms.DeNormalize(*mean_std), extended_transforms.channel_4_to_channel_3(4, 3), ##默认3通道如果四通道会转成三通道 standard_transforms.ToPILImage(), ]) # DeNomorlized,出来是pil图片了 visualize = standard_transforms.Compose([ standard_transforms.Resize(256), standard_transforms.CenterCrop(256), ##中心裁剪,此处可以删除 standard_transforms.ToTensor() ]) # resize 大小之后转tensor """#################---数据加载------########""" train_set = yaogan(mode='train', cls=train_args['training_cls'], joint_transform=None, input_transform=input_transform, target_transform=target_transform) train_loader = DataLoader(train_set, batch_size=train_args['batch_size'], num_workers=train_args['num_works'], shuffle=True) val_set = yaogan(mode='val', cls=train_args['training_cls'], input_transform=input_transform, target_transform=target_transform) val_loader = DataLoader(val_set, batch_size=1, num_workers=train_args['num_works'], shuffle=False) # test_set=yaogan(mode='test',cls=train_args['training_cls'],joint_transform=None, # input_transform=input_transform,target_transform=None) # test_loader=DataLoader(test_set,batch_size=1, # num_workers=train_args['num_works'], shuffle=False) optimizer = optim.Adadelta(model.parameters(), lr=train_args['lr']) ##define a weighted loss (0weight for 0 label) # weight=[0.09287939 ,0.02091968 ,0.02453979, 0.25752962 ,0.33731845, 1., # 0.09518322, 0.52794035 ,0.24298112 ,0.02657369, 0.15057124 ,0.36864611, # 0.25835161,0.16672758 ,0.40728756 ,0.00751281] """###############-------训练数据权重--------###############""" if train_args['weight'] is not None: weight = [0.1, 1.] weight = torch.Tensor(weight) else: weight = None criterion = nn.CrossEntropyLoss(weight=weight, reduction='elementwise_mean', ignore_index=-100).to(device) # criterion=nn.BCELoss(weight=weight,reduction='elementwise_mean').cuda() check_mkdir(train_args['ckpt_path']) check_mkdir(os.path.join(train_args['ckpt_path'], args['exp'])) check_mkdir( os.path.join(train_args['ckpt_path'], args['exp'], train_args['exp_name'])) open( os.path.join(train_args['ckpt_path'], args['exp'], train_args['exp_name'], str(time.time()) + '.txt'), 'w').write(str(train_args) + '\n\n') """###############-------start training--------###############""" for epoch in range(curr_epoch, train_args['epoch_num'] + 1): adjust_lr(optimizer, epoch) train(train_loader, model, criterion, optimizer, epoch, train_args, device) val_loss = validate(val_loader, model, criterion, optimizer, restore, epoch, train_args, visualize, device) writer.close()
def main(args): print('=' * 10, 'Starting', '=' * 10, '\n') print(device) # Set the seed for reproducing the results random.seed(args.manual_seed) np.random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.manual_seed) cudnn.benchmark = True # Set up results folder if not os.path.exists(os.path.join(ROOT, RESULT, 'saved_val_images')): os.makedirs(os.path.join(ROOT, RESULT, 'saved_val_images')) if not os.path.exists(os.path.join(ROOT, RESULT, 'saved_train_images')): os.makedirs(os.path.join(ROOT, RESULT, 'saved_train_images')) # Setup Dataloader data_loader = get_loader(args.dataset) input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) target_transform = extended_transforms.MaskToTensor() traindata = data_loader('train', n_classes=args.n_classes, transform=input_transform, target_transform=target_transform, do_transform=True) trainloader = data.DataLoader(traindata, batch_size=args.batch_size, num_workers=1, shuffle=True) valdata = data_loader('val', n_classes=args.n_classes, transform=input_transform, target_transform=target_transform) valloader = data.DataLoader(valdata, batch_size=args.batch_size, num_workers=1, shuffle=False) n_classes = traindata.n_classes n_trainsamples_total = len(traindata) n_trainsamples_lip = n_trainsamples_total / 2 n_trainsamples_sbd = n_trainsamples_total / 2 n_iters_per_epoch_common = np.ceil(n_trainsamples_total / float(20)) n_iters_per_epoch_lip = np.ceil(n_trainsamples_lip / float(10)) n_iters_per_epoch_sbd = np.ceil(n_trainsamples_sbd / float(10)) print('# total training samples = {}'.format(n_trainsamples_total)) print('# sbd training samples = {}'.format(n_trainsamples_sbd)) print('# lip training samples = {}'.format(n_trainsamples_lip)) # Setup Model model = get_model(name=args.arch, n_classes=n_classes, ignore_index=traindata.ignore_index, output_stride=args.output_stride, pretrained=args.pretrained, momentum_bn=args.momentum_bn, dprob=args.dprob).to(device) X = [] Y1 = [] Y1_test = [] Y2 = [] Y2_test = [] avg_pixel_acc = 0 mean_class_acc = 0 mIoU = 0 avg_pixel_acc_test = 0 mean_class_acc_test = 0 mIoU_test = 0 # Learning rates: For new layers (such as final layer), we set lr to be 10x the learning rate of layers already trained common_bias = filter(lambda x: ('bias' in x[0]) and ('final' not in x[0]), model.named_parameters()) common_bias = list(map(lambda x: x[1], common_bias)) common_nonbias = filter( lambda x: ('bias' not in x[0]) and ('final' not in x[0]), model.named_parameters()) common_nonbias = list(map(lambda x: x[1], common_nonbias)) final1_bias = filter( lambda x: ('bias' in x[0]) and ('final1' in x[0]) and ('conv' in x[0]), model.named_parameters()) final1_bias = list(map(lambda x: x[1], final1_bias)) final1_nonbias = filter( lambda x: (('bias' not in x[0]) or ('bn' in x[0])) and ('final1' in x[0]), model.named_parameters()) final1_nonbias = list(map(lambda x: x[1], final1_nonbias)) final2_bias = filter( lambda x: ('bias' in x[0]) and ('final2' in x[0]) and ('conv' in x[0]), model.named_parameters()) final2_bias = list(map(lambda x: x[1], final2_bias)) final2_nonbias = filter( lambda x: (('bias' not in x[0]) or ('bn' in x[0])) and ('final2' in x[0]), model.named_parameters()) final2_nonbias = list(map(lambda x: x[1], final2_nonbias)) optimizer_common = torch.optim.SGD([{ 'params': common_bias, 'lr': args.lr }, { 'params': common_nonbias, 'lr': args.lr }], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=(args.optim == 'Nesterov')) optimizer_sbd = torch.optim.SGD( [{ 'params': final1_bias, 'lr': 20 * args.lr if args.pretrained else args.lr }, { 'params': final1_nonbias, 'lr': 10 * args.lr if args.pretrained else args.lr }], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=(args.optim == 'Nesterov')) optimizer_lip = torch.optim.SGD( [{ 'params': final2_bias, 'lr': 20 * args.lr if args.pretrained else args.lr }, { 'params': final2_nonbias, 'lr': 10 * args.lr if args.pretrained else args.lr }], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=(args.optim == 'Nesterov')) optimizers = [optimizer_common, optimizer_sbd, optimizer_lip] # Setting up scheduler total_iters_common = (args.epochs * n_iters_per_epoch_common) total_iters_sbd = (args.epochs * n_iters_per_epoch_sbd) total_iters_lip = (args.epochs * n_iters_per_epoch_lip) lambda_common = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters_common) lambda_sbd = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters_sbd) lambda_lip = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters_lip) scheduler_common = lr_scheduler.LambdaLR(optimizer_common, lr_lambda=[lambda_common] * 2) scheduler_sbd = lr_scheduler.LambdaLR(optimizer_sbd, lr_lambda=[lambda_sbd] * 2) scheduler_lip = lr_scheduler.LambdaLR(optimizer_lip, lr_lambda=[lambda_lip] * 2) schedulers = [scheduler_common, scheduler_sbd, scheduler_lip] global l_avg, totalclasswise_pixel_acc, totalclasswise_gtpixels, totalclasswise_predpixels global l_avg_test, totalclasswise_pixel_acc_test, totalclasswise_gtpixels_test, totalclasswise_predpixels_test global steps, steps_test global bug_counter bug_counter = 0 scheduler_common.step() scheduler_sbd.step() scheduler_lip.step() counter_sizes = [20, 10, 10] global counters counters = [0, 0, 0] criterion_sbd = nn.CrossEntropyLoss(size_average=False, ignore_index=traindata.ignore_index) criterion_lip = nn.CrossEntropyLoss(size_average=False, ignore_index=traindata.ignore_index) criterions = [criterion_sbd, criterion_lip] for epoch in range(args.epochs): print('=' * 10, 'Epoch %d' % (epoch + 1), '=' * 10) l_avg = [0, 0] totalclasswise_pixel_acc = [0, 0] totalclasswise_gtpixels = [0, 0] totalclasswise_predpixels = [0, 0] l_avg_test = [0, 0] totalclasswise_pixel_acc_test = [0, 0] totalclasswise_gtpixels_test = [0, 0] totalclasswise_predpixels_test = [0, 0] steps = [0, 0] steps_test = [0, 0] train(model, optimizers, criterions, trainloader, epoch, schedulers, traindata, counter_sizes) val(model, criterions, valloader, epoch, valdata) # save the model every 5 epochs if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1: if (epoch + 1) > 5: os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}.pkl".format(args.arch, args.dataset, epoch - 4))) os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer0.pkl".format( args.arch, args.dataset, epoch - 4))) os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer1.pkl".format( args.arch, args.dataset, epoch - 4))) os.remove( os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer2.pkl".format( args.arch, args.dataset, epoch - 4))) torch.save( model, os.path.join( ROOT, RESULT, "{}_{}_{}.pkl".format(args.arch, args.dataset, epoch + 1))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizers[0].state_dict() }, os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer0.pkl".format(args.arch, args.dataset, epoch + 1))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizers[1].state_dict() }, os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer1.pkl".format(args.arch, args.dataset, epoch + 1))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizers[2].state_dict() }, os.path.join( ROOT, RESULT, "{}_{}_{}_optimizer2.pkl".format(args.arch, args.dataset, epoch + 1))) # remove old loss & accuracy files if os.path.isfile(os.path.join(ROOT, RESULT, "saved_loss.p")): os.remove(os.path.join(ROOT, RESULT, "saved_loss.p")) if os.path.isfile(os.path.join(ROOT, RESULT, "saved_accuracy.p")): os.remove(os.path.join(ROOT, RESULT, "saved_accuracy.p")) # save train and validation loss X.append(epoch + 1) Y1.append(l_avg[0] / steps[0]) Y1_test.append(l_avg_test[0] / steps_test[0]) Y2.append(l_avg[1] / steps[1]) Y2_test.append(l_avg_test[1] / steps_test[1]) saved_loss = { "X": X, "Y1": Y1, "Y2": Y2, "Y1_test": Y1_test, "Y2_test": Y2_test } pickle.dump(saved_loss, open(os.path.join(ROOT, RESULT, "saved_loss.p"), "wb")) # pixel accuracy totalclasswise_pixel_acc[0] = totalclasswise_pixel_acc[0].reshape( (-1, n_classes[0])).astype(np.float32) totalclasswise_gtpixels[0] = totalclasswise_gtpixels[0].reshape( (-1, n_classes[0])) totalclasswise_predpixels[0] = totalclasswise_predpixels[0].reshape( (-1, n_classes[0])) totalclasswise_pixel_acc_test[0] = totalclasswise_pixel_acc_test[ 0].reshape((-1, n_classes[0])).astype(np.float32) totalclasswise_gtpixels_test[0] = totalclasswise_gtpixels_test[ 0].reshape((-1, n_classes[0])) totalclasswise_predpixels_test[0] = totalclasswise_predpixels_test[ 0].reshape((-1, n_classes[0])) totalclasswise_pixel_acc[1] = totalclasswise_pixel_acc[1].reshape( (-1, n_classes[1])).astype(np.float32) totalclasswise_gtpixels[1] = totalclasswise_gtpixels[1].reshape( (-1, n_classes[1])) totalclasswise_predpixels[1] = totalclasswise_predpixels[1].reshape( (-1, n_classes[1])) totalclasswise_pixel_acc_test[1] = totalclasswise_pixel_acc_test[ 1].reshape((-1, n_classes[1])).astype(np.float32) totalclasswise_gtpixels_test[1] = totalclasswise_gtpixels_test[ 1].reshape((-1, n_classes[1])) totalclasswise_predpixels_test[1] = totalclasswise_predpixels_test[ 1].reshape((-1, n_classes[1])) if isinstance(avg_pixel_acc, list): avg_pixel_acc[0] = np.vstack( (avg_pixel_acc[0], np.sum(totalclasswise_pixel_acc[0], axis=1) / np.sum(totalclasswise_gtpixels[0], axis=1))) mean_class_acc[0] = np.vstack( (mean_class_acc[0], np.mean(totalclasswise_pixel_acc[0] / totalclasswise_gtpixels[0], axis=1))) mIoU[0] = np.vstack( (mIoU[0], np.mean( totalclasswise_pixel_acc[0] / (totalclasswise_gtpixels[0] + totalclasswise_predpixels[0] - totalclasswise_pixel_acc[0]), axis=1))) avg_pixel_acc[1] = np.vstack( (avg_pixel_acc[1], np.sum(totalclasswise_pixel_acc[1], axis=1) / np.sum(totalclasswise_gtpixels[1], axis=1))) mean_class_acc[1] = np.vstack( (mean_class_acc[1], np.mean(totalclasswise_pixel_acc[1] / totalclasswise_gtpixels[1], axis=1))) mIoU[1] = np.vstack( (mIoU[1], np.mean( totalclasswise_pixel_acc[1] / (totalclasswise_gtpixels[1] + totalclasswise_predpixels[1] - totalclasswise_pixel_acc[1]), axis=1))) avg_pixel_acc_test[0] = np.vstack( (avg_pixel_acc_test[0], np.sum(totalclasswise_pixel_acc_test[0], axis=1) / np.sum(totalclasswise_gtpixels_test[0], axis=1))) mean_class_acc_test[0] = np.vstack( (mean_class_acc_test[0], np.mean(totalclasswise_pixel_acc_test[0] / totalclasswise_gtpixels_test[0], axis=1))) mIoU_test[0] = np.vstack( (mIoU_test[0], np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1))) avg_pixel_acc_test[1] = np.vstack( (avg_pixel_acc_test[1], np.sum(totalclasswise_pixel_acc_test[1], axis=1) / np.sum(totalclasswise_gtpixels_test[1], axis=1))) mean_class_acc_test[1] = np.vstack( (mean_class_acc_test[1], np.mean(totalclasswise_pixel_acc_test[1] / totalclasswise_gtpixels_test[1], axis=1))) mIoU_test[1] = np.vstack( (mIoU_test[1], np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1))) else: avg_pixel_acc = [] mean_class_acc = [] mIoU = [] avg_pixel_acc.append( np.sum(totalclasswise_pixel_acc[0], axis=1) / np.sum(totalclasswise_gtpixels[0], axis=1)) mean_class_acc.append( np.mean(totalclasswise_pixel_acc[0] / totalclasswise_gtpixels[0], axis=1)) mIoU.append( np.mean( totalclasswise_pixel_acc[0] / (totalclasswise_gtpixels[0] + totalclasswise_predpixels[0] - totalclasswise_pixel_acc[0]), axis=1)) avg_pixel_acc.append( np.sum(totalclasswise_pixel_acc[1], axis=1) / np.sum(totalclasswise_gtpixels[1], axis=1)) mean_class_acc.append( np.mean(totalclasswise_pixel_acc[1] / totalclasswise_gtpixels[1], axis=1)) mIoU.append( np.mean( totalclasswise_pixel_acc[1] / (totalclasswise_gtpixels[1] + totalclasswise_predpixels[1] - totalclasswise_pixel_acc[1]), axis=1)) avg_pixel_acc_test = [] mean_class_acc_test = [] mIoU_test = [] avg_pixel_acc_test.append( np.sum(totalclasswise_pixel_acc_test[0], axis=1) / np.sum(totalclasswise_gtpixels_test[0], axis=1)) mean_class_acc_test.append( np.mean(totalclasswise_pixel_acc_test[0] / totalclasswise_gtpixels_test[0], axis=1)) mIoU_test.append( np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1)) avg_pixel_acc_test.append( np.sum(totalclasswise_pixel_acc_test[1], axis=1) / np.sum(totalclasswise_gtpixels_test[1], axis=1)) mean_class_acc_test.append( np.mean(totalclasswise_pixel_acc_test[1] / totalclasswise_gtpixels_test[1], axis=1)) mIoU_test.append( np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1)) saved_accuracy = { "X": X, "P1": avg_pixel_acc[0], "P2": avg_pixel_acc[1], "M1": mean_class_acc[0], "M2": mean_class_acc[1], "I1": mIoU[0], "I2": mIoU[1], "P1_test": avg_pixel_acc_test[0], "P2_test": avg_pixel_acc_test[1], "M1_test": mean_class_acc_test[0], "M2_test": mean_class_acc_test[1], "I1_test": mIoU_test[0], "I2_test": mIoU_test[1] } pickle.dump(saved_accuracy, open(os.path.join(ROOT, RESULT, "saved_accuracy.p"), "wb")) # print validation mIoU of both tasks this_mIoU1 = np.mean(totalclasswise_pixel_acc_test[0] / (totalclasswise_gtpixels_test[0] + totalclasswise_predpixels_test[0] - totalclasswise_pixel_acc_test[0]), axis=1)[0] this_mIoU2 = np.mean(totalclasswise_pixel_acc_test[1] / (totalclasswise_gtpixels_test[1] + totalclasswise_predpixels_test[1] - totalclasswise_pixel_acc_test[1]), axis=1)[0] print('Val: mIoU_sbd = {}, mIoU_lip = {}'.format( this_mIoU1, this_mIoU2))
'checkpoint': '', 'save_interval': 10000, 'rotate_degree': 10, 'deep_base': True, 'aux_weight': 0.1, 'print_aux': True, } ckpt_path = './ckpt' exp_name = 'model' img_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.3598, 0.3653, 0.3662], [0.2573, 0.2663, 0.2756]) ]) mask_transform = extend_transforms.MaskToTensor() train_joint_transform = extend_transforms.Compose([ extend_transforms.RandomScale(), extend_transforms.RandomSizedRatio(760, 842, 274, 304), extend_transforms.RandomRotate(args['rotate_degree']), extend_transforms.RandomCrop(args['train_crop_size']), ]) train_set = culane.CULANE('train', joint_transform=train_joint_transform, transform=img_transform, mask_transform=mask_transform) train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=10,
def main(args): print('=' * 10, 'Starting', '=' * 10, '\n') print(device) # Set the seed for reproducing the results random.seed(args.manual_seed) np.random.seed(args.manual_seed) torch.manual_seed(args.manual_seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.manual_seed) cudnn.benchmark = True # Set up results folder if not os.path.exists( os.path.join(ROOT_ADDRESS, 'results_parts/saved_val_images')): os.makedirs( os.path.join(ROOT_ADDRESS, 'results_parts/saved_val_images')) if not os.path.exists( os.path.join(ROOT_ADDRESS, 'results_parts/saved_train_images')): os.makedirs( os.path.join(ROOT_ADDRESS, 'results_parts/saved_train_images')) # Setup Dataloader data_loader = get_loader(args.dataset) # data_path = get_data_path(args.dataset) # traindata = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) # trainloader = data.DataLoader(traindata, batch_size=args.batch_size, num_workers=7, shuffle=True) # valdata = data_loader(data_path, split="val", is_transform=False, img_size=(args.img_rows, args.img_cols)) # valloader = data.DataLoader(valdata, batch_size=args.batch_size, num_workers=7, shuffle=False) mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) input_transform = standard_transforms.Compose([ standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) target_transform = extended_transforms.MaskToTensor() traindata = data_loader('train', transform=input_transform, target_transform=target_transform, do_transform=True) trainloader = data.DataLoader(traindata, batch_size=args.batch_size, num_workers=1, shuffle=True) valdata = data_loader('val', transform=input_transform, target_transform=target_transform) valloader = data.DataLoader(valdata, batch_size=args.batch_size, num_workers=1, shuffle=False) n_classes = traindata.n_classes n_trainsamples = len(traindata) n_iters_per_epoch = np.ceil(n_trainsamples / float(args.batch_size * args.iter_size)) # Setup Model model = get_model(args.arch, n_classes, ignore_index=traindata.ignore_index, output_stride=args.output_stride).to(device) epochs_done = 0 X = [] Y = [] Y_test = [] avg_pixel_acc = 0 mean_class_acc = 0 mIoU = 0 avg_pixel_acc_test = 0 mean_class_acc_test = 0 mIoU_test = 0 best_mIoU = 0 best_epoch = 0 if args.model_path: model_name = args.model_path.split('.') checkpoint_name = model_name[0] + '_optimizer.pkl' checkpoint = torch.load(checkpoint_name) optm = checkpoint['optimizer'] model.load_state_dict(checkpoint['state_dict']) split_str = model_name[0].split('_') epochs_done = int(split_str[-1]) saved_loss = pickle.load( open(os.path.join(ROOT_ADDRESS, "results_parts/saved_loss.p"), "rb")) saved_accuracy = pickle.load( open(os.path.join(ROOT_ADDRESS, "results_parts/saved_accuracy.p"), "rb")) X = saved_loss["X"][:epochs_done] Y = saved_loss["Y"][:epochs_done] Y_test = saved_loss["Y_test"][:epochs_done] avg_pixel_acc = saved_accuracy["P"][:epochs_done, :] mean_class_acc = saved_accuracy["M"][:epochs_done, :] mIoU = saved_accuracy["I"][:epochs_done, :] avg_pixel_acc_test = saved_accuracy["P_test"][:epochs_done, :] mean_class_acc_test = saved_accuracy["M_test"][:epochs_done, :] mIoU_test = saved_accuracy["I_test"][:epochs_done, :] if args.best_model_path: best_model_name = args.best_model_path.split('.')[0].split('_') best_mIoU = float(best_model_name[-2]) best_epoch = int(best_model_name[-3]) # Learning rates: For new layers (such as final layer), we set lr to be 10x the learning rate of layers already trained bias_10x_params = filter( lambda x: ('bias' in x[0]) and ('final' in x[0]) and ('conv' in x[0]), model.named_parameters()) bias_10x_params = list(map(lambda x: x[1], bias_10x_params)) bias_params = filter(lambda x: ('bias' in x[0]) and ('final' not in x[0]), model.named_parameters()) bias_params = list(map(lambda x: x[1], bias_params)) nonbias_10x_params = filter( lambda x: (('bias' not in x[0]) or ('bn' in x[0])) and ('final' in x[0]), model.named_parameters()) nonbias_10x_params = list(map(lambda x: x[1], nonbias_10x_params)) nonbias_params = filter( lambda x: ('bias' not in x[0]) and ('final' not in x[0]), model.named_parameters()) nonbias_params = list(map(lambda x: x[1], nonbias_params)) optimizer = torch.optim.SGD([ { 'params': bias_params, 'lr': args.lr }, { 'params': bias_10x_params, 'lr': args.lr }, { 'params': nonbias_10x_params, 'lr': args.lr }, { 'params': nonbias_params, 'lr': args.lr }, ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=(args.optim == 'Nesterov')) num_param_groups = 4 # Setting up scheduler if args.model_path and args.restore: # Here we restore all states of optimizer optimizer.load_state_dict(optm) total_iters = n_iters_per_epoch * args.epochs lambda1 = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters) scheduler = lr_scheduler.LambdaLR( optimizer, lr_lambda=[lambda1] * num_param_groups, last_epoch=epochs_done * n_iters_per_epoch) else: # Here we simply restart the training if args.T0: total_iters = args.T0 * n_iters_per_epoch else: total_iters = ((args.epochs - epochs_done) * n_iters_per_epoch) lambda1 = lambda step: 0.5 + 0.5 * math.cos(np.pi * step / total_iters) scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=[lambda1] * num_param_groups) global l_avg, totalclasswise_pixel_acc, totalclasswise_gtpixels, totalclasswise_predpixels global l_avg_test, totalclasswise_pixel_acc_test, totalclasswise_gtpixels_test, totalclasswise_predpixels_test global steps, steps_test scheduler.step() criterion = nn.CrossEntropyLoss(size_average=False, ignore_index=traindata.ignore_index) print('=' * 10, 'Entering epoch loop', '=' * 10, '\n') for epoch in range(epochs_done, args.epochs): print('=' * 10, 'Epoch %d' % (epoch + 1), '=' * 10) l_avg = 0 totalclasswise_pixel_acc = 0 totalclasswise_gtpixels = 0 totalclasswise_predpixels = 0 l_avg_test = 0 totalclasswise_pixel_acc_test = 0 totalclasswise_gtpixels_test = 0 totalclasswise_predpixels_test = 0 steps = 0 steps_test = 0 train(model, optimizer, criterion, trainloader, epoch, scheduler, traindata) val(model, criterion, valloader, epoch, valdata) # save the model every 5 epochs if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1: if (epoch + 1) > 5: os.remove( os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}.pkl".format( args.arch, args.dataset, epoch - 4))) os.remove( os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_optimizer.pkl".format( args.arch, args.dataset, epoch - 4))) torch.save( model, os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}.pkl".format( args.arch, args.dataset, epoch + 1))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_optimizer.pkl".format( args.arch, args.dataset, epoch + 1))) # remove old loss & accuracy files if os.path.isfile( os.path.join(ROOT_ADDRESS, "results_parts/saved_loss.p")): os.remove(os.path.join(ROOT_ADDRESS, "results_parts/saved_loss.p")) if os.path.isfile( os.path.join(ROOT_ADDRESS, "results_parts/saved_accuracy.p")): os.remove( os.path.join(ROOT_ADDRESS, "results_parts/saved_accuracy.p")) # save train and validation loss X.append(epoch + 1) Y.append(l_avg / steps) Y_test.append(l_avg_test / steps_test) saved_loss = {"X": X, "Y": Y, "Y_test": Y_test} pickle.dump( saved_loss, open(os.path.join(ROOT_ADDRESS, "results_parts/saved_loss.p"), "wb")) # pixel accuracy totalclasswise_pixel_acc = totalclasswise_pixel_acc.reshape( (-1, n_classes)).astype(np.float32) totalclasswise_gtpixels = totalclasswise_gtpixels.reshape( (-1, n_classes)) totalclasswise_predpixels = totalclasswise_predpixels.reshape( (-1, n_classes)) totalclasswise_pixel_acc_test = totalclasswise_pixel_acc_test.reshape( (-1, n_classes)).astype(np.float32) totalclasswise_gtpixels_test = totalclasswise_gtpixels_test.reshape( (-1, n_classes)) totalclasswise_predpixels_test = totalclasswise_predpixels_test.reshape( (-1, n_classes)) if isinstance(avg_pixel_acc, np.ndarray): avg_pixel_acc = np.vstack( (avg_pixel_acc, np.sum(totalclasswise_pixel_acc, axis=1) / np.sum(totalclasswise_gtpixels, axis=1))) mean_class_acc = np.vstack( (mean_class_acc, np.mean(totalclasswise_pixel_acc / totalclasswise_gtpixels, axis=1))) mIoU = np.vstack( (mIoU, np.mean(totalclasswise_pixel_acc / (totalclasswise_gtpixels + totalclasswise_predpixels - totalclasswise_pixel_acc), axis=1))) avg_pixel_acc_test = np.vstack( (avg_pixel_acc_test, np.sum(totalclasswise_pixel_acc_test, axis=1) / np.sum(totalclasswise_gtpixels_test, axis=1))) mean_class_acc_test = np.vstack( (mean_class_acc_test, np.mean(totalclasswise_pixel_acc_test / totalclasswise_gtpixels_test, axis=1))) mIoU_test = np.vstack((mIoU_test, np.mean(totalclasswise_pixel_acc_test / (totalclasswise_gtpixels_test + totalclasswise_predpixels_test - totalclasswise_pixel_acc_test), axis=1))) else: avg_pixel_acc = np.sum(totalclasswise_pixel_acc, axis=1) / np.sum( totalclasswise_gtpixels, axis=1) mean_class_acc = np.mean(totalclasswise_pixel_acc / totalclasswise_gtpixels, axis=1) mIoU = np.mean( totalclasswise_pixel_acc / (totalclasswise_gtpixels + totalclasswise_predpixels - totalclasswise_pixel_acc), axis=1) avg_pixel_acc_test = np.sum( totalclasswise_pixel_acc_test, axis=1) / np.sum( totalclasswise_gtpixels_test, axis=1) mean_class_acc_test = np.mean(totalclasswise_pixel_acc_test / totalclasswise_gtpixels_test, axis=1) mIoU_test = np.mean( totalclasswise_pixel_acc_test / (totalclasswise_gtpixels_test + totalclasswise_predpixels_test - totalclasswise_pixel_acc_test), axis=1) saved_accuracy = { "X": X, "P": avg_pixel_acc, "M": mean_class_acc, "I": mIoU, "P_test": avg_pixel_acc_test, "M_test": mean_class_acc_test, "I_test": mIoU_test } pickle.dump( saved_accuracy, open(os.path.join(ROOT_ADDRESS, "results_parts/saved_accuracy.p"), "wb")) # save the best model this_mIoU = np.mean( totalclasswise_pixel_acc_test / (totalclasswise_gtpixels_test + totalclasswise_predpixels_test - totalclasswise_pixel_acc_test), axis=1)[0] if this_mIoU > best_mIoU: if best_mIoU > 0: os.remove( os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_{}_best.pkl".format( args.arch, args.dataset, best_epoch, float2str(best_mIoU)))) os.remove( os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_{}_optimizer_best.pkl".format( args.arch, args.dataset, best_epoch, float2str(best_mIoU)))) best_mIoU = this_mIoU best_epoch = epoch + 1 torch.save( model, os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_{}_best.pkl".format( args.arch, args.dataset, best_epoch, float2str(best_mIoU)))) torch.save( { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, os.path.join( ROOT_ADDRESS, "results_parts/{}_{}_{}_{}_optimizer_best.pkl".format( args.arch, args.dataset, best_epoch, float2str(best_mIoU))))