def eval(args):
frames = args.frames
caption_length = args.caption_length
glove_file = args.glove_file
if args.cuda:
ctx = mx.gpu()
else:
ctx = mx.cpu()
if args.load_pretrain:
pretrain_model = vision.vgg16_bn(pretrained=True,ctx=ctx)
transform = utils.Compose([utils.ToTensor(ctx),
utils.normalize(ctx),
utils.extractFeature(ctx,pretrain_model)
])
else:
pretrain_model = None
transform = utils.Compose([utils.ToTensor(ctx),
utils.normalize(ctx),
])
target_transform = utils.targetCompose([utils.WordToTensor(ctx)])
val_dataset = videoFolder(args.val_folder,args.val_dict, frames, glove_file, caption_length, ctx, transform=transform, target_transform=target_transform)
val_loader = gluon.data.DataLoader(val_dataset, batch_size=args.batch_size,last_batch='discard',shuffle=True)
def get_coco(root, image_set, transforms, mode='instances'):
anno_file_template = "{}_{}2017.json"
PATHS = {
"train": ("train2017",
os.path.join("annotations",
anno_file_template.format(mode, "train"))),
"val": ("val2017",
os.path.join("annotations",
anno_file_template.format(mode, "val"))),
# "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
}
t = [ConvertCocoPolysToMask()]
if transforms is not None:
t.append(transforms)
transforms = T.Compose(t)
img_folder, ann_file = PATHS[image_set]
img_folder = os.path.join(root, img_folder)
ann_file = os.path.join(root, ann_file)
dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
if image_set == "train":
dataset = _coco_remove_images_without_annotations(dataset)
# dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
return dataset
def test(args):
if args.gpu:
ctx = [mx.gpu(0)]
else:
ctx = [mx.cpu(0)]
if args.dataset == "Sony":
out_channels = 12
scale = 2
else:
out_channels = 27
scale = 3
# load data
test_transform = utils.Compose([utils.ToTensor()])
test_dataset = data.MyDataset(args.dataset,
"test",
transform=test_transform)
test_loader = gluon.data.DataLoader(test_dataset,
batch_size=1,
last_batch='discard')
unet = net.UNet(out_channels, scale)
unet.load_params(args.model, ctx=ctx)
batches = 0
avg_psnr = 0.
for img, gt in test_loader:
batches += 1
imgs = gluon.utils.split_and_load(img[0], ctx)
label = gluon.utils.split_and_load(gt[0], ctx)
outputs = []
for x in imgs:
outputs.append(unet(x))
metric.update(label, outputs)
avg_psnr += 10 * math.log10(1 / metric.get()[1])
metric.reset()
avg_psnr /= batches
print('Test avg psnr: {:.3f}'.format(avg_psnr))
def train(args):
np.random.seed(args.seed)
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# dataloader
transform = utils.Compose([utils.Scale(args.image_size),
utils.CenterCrop(args.image_size),
utils.ToTensor(ctx),
])
train_dataset = data.ImageFolder(args.dataset, transform)
train_loader = gluon.data.DataLoader(train_dataset, batch_size=args.batch_size, last_batch='discard')
style_loader = utils.StyleLoader(args.style_folder, args.style_size, ctx=ctx)
print('len(style_loader):',style_loader.size())
# models
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
style_model = net.Net(ngf=args.ngf)
style_model.initialize(init=mx.initializer.MSRAPrelu(), ctx=ctx)
if args.resume is not None:
print('Resuming, initializing using weight from {}.'.format(args.resume))
style_model.collect_params().load(args.resume, ctx=ctx)
print('style_model:',style_model)
# optimizer and loss
trainer = gluon.Trainer(style_model.collect_params(), 'adam',
{'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
for e in range(args.epochs):
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
# prepare data
style_image = style_loader.get(batch_id)
style_v = utils.subtract_imagenet_mean_preprocess_batch(style_image.copy())
style_image = utils.preprocess_batch(style_image)
features_style = vgg(style_v)
gram_style = [net.gram_matrix(y) for y in features_style]
xc = utils.subtract_imagenet_mean_preprocess_batch(x.copy())
f_xc_c = vgg(xc)[1]
with autograd.record():
style_model.setTarget(style_image)
y = style_model(x)
y = utils.subtract_imagenet_mean_batch(y)
features_y = vgg(y)
content_loss = 2 * args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
_, C, _ = gram_style[m].shape
gram_s = F.expand_dims(gram_style[m], 0).broadcast_to((args.batch_size, 1, C, C))
style_loss = style_loss + 2 * args.style_weight * mse_loss(gram_y, gram_s[:n_batch, :, :])
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(args.batch_size)
mx.nd.waitall()
agg_content_loss += content_loss[0]
agg_style_loss += style_loss[0]
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.3f}\tstyle: {:.3f}\ttotal: {:.3f}".format(
time.ctime(), e + 1, count, len(train_dataset),
agg_content_loss.asnumpy()[0] / (batch_id + 1),
agg_style_loss.asnumpy()[0] / (batch_id + 1),
(agg_content_loss + agg_style_loss).asnumpy()[0] / (batch_id + 1)
)
print(mesg)
if (batch_id + 1) % (4 * args.log_interval) == 0:
# save model
save_model_filename = "Epoch_" + str(e) + "iters_" + str(count) + "_" + str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".params"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
style_model.collect_params().save(save_model_path)
print("\nCheckpoint, trained model saved at", save_model_path)
# save model
save_model_filename = "Final_epoch_" + str(args.epochs) + "_" + str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".params"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
style_model.collect_params().save(save_model_path)
print("\nDone, trained model saved at", save_model_path)
def train(args):
frames = args.frames
caption_length = args.caption_length
glove_file = args.glove_file
#CPU_COUNT = multiprocessing.cpu_count()
if args.cuda:
ctx = mx.gpu()
else:
ctx = mx.cpu()
if args.load_pretrain:
pretrain_model = vision.vgg16_bn(pretrained=True,ctx=ctx)
transform = utils.Compose([utils.ToTensor(ctx),
utils.normalize(ctx),
utils.extractFeature(ctx,pretrain_model)
])
else:
pretrain_model = None
transform = utils.Compose([utils.ToTensor(ctx),
utils.normalize(ctx),
])
target_transform = utils.targetCompose([utils.WordToTensor(ctx)])
train_dataset = videoFolder(args.train_folder,args.train_dict, frames, glove_file,
caption_length, ctx, transform=transform, target_transform=target_transform)
test_dataset = videoFolder(args.test_folder,args.test_dict, frames, glove_file,
caption_length, ctx, transform=transform, target_transform=target_transform)
train_loader = gluon.data.DataLoader(train_dataset,batch_size=args.batch_size,
last_batch='discard',shuffle=True)
test_loader = gluon.data.DataLoader(test_dataset,batch_size=args.batch_size,
last_batch='discard',shuffle=False)
#loss = L2Loss_cos()
loss = L2Loss_2()
net = lstm_net(frames,caption_length,ctx,pretrained=args.load_pretrain)
#net = resnet18_v2(caption_length=caption_length,ctx=ctx)
net.collect_params().initialize(init=mx.initializer.MSRAPrelu(), ctx=ctx)
trainer = gluon.Trainer(net.collect_params(), 'adam',
{'learning_rate': args.lr})
smoothing_constant = 0.01
for e in range(args.epochs):
epoch_loss = 0
for batch_id, (x,_) in enumerate(train_loader):
with autograd.record():
pred = net(x)
batch_loss = loss(pred,_)
trainer.step(x.shape[0],ignore_stale_grad=True)
batch_loss.backward()
mx.nd.waitall()
batch_loss = F.mean(batch_loss).asscalar()
if batch_id % 100 == 0:
print("Train Batch:{}, batch_loss:{}".format(batch_id+1, batch_loss))
epoch_loss = (batch_loss if ((batch_id == 0) and (e == 0))
else (1 - smoothing_constant)*epoch_loss + smoothing_constant*batch_loss)
epoch_loss_1 = 0
for batch_id, (x,_) in enumerate(test_loader):
with autograd.predict_mode():
predict = net(x)
batch_loss_1 = loss(pred,_)
batch_loss_1 = F.mean(batch_loss_1).asscalar()
if batch_id % 100 == 0:
print("Test Batch:{}, batch_loss:{}".format(batch_id+1, batch_loss_1))
epoch_loss_1 = (batch_loss_1 if ((batch_id == 0) and (e == 0))
else (1 - smoothing_constant)*epoch_loss_1 + smoothing_constant*batch_loss_1)
print("Epoch {}, train_loss:{}, test_loss:{}".format(e+1, epoch_loss, epoch_loss_1))
if args.save_model == True:
file_name = "./saved_model/" + "lstm_pretrain.params"
net.save_parameters(file_name)
def train(args):
np.random.seed(args.seed)
if args.gpu:
ctx = [mx.gpu(0)]
else:
ctx = [mx.cpu(0)]
if args.dataset == "Sony":
out_channels = 12
scale = 2
else:
out_channels = 27
scale = 3
# load data
train_transform = utils.Compose([
utils.RandomCrop(args.patch_size, scale),
utils.RandomFlipLeftRight(),
utils.RandomFlipTopBottom(),
utils.RandomTranspose(),
utils.ToTensor(),
])
train_dataset = data.MyDataset(args.dataset,
"train",
transform=train_transform)
val_transform = utils.Compose([utils.ToTensor()])
val_dataset = data.MyDataset(args.dataset, "val", transform=val_transform)
train_loader = gluon.data.DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch_size,
last_batch='rollover')
val_loader = gluon.data.DataLoader(val_dataset,
batch_size=1,
last_batch='discard')
unet = net.UNet(out_channels, scale)
unet.initialize(init=initializer.Xavier(), ctx=ctx)
# optimizer and loss
trainer = gluon.Trainer(unet.collect_params(), 'adam',
{'learning_rate': args.lr})
l1_loss = gluon.loss.L1Loss()
print "Start training now.."
for i in range(args.epochs):
total_loss = 0
count = 0
profiler.set_state('run')
for batch_id, (img, gt) in enumerate(train_loader):
batch_size = img.shape[0]
count += batch_size
img_list = gluon.utils.split_and_load(img[0], ctx)
gt_list = gluon.utils.split_and_load(gt[0], ctx)
with autograd.record():
preds = [unet(x) for x in img_list]
losses = []
for ii in range(len(preds)):
loss = l1_loss(gt_list[ii], preds[ii])
losses.append(loss)
for loss in losses:
loss.backward()
total_loss += sum([l.sum().asscalar() for l in losses])
avg_loss = total_loss / count
trainer.step(batch_size)
metric.update(gt_list, preds)
F.waitall()
profiler.set_state('stop')
print profiler.dumps()
break
gt_save = gt_list[0]
output_save = preds[0]
if (batch_id + 1) % 100 == 0:
message = "Epoch {}: [{}/{}]: l1_loss: {:.4f}".format(
i + 1, count, len(train_dataset), avg_loss)
print message
temp = F.concat(gt_save, output_save, dim=3)
temp = temp.asnumpy().reshape(temp.shape[2], temp.shape[3], 3)
scipy.misc.toimage(temp * 255,
high=255,
low=0,
cmin=0,
cmax=255,
mode='RGB').save(args.save_model_dir +
'%04d_%05d_00_train.jpg' %
(i + 1, count))
# evaluate
batches = 0
avg_psnr = 0.
for img, gt in val_loader:
batches += 1
imgs = gluon.utils.split_and_load(img[0], ctx)
label = gluon.utils.split_and_load(gt[0], ctx)
outputs = []
for x in imgs:
outputs.append(unet(x))
metric.update(label, outputs)
avg_psnr += 10 * math.log10(1 / metric.get()[1])
metric.reset()
avg_psnr /= batches
print('Epoch {}: validation avg psnr: {:.3f}'.format(i + 1, avg_psnr))
# save model
if (i + 1) % args.save_freq == 0:
save_model_filename = "Epoch_" + str(i + 1) + ".params"
save_model_path = os.path.join(args.save_model_dir,
save_model_filename)
unet.save_params(save_model_path)
print("\nCheckpoint, trained model saved at", save_model_path)
# save model
save_model_filename = "Final_Epoch_" + str(i + 1) + ".params"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
unet.save_params(save_model_path)
print("\nCheckpoint, trained model saved at", save_model_path)
def train(args):
frames = args.frames
caption_length = args.caption_length
glove_file = args.glove_file
#CPU_COUNT = multiprocessing.cpu_count()
if args.cuda:
ctx = mx.gpu()
else:
ctx = mx.cpu()
if args.load_pretrain:
pretrain_model = vision.vgg16_bn(pretrained=True, ctx=ctx)
transform = utils.Compose([
utils.ToTensor(ctx),
utils.normalize(ctx),
utils.extractFeature(ctx, pretrain_model)
])
else:
pretrain_model = None
transform = utils.Compose([
utils.ToTensor(ctx),
utils.normalize(ctx),
])
target_transform = utils.targetCompose([utils.WordToTensor(ctx)])
train_dataset = videoFolder(args.train_folder,
args.train_dict,
frames,
glove_file,
caption_length,
ctx,
img_size=args.img_size,
transform=transform,
target_transform=target_transform)
test_dataset = videoFolder(args.test_folder,
args.test_dict,
frames,
glove_file,
caption_length,
ctx,
img_size=args.img_size,
transform=transform,
target_transform=target_transform)
train_loader = gluon.data.DataLoader(train_dataset,
batch_size=args.batch_size,
last_batch='discard',
shuffle=True)
test_loader = gluon.data.DataLoader(test_dataset,
batch_size=args.batch_size,
last_batch='discard',
shuffle=False)
loss = L2Loss()
#net = lstm_net(caption_length,ctx,pretrained=args.load_pretrain)
net = resnet18_v2(caption_length=caption_length, ctx=ctx)
net.collect_params().initialize(init=mx.initializer.MSRAPrelu(), ctx=ctx)
trainer = gluon.Trainer(net.collect_params(), 'adam',
{'learning_rate': args.lr})
train_loss = []
test_loss = []
train_loss_batch = []
test_loss_batch = []
smoothing_constant = 0.01
for e in range(args.epochs):
epoch_loss = 0.
for batch_id, (x, _) in enumerate(train_loader):
with autograd.record():
pred = net(x)
batch_loss = loss(pred, _)
trainer.step(x.shape[0], ignore_stale_grad=True)
batch_loss.backward()
mx.nd.waitall()
#print(batch_loss.shape)
batch_loss = nd.mean(batch_loss).asscalar()
#print(batch_id,batch_loss)
if ((batch_id == 0) or (e == 0)):
epoch_loss = batch_loss
else:
epoch_loss = (1 - smoothing_constant
) * epoch_loss + smoothing_constant * batch_loss
train_loss_batch.append(batch_loss)
if (batch_id + 1) % 200 == 0:
print("Train Batch:{}, batch_loss:{}".format(
batch_id + 1, batch_loss))
if ((e + 1) * (batch_id + 1)) % (2 * args.log_interval) == 0:
# save model
save_model_filename = "Epoch_" + str(e) + "_iters_" + str(
batch_id + 1) + '_' + str(time.ctime()).replace(
' ', '_') + "_" + ".params"
save_model_path = os.path.join(args.model_path,
save_model_filename)
net.save_parameters(save_model_path)
print("\nCheckpoint, trained model saved at", save_model_path)
train_loss_filename = "Epoch_" + str(e) + "_iters_" + str(
batch_id + 1) + str(time.ctime()).replace(
' ', '_') + "_train_loss" + ".txt"
train_loss_path = os.path.join(args.log_path,
train_loss_filename)
np.savetxt(train_loss_path, np.array(train_loss_batch))
epoch_loss_1 = 0.
for batch_id, (x, _) in enumerate(test_loader):
with autograd.predict_mode():
predict = net(x)
batch_loss_1 = loss(predict, _)
#batch_loss_1 = F.mean(batch_loss_1.asscalar())
batch_loss_1 = nd.mean(batch_loss_1).asscalar()
if ((batch_id == 0) or (e == 0)):
epoch_loss_1 = batch_loss_1
else:
epoch_loss_1 = (
1 - smoothing_constant
) * epoch_loss_1 + smoothing_constant * batch_loss_1
test_loss_batch.append(batch_loss_1)
if ((e + 1) * (batch_id + 1)) % (args.log_interval) == 0:
test_loss_file_name = "Epoch_" + str(e) + "_iters_" + str(
batch_id + 1) + str(time.ctime()).replace(
' ', '_') + "_test_loss" + ".txt"
test_loss_path = os.path.join(args.log_path,
test_loss_file_name)
np.savetxt(test_loss_path, np.array(test_loss_batch))
train_loss.append(epoch_loss)
test_loss.append(epoch_loss_1)
print("Epoch {}, train_loss:{}, test_loss:{}".format(
e + 1, epoch_loss, epoch_loss_1))
# save model
save_model_filename = "Final_epoch_" + str(args.epochs) + "_" + str(
time.ctime()).replace(' ', '_') + "_" + ".params"
save_model_path = os.path.join(args.model_path, save_model_filename)
net.save_parameters(save_model_path)
print("\nDone, trained model saved at", save_model_path)
train_epoch_loss_file_name = 'train_epoch_loss.txt'
test_epoch_loss_file_name = 'test_epoch_loss.txt'
train_epoch_loss_path = os.path.join(args.log_path,
train_epoch_loss_file_name)
test_epoch_loss_path = os.path.join(args.log_path,
test_epoch_loss_file_name)
np.savetxt(train_epoch_loss_path, train_loss)
np.savetxt(test_epoch_loss_path, test_loss)
def train(args):
train_logger = None
if args.log_dir is not None:
train_logger = tb.SummaryWriter(
os.path.join(args.log_dir, "train"), flush_secs=1
)
transform = utils.Compose(
[
utils.Rotation(),
utils.Crop(),
utils.Resize((32, 32)),
utils.IntensityNormalize(),
utils.ToTensor(),
]
)
train_loader = dataset.load_data(
args.data_dir, batch_size=args.batch_size, transform=transform
)
print("Train length:", len(train_loader))
model = FacialModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-3)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
print_every = 10
num_train = len(train_loader)
for epoch in range(args.epochs):
model.train()
lr = adjust_learning_rate_poly(optimizer, 1e-3, epoch, args.epochs)
running_print_loss = 0
running_print_accuracy = 0
running_accuracy = 0
for i, data in enumerate(train_loader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
acc = accuracy(labels, outputs)
running_print_accuracy += acc.item()
running_accuracy += acc.item()
running_print_loss += loss.item()
running_loss = loss.item()
if (i + 1) % print_every == 0: # print every 2000 mini-batches
print(
"[%d, %5d] loss: %.3f accuracy: %.3f lr: %.8f"
% (
epoch + 1,
i + 1,
running_print_loss / print_every,
running_print_accuracy / print_every,
lr,
)
)
running_print_loss = 0
running_print_accuracy = 0
# write train loss summaries
train_logger.add_scalar("loss", running_loss, epoch * num_train + i + 1)
train_logger.add_scalar("accuracy", running_accuracy / num_train, epoch + 1)