def main(args): torch.manual_seed(1) # crop input image and ground truth and save on disk cropped_input_images_path = os.path.join(args.save_cropped, 'input_images') cropped_gt_images_path = os.path.join(args.save_cropped, 'gt_images') if args.crop_images: crop_and_save(args, cropped_input_images_path, cropped_gt_images_path) seg_dataset = SegmentationData(cropped_input_images_path, cropped_gt_images_path, args.n_classes, args.phase) train_loader = DataLoader(seg_dataset, shuffle=True, num_workers=4, batch_size=args.batch_size) model = FCN(args.n_classes) use_gpu = torch.cuda.is_available() num_gpu = list(range(torch.cuda.device_count())) if use_gpu : model = model.cuda() model = nn.DataParallel(model, device_ids=num_gpu) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = nn.BCEWithLogitsLoss() losses = [] for epoch in range(args.n_epoch): for i, (image, segement_im) in enumerate(train_loader): image = image.float() images = Variable(image.cuda()) labels = Variable(segement_im.cuda()) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels.float()) loss.backward() optimizer.step() # add loss to a list for plotting it later if i == 0: losses.append(loss) print("epoch{} iteration {} loss: {}".format(epoch, i, loss.data.item())) if epoch%5 == 0: pred = outputs.data.max(1)[1].cpu().numpy()[0] decoded = decode_segmap(pred) decoded = Image.fromarray(decoded) path = os.path.join(args.output_path, 'output_%d_%d.png' % (epoch, i)) decoded.save(path) # plot loss plot(losses, args) # save model if not os.path.exists(args.model_path): os.makedirs(args.model_path) model_name = os.path.join(args.model_path, 'fcn.pt') torch.save(model, model_name)
def train(name, *args): kwargs = {} for arg in args: key, val = arg.split('=') kwargs[key] = val OVERRODE[key] = val with get_session() as sess: fcn = FCN(sess=sess, name=name, kwargs=kwargs) sess.run(tf.global_variables_initializer()) fcn.train(EPOCHS)
def test(name, ckpt, image_pack_name=None, output_filename=None): external_image = image_pack_name.index('.') != -1 if image_pack_name is None: data = None elif not external_image: data = load_data(image_pack_name.split(',')) with get_session() as sess: fcn = FCN(sess=sess, name=name) if ckpt != "-1": fcn.load(ckpt) else: fcn.load_absolute(name) if not external_image: errors, _, _, _, ret, conf = fcn.test(scales=[0.5], summary=True, summary_key=123, data=data, eval_speed=False, visualize=True) if output_filename is not None: with open('outputs/%s.pkl' % output_filename, 'wb') as f: pickle.dump(ret, f) with open('outputs/%s_err.pkl' % output_filename, 'wb') as f: pickle.dump(errors, f) with open('outputs/%s_conf.pkl' % output_filename, 'wb') as f: pickle.dump(conf, f) print ret print 'results dumped' else: img = cv2.imread(image_pack_name) # reverse gamma correction for sRGB img = (img / 255.0)**2.2 * 16384 images = [img] fcn.test_external(images=images, fns=[image_pack_name])
def make_model(): model = FCN() if torch.cuda.is_available(): model.to('cuda') model = nn.DataParallel(model) return model
def testFCN(): from gpu_util import log_mem from fcn import FCN, device bs = 1 image = torch.randn([bs, 3, 512, 512]).to(device) rays_o = torch.randn([bs, 3, 512, 512]).to(device) rays_d = torch.randn([bs, 3, 512, 512]).to(device) model = FCN() model.to(device) out = model(image, rays_o, rays_d)
def dump_errors(name, ckpt, fold, output_filename, method='full', samples=0, pooling='median'): samples = int(samples) with get_session() as sess: kwargs = {'dataset_name': 'gehler', 'subset': 0, 'fold': fold} fcn = FCN(sess=sess, name=name, kwargs=kwargs) fcn.load(ckpt) for i in range(4): if method == 'full': errors, t, _, _, _ = fcn.test(scales=[0.5]) elif method == 'resize': errors, t = fcn.test_resize() elif method == 'patches': errors, t = fcn.test_patch_based(scale=0.5, patches=samples, pooling=pooling) else: assert False utils.print_angular_errors(errors) with open(output_filename, 'w') as f: pickle.dump({'e': errors, 't': t}, f)
def train(): ROOT = './' VGG16_WEIGHT_PATH = './vgg/vgg16_weights.npz' DATASET_PATH = os.path.join(ROOT, 'VOC2012/') CHECKPOINT_DIR = os.path.join(DATASET_PATH, 'saved_model') IMAGE_SHAPE = (512, 512) N_CLASSES = 21 N_EPOCHS = 100 BATCH_SIZE = 1 LEARNING_RATE = 1e-5 DECAY_RATE = 0.95 DECAY_EPOCH = 10 DROPOUT_RATE = 0.5 print('Starting end-to-end training FCN-8s') session_config = tf.ConfigProto(gpu_options=tf.GPUOptions( allow_growth=True)) session = tf.compat.v1.InteractiveSession(config=session_config) session.as_default() # ------------- Load VOC from TFRecord --------------- dataset = VOCDataset() dataset_train = dataset.load_dataset(DATASET_PATH, BATCH_SIZE, is_training=True) dataset_val = dataset.load_dataset(DATASET_PATH, BATCH_SIZE, is_training=False) # ------------- Build fcn model ------------ fcn = FCN(IMAGE_SHAPE, N_CLASSES, VGG16_WEIGHT_PATH) fcn.build_from_vgg() learning_rate_fn = learning_rate_with_exp_decay(BATCH_SIZE, dataset.n_images['train'], DECAY_EPOCH, DECAY_RATE, LEARNING_RATE) compile_model(fcn, learning_rate_fn) fit_model(fcn, N_EPOCHS, BATCH_SIZE, dataset_train, dataset_val, CHECKPOINT_DIR, DROPOUT_RATE)
def dump_result(name, ckpt, image_pack_name=None): if image_pack_name is None: data = None else: data = load_data(image_pack_name.split(',')) outputs = [] gts = [] with get_session() as sess: fcn = FCN(sess=sess, name=name) fcn.load(ckpt) _, _, outputs, gts = fcn.test(scales=[0.5], summary=True, summary_key=123, data=data) result = { 'outputs': np.array(outputs), 'gts': np.array(gts), } pickle.dump( result, open("outputs/%s-%s-%s.pkl" % (name, ckpt, image_pack_name), "wb"))
def cont(name, preload, key): with get_session() as sess: fcn = FCN(sess=sess, name=name) sess.run(tf.global_variables_initializer()) if preload is not None: fcn.load(name=preload, key=key) fcn.train(EPOCHS)
def test_input_gamma(name, ckpt, input_gamma, image_pack_name=None, output_filename=None): config_set_input_gamma(float(input_gamma)) if image_pack_name is None: data = None else: data = load_data(image_pack_name.split(',')) with get_session() as sess: fcn = FCN(sess=sess, name=name) fcn.load(ckpt) _, _, _, _, ret = fcn.test(scales=[0.5], summary=True, summary_key=123, data=data) if output_filename is not None: with open('outputs/%s.pkl' % output_filename, 'wb') as f: pickle.dump(ret, f) print ret print 'results dumped'
def f(input_shape, **kwargs): if K.image_data_format() == 'channels_first': input_shape = (input_shape[2], input_shape[0], input_shape[1]) model = FCN(input_shape=input_shape, **kwargs) return model
def vis_fcn_vgg16(): input_shape = (500, 500, 3) fcn_vgg16 = FCN(input_shape=input_shape) plot_model(fcn_vgg16, to_file='fcn_vgg16.png')
def test_naive(): return FCN.test_naive()
from camera import CameraHandler #from camera.camera_handler_2 import CameraHandler from policy import Detector import numpy as np from fcn import FCN import cv2 '''FCN''' VGG_WEIGHTS = '/home/nvidia/RobotControl/src/fcn/vgg16.npy' MODEL_PATH = '/home/nvidia/RobotControl/src/fcn/model_save/old_2/' model = FCN((240, 320), 1, VGG_WEIGHTS) model.load(path=MODEL_PATH) img = None ct = 0 mask = np.zeros((480, 640), dtype=np.uint8) mask = cv2.fillConvexPoly( mask, np.array([[275, 285], [365, 285], [430, 166], [210, 166]]), 1, 1) def preprocessor(img): # mask=np.zeros((480, 640), dtype=np.uint8) # mask = cv2.fillConvexPoly(mask, np.array([[275, 285], [365, 285], [430, 166], [210, 166]]), 1, 1) return model.extract_hmap(img)[0] * mask # return img * np.expand_dims(mask,-1) handler = CameraHandler(2, frame_rate=20, preprocessor=lambda img: preprocessor(img)) handler.start()
import matplotlib.pyplot as plt logging.basicConfig(level=logging.DEBUG) def print_usage(): print('Usage:') print('python train.py MAX_ITERATIONS_COARSE MAX_ITERATIONS_FINE ' 'SAVE_PARAMS_AFTER RESTORE_COARSE_PARAMS_PATH') if len(sys.argv) not in [4, 5]: print_usage() exit(1) TRAINVAL_ROOT_DIR = '/home/paperspace/PASCAL-VOC-Dataset/TrainVal' TEST_ROOT_DIR = '/home/paperpsace/PASCAL-VOC-Dataset/Test' VGG_PARAMS_ROOT_DIR = '/home/paperspace/FCN/vgg-weights' MAX_ITERATIONS_COARSE = int(sys.argv[1]) MAX_ITERATIONS_FINE = int(sys.argv[2]) SAVE_PARAMS_AFTER = int(sys.argv[3]) if len(sys.argv) == 5: RESTORE_CKPT = sys.argv[4] else: RESTORE_CKPT = None fcn = FCN(TRAINVAL_ROOT_DIR, TEST_ROOT_DIR, VGG_PARAMS_ROOT_DIR) fcn.train(MAX_ITERATIONS_COARSE, MAX_ITERATIONS_FINE, SAVE_PARAMS_AFTER, RESTORE_CKPT)
from fcn import FCN # set PATH to location of Pacal VOC dataset PATH = '/home/seth/Datasets/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/' fcn = FCN(path=PATH) #fcn.load('my_model') fcn.train(epochs=75) fcn.save('my_model') fcn.evaluate(val=False) fcn.evaluate() for id in fcn.train_list[:10]: fcn.data.show_seg(id) fcn.predict(id) for id in fcn.val_list[:10]: fcn.data.show_seg(id) fcn.predict(id)
if __name__ == '__main__': target_dir = sys.argv[1] if len(sys.argv) > 2: start_ind = int(sys.argv[2]) else: start_ind = 0 image_dir = osp.join(target_dir, "*") print(image_dir) imgpaths = glob(image_dir)[start_ind:] print(len(imgpaths)) save_dir = osp.join(osp.dirname(target_dir), osp.basename(target_dir) + "_colorcorrected") print(save_dir) ckpt = "pretrained/colorchecker_fold1and2.ckpt" with get_session() as sess: fcn = FCN(sess=sess, name=ckpt) fcn.load_absolute(ckpt) errors = test_external(fcn, imgpaths=imgpaths, save_prefix=save_dir, show=False, write_compare=False) with open(osp.join(target_dir, "errored_images.txt"), "w") as f: for p in errors: f.write(p)
from arguments import Arguments from fcn import FCN from svm import SVM import torch INIT_FCN = True INIT_SVM = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args = Arguments() for dataset in args.input_sizes: if INIT_FCN: model = FCN(args.input_sizes[dataset], args.output_sizes[dataset]).to(device) print(model.input_size, model.output_size) init_path = '../init/{}_fcn.init'.format(dataset) torch.save(model.state_dict(), init_path) print('Save init: {}'.format(init_path)) if INIT_SVM: model = SVM(args.input_sizes[dataset], args.output_sizes[dataset]).to(device) print(model.n_feature, model.n_class) init_path = '../init/{}_svm.init'.format(dataset) torch.save(model.state_dict(), init_path) print('Save init: {}'.format(init_path))
import tensorflow as tf print('Tensorflow version: {}'.format(tf.__version__)) from dataset import VOCDataset from fcn import FCN # Killing optional CPU driver warnings os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' tf.logging.set_verbosity(tf.logging.ERROR) if __name__ == '__main__': image_shape = (512, 512) n_classes = 21 vgg16_weights_path = './vgg/vgg16_weights.npz' model = FCN(image_shape, n_classes, vgg16_weights_path) model.build_from_vgg() root_path = './' dataset_path = os.path.join(root_path, 'VOC2012/') dataset = VOCDataset(augmentation_params=None) dataset_val = dataset.load_dataset(dataset_path, batch_size=8, is_training=False) iterator = tf.data.Iterator.from_structure(dataset_val.output_types, dataset_val.output_shapes) next_batch = iterator.get_next() val_init_op = iterator.make_initializer(dataset_val)
'_stratify_True_uniform_True_repeat_{}.pkl'.format( dataset, args.num_workers, args.repeat) else: data_file = '../ckpts/data_{}_non_iid_{}_num_workers_{}' \ '_stratify_True_uniform_True_repeat_{}.pkl'.format( dataset, non_iid, args.num_workers, args.repeat) print('Loading data: {}'.format(data_file)) X_trains, y_trains = pkl.load( open(data_file, 'rb')) print(fog_graph) best = 0 # Fire the engines model = FCN(args.input_size, args.output_size).to(device) model.load_state_dict(torch.load(init_path)) print('Load init: {}'.format(init_path)) best = 0 x_ax = [] y_ax = [] l_test = [] grad_tr = [] for epoch in range(1, args.epochs + 1): grad = train(args, model, fog_graph, workers, X_trains, y_trains, device, epoch, loss_fn='hinge') acc, loss = test(args, model, device, test_loader, best, epoch, loss_fn='hinge') y_ax.append(acc) x_ax.append(epoch)
padding0 = 32 - w % 32 padding1 = 32 - h % 32 image = cv2.copyMakeBorder(image, padding0 // 2, padding0 - padding0 // 2, padding1 // 2, padding1 - padding1 // 2, cv2.BORDER_REFLECT) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) image = transform(image).unsqueeze(0) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) state_dict = torch.load(path, map_location='cpu') model = FCN(21).to(device).eval() model.load_state_dict(state_dict) del state_dict with torch.no_grad(): image = image.to(device) output = model(image) predict = output.data.cpu().squeeze(0).numpy() predict = np.argmax(predict, axis=0) predict = colormap[predict] cv2.imwrite("res.png", predict)
])), batch_size=args.test_batch_size, shuffle=True, **kwargs) X_trains, y_trains = pkl.load( open( '../ckpts/data_non_iid_{}_num_workers_{}_stratify_True_uniform_True.pkl' .format(non_iid, args.num_workers), 'rb')) print(fog_graph) best = 0 # Fire the engines model = FCN().to(device) model.load_state_dict(torch.load(init_path)) print('Load init: {}'.format(init_path)) best = 0 x_ax = [] y_ax = [] l_test = [] for epoch in range(1, args.epochs + 1): train(args, model, fog_graph, workers, X_trains, y_trains,
#制定训练的硬件场景咯 device = t.device('cuda') if t.cuda.is_available() else t.device('cpu') num_class = cfg.DATASET[1] train_data = DataLoader(Cam_train, batch_size=cfg.BATCH_SIZE, shuffle=True, num_workers=1) val_data = DataLoader(Cam_val, batch_size=cfg.BATCH_SIZE, shuffle=True, num_workers=1) #FCN模型类实例化 fcn = FCN.FCN(num_class) fcn = fcn.to(device) #Loss function 因为之后有可能会尝试别的激活函数,所以不用封装好的CEL #CrossEntropyLoss() 函数联合调用了 nn.LogSoftmax() 和 nn.NLLLoss() More on : #https://blog.csdn.net/jasonleesjtu/article/details/89141554?utm_medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase&depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase criterion = nn.NLLLoss().to(device) #optimizer 经验:一般2D数据集Adam效果比较好 如果是3D数据集SGD会好些 optimizer = optim.Adam(fcn.parameters(), lr=1e-4) def train(model): best = [0] #初始化 保存最好的值 net = model.train() #torch 训练模式 # 训练轮次 for epoch in range(cfg.EPOCH_NUMBER):
x_data, y_data = load_dataset('classification_data.csv', device) n_folds = 5 kf = KFold(n_splits=n_folds, shuffle=True, random_state=2) scores = [] idx = 1 for train_index, test_index in tqdm(kf.split(x_data)): # init model nets = [] if resnet_model is not None: net = ResNet() cp_name = resnet_model + str(idx) load_path = os.path.join(save_dir, cp_name) net.load_state_dict(torch.load(load_path)['state_dict']) nets.append(net) if fcn_model is not None: net = FCN() cp_name = fcn_model + str(idx) load_path = os.path.join(save_dir, cp_name) net.load_state_dict(torch.load(load_path)['state_dict']) nets.append(net) if encoder_model is not None: net = Encoder() cp_name = encoder_model + str(idx) load_path = os.path.join(save_dir, cp_name) net.load_state_dict(torch.load(load_path)['state_dict']) nets.append(net) x_train_data = [x_data[idx] for idx in train_index] y_train_data = [y_data[idx] for idx in train_index] x_test_data = [x_data[idx] for idx in test_index] y_test_data = [y_data[idx] for idx in test_index]
if args.batch_size == 0: args.batch_size = args.num_train print("Resetting batch size: {}...".format(args.batch_size)) train_loader = get_trainloader(args.dataset, args.batch_size, False) test_loader = get_testloader(args.dataset, args.test_batch_size, noise=args.noise) print('+' * 80) # Fire the engines # Fire the engines if args.clf == 'fcn': print('Initializing FCN...') model = FCN(args.input_size, args.output_size) elif args.clf == 'svm': print('Initializing SVM...') model = SVM(args.input_size, args.output_size) elif args.clf == 'resnet18': print('Initializing ResNet18...') model = resnet.resnet18(num_channels=args.num_channels, num_classes=args.output_size) model.load_state_dict(torch.load(init_path)) print('Load init: {}'.format(init_path)) model = nn.DataParallel(model.to(device), device_ids=args.device_id) if 'sgd' in args.paradigm: optim = optim.SGD(params=model.parameters(), lr=args.lr) elif 'adam' in args.paradigm:
def test_multi(name, ckpt): with get_session() as sess: fcn = FCN(sess=sess, name=name) fcn.load(ckpt) fcn.test_multi()
def test_external(self, imgpaths, scale=1.0, show=True, save_prefix="cc_outputs", write=False, write_compare=False): illums = [] confidence_maps = [] errors = [] for i, imgp in enumerate(imgpaths): filename = osp.basename(imgp) + ".jpg" # force save as jpg print "" print i, filename try: img = cv2.imread(imgp) except Exception as e: print("ERROR can't read image {}\n{}".format(filename, e)) errors.append(imgp) continue if img is None: print("ERROR read None for {}".format(filename)) errors.append(imgp) continue # reverse gamma correction for sRGB img_linear = (img / 255.0)**2.2 * 65536 if scale != 1.0: img_linear = cv2.resize(img_linear, (0, 0), fx=scale, fy=scale) shape = img_linear.shape[:2] if shape not in self.test_nets: #print("shape {}".format(shape)) aspect_ratio = 1.0 * shape[1] / shape[0] if aspect_ratio < 1: target_shape = (MERGED_IMAGE_SIZE, MERGED_IMAGE_SIZE * aspect_ratio) else: target_shape = (MERGED_IMAGE_SIZE / aspect_ratio, MERGED_IMAGE_SIZE) target_shape = tuple(map(int, target_shape)) test_net = {} test_net['illums'] = tf.placeholder(tf.float32, shape=(None, 3), name='test_illums') test_net['images'] = tf.placeholder(tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): try: test_net['pixels'] = FCN.build_branches( test_net['images'], 1.0) except ValueError as e: print("ERROR {}".format(e)) errors.append(imgp) continue test_net['est'] = tf.reduce_sum(test_net['pixels'], axis=(1, 2)) test_net['merged'] = get_visualization(test_net['images'], test_net['pixels'], test_net['est'], test_net['illums'], target_shape) self.test_nets[shape] = test_net test_net = self.test_nets[shape] pixels, est, merged = self.sess.run( [test_net['pixels'], test_net['est'], test_net['merged']], feed_dict={ test_net['images']: img_linear[None, :, :, :], test_net['illums']: [[1, 1, 1]] }) est = est[0] est /= np.linalg.norm(est) pixels = pixels[0] confidences = np.linalg.norm(pixels, axis=2) #confidence_maps.append(confidences) ind = int(confidences.flatten().shape[0] * 0.95) print("Confidence: {:0.1f} mean, {:0.1f} max, {:0.1f} 95%".format( confidences.mean(), confidences.max(), sorted(confidences.flatten())[ind])) merged = merged[0] #illums.append(est) if show: cv2.imshow('Ret', merged[:, :, ::-1]) k = cv2.waitKey(0) % (2**20) try: os.makedirs(save_prefix) except: pass corrected = np.power( img_linear[:, :, ::-1] / 65535 / est[None, None, :] * np.mean(est), 1 / 2.2)[:, :, ::-1] if show: cv2.imshow("corrected", corrected) cv2.waitKey(0) corrected = corrected * 255.0 output_img = corrected if write_compare: orig_img = img #np.power(img / 65535, 1/2.2) * 255 output_img = np.concatenate((orig_img, corrected), axis=1) cv2.imwrite( osp.join(save_prefix, 'corrected_%s' % filename), output_img, ) return errors
def test_network(name, ckpt): with get_session() as sess: fcn = FCN(sess=sess, name=name) fcn.load(ckpt) fcn.test_network()
factor=np.sqrt(0.1), cooldown=0, patience=5, min_lr=0.5e-6) early_stopper = EarlyStopping(monitor='val_loss', min_delta=0.001, patience=100) csv_logger = CSVLogger('output/{}_fcn_vgg16.csv'.format( datetime.datetime.now().isoformat())) datagen = PascalVocGenerator(**init_args['pascal_voc_generator']['train']) train_loader = ImageSetLoader(**init_args['image_set_loader']['train']) val_loader = ImageSetLoader(**init_args['image_set_loader']['val']) fcn_vgg16 = FCN(basenet='vgg16', input_shape=(500, 500, 3), num_output=21) fcn_vgg16.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy', 'categorical_accuracy']) flow_args = init_args['pascal_voc_generator']['flow_from_imageset'] train_flow_args = flow_args.copy() train_flow_args['image_set_loader'] = train_loader val_flow_args = flow_args.copy() val_flow_args['image_set_loader'] = val_loader fcn_vgg16.fit_generator( datagen.flow_from_imageset(**train_flow_args), steps_per_epoch=1112, epochs=100, validation_data=datagen.flow_from_imageset(**val_flow_args),
def main(): root = "./data/VOCdevkit/VOC2012" batch_size = 4 num_workers = 4 num_classes = 21 lr = 0.0025 # lr = 5e-4 # fine-tune epoches = 100 writer = SummaryWriter(comment="-fcn") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_transform = A.Compose([ A.HorizontalFlip(), # 注意这个先后顺序 A.VerticalFlip(), # A.transpose(p=0.5), A.RandomRotate90(), # A.ElasticTransform(p=1, alpha=120, # sigma=120 * 0.05, # alpha_affine=120 * 0.03), A.RandomResizedCrop(320, 480), ]) val_transform = A.Compose([ A.RandomResizedCrop(320, 480)]) train_set = VOCdataset(root, mode="train", transform=train_transform) val_set = VOCdataset(root, mode="val", transform=val_transform) train_loader = data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers) val_loader = data.DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=num_workers) model = FCN(num_classes).to(device) # state_dict = torch.load("./model/best.pth") # print("loading pretrained parameters") # model.load_state_dict(state_dict) # del state_dict criteria = nn.CrossEntropyLoss() # optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=2e-4) # optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, # weight_decay=2e-4) vgg_parameters = (list(map(id, model.encode1.parameters()))+ list(map(id, model.encode2.parameters()))+ list(map(id, model.encode3.parameters()))+ list(map(id, model.encode4.parameters()))+ list(map(id, model.encode5.parameters()))) encode_parameters = (list(model.encode1.parameters())+ list(model.encode2.parameters())+ list(model.encode3.parameters())+ list(model.encode4.parameters())+ list(model.encode5.parameters())) decode_parameters = filter(lambda p: id(p) not in vgg_parameters, model.parameters()) optimizer = optim.SGD([{'params': encode_parameters, 'lr': 0.1 * lr}, {'params': decode_parameters, 'lr': lr}], momentum=0.9, weight_decay=2e-3) # optimizer = optim.Adam([{'params': encode_parameters, 'lr': 0.1 * lr}, # {'params': decode_parameters, 'lr': lr}], # weight_decay=2e-4) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.85) # scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, # T_0=100, # T_mult=1, # eta_min=0.0001) # scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10) best_miou = 0.0 for epoch in range(1, epoches+1): print("Epoch: ", epoch) scheduler.step() train_info = train(train_loader, model, criteria, optimizer, device, batch_size) val_info = validate(val_loader, model, criteria, device, batch_size) string = "loss: {}, pixel acc: {}, mean acc: {} miou: {}" print("train", end=' ') print(string.format(train_info['loss'], train_info["pixel acc"], train_info['mean acc'], train_info['miou'])) print("val", end=' ') print(string.format(val_info['loss'], val_info['pixel acc'], val_info['mean acc'], val_info['miou'])) writer.add_scalar("lr", optimizer.state_dict()['param_groups'][0]['lr'], epoch) writer.add_scalar('train/loss', train_info['loss'], epoch) writer.add_scalar('train/pixel acc', train_info['pixel acc'], epoch) writer.add_scalar('train/mean acc', train_info['mean acc'], epoch) writer.add_scalar('train/miou', train_info['miou'], epoch) writer.add_scalar('val/loss', val_info['loss'], epoch) writer.add_scalar('val/pixel acc', val_info['pixel acc'], epoch) writer.add_scalar('val/mean acc', val_info['mean acc'], epoch) writer.add_scalar('val/miou', val_info['miou'], epoch) if val_info['miou'] > best_miou: best_miou = val_info['miou'] torch.save(model.state_dict(), './model/best.pth') print("best model find at {} epoch".format(epoch))