def train(**kwargs):
cfg = Config()
for k, v in kwargs.items():
setattr(cfg, k, v)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
content = load_image(cfg.content, transform, max_size=cfg.max_size)
style = load_image(cfg.style,
transform,
shape=[content.size(3),
content.size(2)])
target = Variable(content.clone(), requires_grad=True)
optimizer = torch.optim.Adam([target], lr=cfg.lr, betas=[0.5, 0.999])
vgg = VGGNet()
if cfg.use_gpu:
vgg.cuda()
for step in range(cfg.total_step):
target_features = vgg(target)
content_features = vgg(Variable(content))
style_features = vgg(Variable(style))
style_loss = 0
content_loss = 0
for f1, f2, f3 in zip(target_features, content_features,
style_features):
# Compute content loss
content_loss += torch.mean((f1 - f2)**2)
_, c, h, w = f1.size()
f1 = f1.view(c, h * w)
f3 = f3.view(c, h * w)
# Compute gram matrix
f1 = torch.mm(f1, f1.t())
f3 = torch.mm(f3, f3.t())
style_loss += torch.mean((f1 - f3)**2) / (c * h * w)
loss = content_loss + cfg.style_weight * style_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step + 1) % cfg.log_step == 0:
print('Step [%d/%d], Content Loss: %.4f, Style Loss: %.4f' %
(step + 1, cfg.total_step, content_loss.data[0],
style_loss.data[0]))
if (step + 1) % cfg.sample_step == 0:
denorm = transforms.Normalize((-2.12, -2.04, -1.80),
(4.37, 4.46, 4.44))
img = target.clone().cpu().squeeze()
img = denorm(img.data).clamp_(0, 1)
torchvision.utils.save_image(img, 'output-%d.png' % (step + 1))
content_image = content_image.resize(size.astype(int), Image.ANTIALIAS)
content_image = transform(content_image).unsqueeze(0).cuda()
# Style Image processing
style_image = Image.open(style_image)
style_image = style_image.resize(
[content_image.size(2), content_image.size(3)], Image.LANCZOS)
style_image = transform(style_image).unsqueeze(0).cuda()
# Initialize result and optimizer
result_image = Variable(content_image.clone(), requires_grad=True)
optimizer = torch.optim.Adam([result_image], lr=0.003, betas=[0.5, 0.999])
# Model
vgg = VGGNet()
vgg = vgg.cuda()
# Train
for step in range(epochs):
target_features = vgg(result_image)
content_features = vgg(Variable(content_image))
style_features = vgg(Variable(style_image))
style_loss = 0
content_loss = 0
for f1, f2, f3 in zip(target_features, content_features, style_features):
# Content loss
content_loss += torch.mean((f1 - f2)**2)
train_data_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_data_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
if sys.argv[1] == 'vgg':
model = VGGNet()
elif sys.argv[1] == 'mobile':
model = MobileNet()
elif sys.argv[1] == 'custom':
model = CifarClassifier()
else:
raise ValueError(f'Unknown network type {sys.argv[1]}')
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(num_epochs):
for i, (x, x_class) in enumerate(train_data_loader):
# Forward pass
x = x.cuda() #.view(-1, img_size)
class_logits = model(x)
# Backprop and optimize
loss = loss_fn(class_logits, x_class.cuda())
# darc1 regularizer (optional)
darc1_loss = 0 #1e-3*torch.max(torch.sum(torch.abs(class_logits), dim=0))
loss = darc1_loss + loss
class Solver(object):
DEFAULTS = {}
def __init__(self, version, data_loader, config, output_txt):
"""
Initializes a Solver object
"""
# data loader
self.__dict__.update(Solver.DEFAULTS, **config)
self.version = version
self.data_loader = data_loader
self.output_txt = output_txt
self.build_model()
# start with a pre-trained model
if self.pretrained_model:
self.load_pretrained_model()
def build_model(self):
"""
Instantiates the model, loss criterion, and optimizer
"""
# instantiate model
self.model = VGGNet(self.config, self.use_batch_norm,
self.input_channels, self.class_count,
self.init_weights)
# instantiate loss criterion
self.criterion = nn.CrossEntropyLoss()
# instantiate optimizer
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
# print network
self.print_network(self.model, 'VGGNet')
# use gpu if enabled
if torch.cuda.is_available() and self.use_gpu:
self.model.cuda()
self.criterion.cuda()
def print_network(self, model, name):
"""
Prints the structure of the network and the total number of parameters
"""
num_params = 0
for p in model.parameters():
num_params += p.numel()
write_print(self.output_txt, name)
write_print(self.output_txt, str(model))
write_print(self.output_txt,
'The number of parameters: {}'.format(num_params))
def load_pretrained_model(self):
"""
loads a pre-trained model from a .pth file
"""
self.model.load_state_dict(
torch.load(
os.path.join(self.model_save_path,
'{}.pth'.format(self.pretrained_model))))
write_print(self.output_txt,
'loaded trained model {}'.format(self.pretrained_model))
def print_loss_log(self, start_time, iters_per_epoch, e, i, loss):
"""
Prints the loss and elapsed time for each epoch
"""
total_iter = self.num_epochs * iters_per_epoch
cur_iter = e * iters_per_epoch + i
elapsed = time.time() - start_time
total_time = (total_iter - cur_iter) * elapsed / (cur_iter + 1)
epoch_time = (iters_per_epoch - i) * elapsed / (cur_iter + 1)
epoch_time = str(datetime.timedelta(seconds=epoch_time))
total_time = str(datetime.timedelta(seconds=total_time))
elapsed = str(datetime.timedelta(seconds=elapsed))
log = "Elapsed {}/{} -- {}, Epoch [{}/{}], Iter [{}/{}], " \
"loss: {:.4f}".format(elapsed,
epoch_time,
total_time,
e + 1,
self.num_epochs,
i + 1,
iters_per_epoch,
loss)
write_print(self.output_txt, log)
def save_model(self, e):
"""
Saves a model per e epoch
"""
path = os.path.join(self.model_save_path,
'{}/{}.pth'.format(self.version, e + 1))
torch.save(self.model.state_dict(), path)
def model_step(self, images, labels):
"""
A step for each iteration
"""
# set model in training mode
self.model.train()
# empty the gradients of the model through the optimizer
self.optimizer.zero_grad()
# forward pass
output = self.model(images)
# compute loss
loss = self.criterion(output, labels.squeeze())
# compute gradients using back propagation
loss.backward()
# update parameters
self.optimizer.step()
# return loss
return loss
def train(self):
"""
Training process
"""
self.losses = []
self.top_1_acc = []
self.top_5_acc = []
iters_per_epoch = len(self.data_loader)
# start with a trained model if exists
if self.pretrained_model:
start = int(self.pretrained_model.split('/')[-1])
else:
start = 0
# start training
start_time = time.time()
for e in range(start, self.num_epochs):
for i, (images, labels) in enumerate(tqdm(self.data_loader)):
images = to_var(images, self.use_gpu)
labels = to_var(torch.LongTensor(labels), self.use_gpu)
loss = self.model_step(images, labels)
# print out loss log
if (e + 1) % self.loss_log_step == 0:
self.print_loss_log(start_time, iters_per_epoch, e, i, loss)
self.losses.append((e, loss))
# save model
if (e + 1) % self.model_save_step == 0:
self.save_model(e)
# evaluate on train dataset
# if (e + 1) % self.train_eval_step == 0:
# top_1_acc, top_5_acc = self.train_evaluate(e)
# self.top_1_acc.append((e, top_1_acc))
# self.top_5_acc.append((e, top_5_acc))
# print losses
write_print(self.output_txt, '\n--Losses--')
for e, loss in self.losses:
write_print(self.output_txt, str(e) + ' {:.4f}'.format(loss))
# print top_1_acc
write_print(self.output_txt, '\n--Top 1 accuracy--')
for e, acc in self.top_1_acc:
write_print(self.output_txt, str(e) + ' {:.4f}'.format(acc))
# print top_5_acc
write_print(self.output_txt, '\n--Top 5 accuracy--')
for e, acc in self.top_5_acc:
write_print(self.output_txt, str(e) + ' {:.4f}'.format(acc))
def eval(self, data_loader):
"""
Returns the count of top 1 and top 5 predictions
"""
# set the model to eval mode
self.model.eval()
top_1_correct = 0
top_5_correct = 0
total = 0
with torch.no_grad():
for images, labels in data_loader:
images = to_var(images, self.use_gpu)
labels = to_var(torch.LongTensor(labels), self.use_gpu)
output = self.model(images)
total += labels.size()[0]
# top 1
# get the max for each instance in the batch
_, top_1_output = torch.max(output.data, dim=1)
top_1_correct += torch.sum(
torch.eq(labels.squeeze(), top_1_output))
# top 5
_, top_5_output = torch.topk(output.data, k=5, dim=1)
for i, label in enumerate(labels):
if label in top_5_output[i]:
top_5_correct += 1
return top_1_correct.item(), top_5_correct, total
def train_evaluate(self, e):
"""
Evaluates the performance of the model using the train dataset
"""
top_1_correct, top_5_correct, total = self.eval(self.data_loader)
log = "Epoch [{}/{}]--top_1_acc: {:.4f}--top_5_acc: {:.4f}".format(
e + 1, self.num_epochs, top_1_correct / total,
top_5_correct / total)
write_print(self.output_txt, log)
return top_1_correct / total, top_5_correct / total
def test(self):
"""
Evaluates the performance of the model using the test dataset
"""
top_1_correct, top_5_correct, total = self.eval(self.data_loader)
log = "top_1_acc: {:.4f}--top_5_acc: {:.4f}".format(
top_1_correct / total, top_5_correct / total)
write_print(self.output_txt, log)
def train(cfg):
n_class = int(cfg["data"]["n_class"])
img_h = int(cfg["data"]["img_h"])
img_w = int(cfg["data"]["img_w"])
batch_size = int(cfg["training"]["batch_size"])
epochs = int(cfg["training"]["epochs"])
lr = float(cfg["training"]["optimizer"]["lr"])
momentum = float(cfg["training"]["optimizer"]["momentum"])
w_decay = float(cfg["training"]["optimizer"]["weight_decay"])
step_size = int(cfg["training"]["lr_schedule"]["step_size"])
gamma = float(cfg["training"]["lr_schedule"]["gamma"])
configs = "FCNs-BCEWithLogits_batch{}_epoch{}_RMSprop_scheduler-step{}-gamma{}_lr{}_momentum{}_w_decay{}_input_size{}_03091842".format(
batch_size, epochs, step_size, gamma, lr, momentum, w_decay, img_h)
print("Configs:", configs)
root_dir = cfg["data"]["root_dir"]
train_filename = cfg["data"]["train_file"]
val_filename = cfg["data"]["val_file"]
mean_filename = cfg["data"]["mean_file"]
class_weight_filename = cfg["data"]["class_weight_file"]
train_file = os.path.join(root_dir, train_filename)
print(train_file)
val_file = os.path.join(root_dir, val_filename)
mean_file = os.path.join(root_dir, mean_filename)
class_weight_file = os.path.join(root_dir, class_weight_filename)
model_dir = cfg["training"]["model_dir"]
if not os.path.exists(model_dir):
os.makedirs(model_dir)
model_path = os.path.join(model_dir, configs)
use_gpu = torch.cuda.is_available()
num_gpu = list(range(torch.cuda.device_count()))
continue_train = False
#MeanRGB_train = ComputeMeanofInput(train_file)
#MeanRGB_train = np.load(mean_file)
MeanRGB_train = np.array([0.0, 0.0, 0.0])
print("MeanRGB_train: {}".format(MeanRGB_train))
train_data = ScanNet2d(csv_file=train_file,
phase='train',
trainsize=(img_h, img_w),
MeanRGB=MeanRGB_train)
val_data = ScanNet2d(csv_file=val_file,
phase='val',
trainsize=(img_h, img_w),
MeanRGB=MeanRGB_train)
train_loader = DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
num_workers=1)
val_loader = DataLoader(val_data,
batch_size=batch_size,
shuffle=False,
num_workers=1)
#class_weight = trainer.computer_class_weights(train_file)
class_weight = np.load(class_weight_file)
print("class_weight: {}".format(class_weight))
class_weight = torch.from_numpy(class_weight)
print("shape of class weight {}".format(class_weight.shape))
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
fcn_model = FCN8s(encoder_net=vgg_model, n_class=n_class)
if use_gpu:
ts = time.time()
vgg_model = vgg_model.cuda()
fcn_model = fcn_model.cuda()
fcn_model = nn.DataParallel(fcn_model, device_ids=num_gpu)
class_weight = class_weight.cuda()
print("Finish cuda loading, tme elapsed {}".format(time.time() - ts))
L = nn.BCEWithLogitsLoss(reduction='none')
optimizer = optim.RMSprop(fcn_model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=w_decay)
#optimizer = optim.SGD(fcn_model.parameters(), lr=lr, momentum= momentum, weight_decay=w_decay)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=step_size,
gamma=gamma)
score_dir = os.path.join("scores", configs)
if not os.path.exists(score_dir):
os.makedirs(score_dir)
log_headers = [
'epoch', 'train/loss', 'train/acc', 'train/acc_cls', 'train/mean_iu',
'train/fwavacc', 'val/loss', 'val/acc', 'val/acc_cls', 'val/mean_iu',
'val/fwavacc', 'elapsed_time'
]
if not os.path.exists(os.path.join(score_dir, 'log.csv')):
with open(os.path.join(score_dir, 'log.csv'), 'w') as f:
f.write(','.join(log_headers) + '\n')
IU_scores = np.zeros((epochs, n_class + 1))
pixel_scores = np.zeros(epochs)
writer = SummaryWriter()
# color_mapping = util.GenerateColorMapping(n_class)
best_mean_iu = 0
epoch_loss = 0.0
if continue_train:
model_path = "C:\\Users\\ji\\Documents\\FCN-VGG16\\models\\FCNs-BCEWithLogits_batch1_epoch500_RMSprop_scheduler-step50-gamma0.5_lr0.0001_momentum0.0_w_decay1e-05"
fcn_model = torch.load(model_path)
fcn_model.train()
for epoch in range(epochs):
fcn_model.train()
scheduler.step()
ts = time.time()
running_loss = 0.0
######
label_preds = []
label_trues = []
######
for i, batch in enumerate(train_loader):
optimizer.zero_grad()
if use_gpu:
inputs = Variable(batch['X'].cuda())
labels = Variable(batch['Y'].cuda())
else:
inputs, labels = Variable(batch['X']), Variable(batch['Y'])
outputs = fcn_model(inputs)
#print("out: {}".format(outputs.shape))
#print("label: {}".format(labels.shape))
#print(outputs)
#print(labels)
loss = L(outputs, labels)
# print(loss.shape)
loss = loss.permute(0, 2, 3,
1).reshape(-1,
n_class + 1) #.view(-1,n_class+1)
# print(loss.shape)
loss = torch.mean(loss, dim=0)
# print(loss.shape)
loss = torch.mul(loss, class_weight)
# print(loss.shape)
loss = torch.mean(loss)
# print(loss)
loss.backward()
# print("grad")
# print(fcn_model.outp.weight.grad)
# print(fcn_model.embs[0].weight.grad)
optimizer.step()
#scheduler.step()
if i == 0 and epoch == 0:
# count= util.count_parameters(fcn_model)
# print("number of parameters in model {}".format(count))
visIn = inputs[:3]
#print('shape of in {}'.format(visIn[:5].shape))
visLabel = batch['l'][:3]
epoch_loss += loss.item()
running_loss += loss.item()
# print("loss: {}".format(loss.data))
if i % 10 == 9 and i != 0:
print("epoch{}, iter{}, Iterloss: {}".format(
epoch, i, running_loss / 10))
writer.add_scalar('train/iter_loss', running_loss / 10,
epoch * len(train_loader) + i)
running_loss = 0.0
# N, _, h, w = outputs.shape
# targets = batch['l'].cpu().numpy().reshape(N,h,w)
# outputs = outputs.data.cpu().numpy()
# preds_v, targets_v = util.visulaize_output(outputs,targets,color_mapping,n_class)
# writer.add_images('train/predictions',torch.from_numpy(preds_v),dataformats='NHWC')
# writer.add_images('train/targets',torch.from_numpy(targets_v),dataformats='NHWC')
#####################################
outputs = outputs.data.cpu().numpy()
N, _, h, w = outputs.shape
pred = outputs.transpose(0, 2, 3, 1).reshape(
-1, n_class + 1).argmax(axis=1).reshape(N, h, w)
target = batch['l'].cpu().numpy().reshape(N, h, w)
#########
for lt, lp in zip(target, pred):
label_trues.append(lt)
label_preds.append(lp)
metrics = util.label_accuracy_score(label_trues, label_preds,
n_class + 1)
with open(os.path.join(score_dir, "log.csv"), 'a') as f:
log = [epoch] + [epoch_loss] + list(metrics) + [''] * 7
log = map(str, log)
f.write(','.join(log) + '\n')
########################################
#scheduler.step()
writer.add_scalar('train/epoch_loss', epoch_loss, epoch)
print("Finish epoch{}, epoch loss {}, time eplapsed {}".format(
epoch, epoch_loss,
time.time() - ts))
epoch_loss = 0.0
####################
writer.add_scalar('train/mean_iu', metrics[2], epoch)
writer.add_scalar('train/acc', metrics[0], epoch)
writer.add_scalar('train/acc_cls', metrics[1], epoch)
######################
#Training precess visulize
visOut = fcn_model(visIn)
preds_v, targets_v = util.visulaize_output(visOut, visLabel, n_class)
writer.add_images('train/predictions',
torch.from_numpy(preds_v),
global_step=epoch,
dataformats='NHWC')
writer.add_images('train/targets',
torch.from_numpy(targets_v),
global_step=epoch,
dataformats='NHWC')
if not os.path.exists(model_path):
os.makedirs(model_path)
torch.save(fcn_model, os.path.join(model_path, str(epoch)))
best_mean_iu = val_model(epoch, val_loader, fcn_model, use_gpu,
n_class, IU_scores, pixel_scores, score_dir,
writer, best_mean_iu, model_path, L)
writer.flush()
writer.close()
