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))
