def run(train_datasets, test_datasets, args=None):
# Number of samples used for estimating fisher #
fisher_estimation_sample_size = 1024
# Define Model
model = Net(args)
if torch.cuda.is_available():
model.cuda()
# Set Model Optimizer
model.optimizer = optim.SGD(model.parameters(),
lr=args.lr,
weight_decay=1e-05,
nesterov=True,
momentum=0.9)
for task, train_dataset in enumerate(train_datasets):
''' Evaluate Current Net '''
test_it(model, 0, test_datasets, args, task)
for epoch in range(1, args.epochs + 1):
train(model, epoch, train_dataset, args)
''' Evaluate Current Net '''
test_it(model, epoch, test_datasets, args, task)
if args.ewc:
# Get fisher inf of parameters and consolidate it in the net
model.consolidate(
model.estimate_fisher(train_dataset,
fisher_estimation_sample_size))
def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image, size=args.content_size, keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
style_model = Net(ngf=args.ngf)
model_dict = torch.load(args.model)
model_dict_clone = model_dict.copy()
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model.load_state_dict(model_dict, False)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def predect(config):
test_dataset = ScitsrDataset(config.dev_data)
test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)
model = Net(config.num_node_features, config.num_class)
model.cuda()
model.load_state_dict(torch.load(config.model_path))
val_iter = iter(test_loader)
# data = val_iter.next()
n_correct = 0
n_total = 0
for data in val_iter:
out_pred = model(data)
_, pred = out_pred.max(dim=1)
label = data.y.detach().cpu().numpy()
pred = pred.detach().cpu().numpy()
epoch_ped = (label == pred).sum()
if (label == pred).all():
n_correct = n_correct + 1
n_total = n_total + 1
logging.info("len_label = %d, epoch_ped = %d, n_correct = %d, n_total = %d"%
(len(label), epoch_ped, n_correct, n_total))
accuracy = n_correct / float(n_total)
logging.info('Test accuray: %f' % (accuracy))
def main():
train_loader = data_loader.load_train(params['data_path'], params['batch_size'])
val_loader = data_loader.load_val(params['data_path'], params['batch_size'])
print('train and val data loaded')
model = Net()
model.cuda()
optimizer = optim.Adam(model.parameters(), lr=params['lr'])
epoch = 0
epoches_to_train = params['epoches']
train_logger = TrainLogger()
# epoch, train_logger = load_checkpoint(model, optimizer, 'checkpoints/1/epoch_3.pth')
# epoches_to_train = 10
train(model, optimizer, epoch, epoches_to_train, train_loader, val_loader, params['checkpoint_path'], train_logger)
class Agent(object):
def __init__(self, envs, args):
self.envs = envs
self.args = args
self.net = Net(self.env.action_space.n)
if self.args.cuda:
self.net.cuda()
self.optimizer = torch.optim.RMSprop(self.net.parameters(), lr=self.args.lr, eps=self.args.eps)
if not os.path.exists(self.args.save_path):
os.mkdir(self.args.save_path)
self.model_path = self.args.save_path + self.args.env_name + '/'
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
def train(save_directory: str,
model_path: str = None,
epochs=10,
validate=0.2):
data = SeedlingsData()
data.load(train_data_paths=[
constants.train_output_resize_file_path,
constants.train_output_rotate_file_path,
constants.train_output_crop_file_path,
constants.train_output_mirror_file_path
],
test_data_paths=[constants.test_output_resize_file_path],
validate=validate)
data.set_batch_size(128)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if model_path is not None:
net = load_model(model_path)
else:
# Create network
print(
"select model in ['resnet50+', 'resnet50','resnet101', 'resnet152', 'densenet161', 'densenet201', 'inception_v3']"
)
model = input("model: ")
net = Net(model)
print(net)
net.cuda()
optimizer = optim.Adam(net.parameters(), lr=1e-4, weight_decay=1e-5)
for epoch in range(0, epochs):
train_epoch(net, data, epoch, normalize, optimizer)
if data.validate >= 0.0001:
accuracy = validate_epoch(net, data, epoch, normalize)
accuracy_list.append(accuracy)
save_model(net, save_directory, accuracy=accuracy, prefix='model')
if accuracy_list.index(
max(accuracy_list)) == len(accuracy_list) - 1:
save_model(net, save_directory, is_best=True, prefix='model')
# validate_analysis(net, data, normalize)
else:
save_model(net, save_directory, is_best=True, prefix='model')
del net
def main():
from train import Trainer
net = Net(VinConfig)
net = net.cuda()
net = net.double()
trainer = Trainer(VinConfig, net)
trainer.train()
def main():
global opt, model
opt = parser.parse_args()
print opt
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
print("===> Loading datasets")
train_set = DatasetFromHdf5("../lapsrn/data/data.h5")
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
print("===> Building model")
model = Net()
criterion = L1_Charbonnier_loss()
print("===> Setting GPU")
if cuda:
model = model.cuda()
criterion = criterion.cuda()
else:
model = model.cpu()
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training")
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
train(training_data_loader, optimizer, model, criterion, epoch)
save_checkpoint(model, epoch)
def main():
global opt, model
opt = parser.parse_args()
print(opt)
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
print("===> Loading datasets")
# train_set = DatasetFromHdf5("path_to_dataset.h5")
train_set = get_training_set()
training_data_loader = DataLoader(dataset=train_set,
batch_size=opt.batchSize,
shuffle=True)
print("===> Building model")
model = Net()
# criterion = nn.L1Loss(size_average=False)
criterion = nn.L1Loss()
print("===> Setting GPU")
if cuda:
model = model.cuda()
criterion = criterion.cuda()
wandb.watch(model)
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = 374
model.load_state_dict(checkpoint.state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
print("===> Setting Optimizer")
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=opt.lr,
weight_decay=opt.weight_decay,
betas=(0.9, 0.999),
eps=1e-08)
# optimizer = optim.Adam(model.parameters(), lr=opt.lr,
# weight_decay=opt.weight_decay, betas=(0.9, 0.999), eps=1e-08)
print("===> Training")
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
train(training_data_loader, optimizer, model, criterion, epoch)
save_checkpoint(model, epoch)
def test_5_class(root):
correct = [0, 0, 0, 0, 0]
test_data = MyDataset(txt=root+'label/mytest.txt', transform=transforms.ToTensor())
test_loader = DataLoader(dataset=test_data, batch_size=64)
test_num_list = [0,0,0,0,0]
file_read = open(root + 'label/mytest.txt', 'r')
file_list = file_read.readlines()
file_read.close()
for i in range(len(file_list)):
index = int(file_list[i].split(' ')[1])
test_num_list[index] += 1
model = Net()
model.load_state_dict(torch.load('weight/result.pkl'))
if(use_gpu):
model = model.cuda()
model.eval()
for batch_x, batch_y in test_loader:
batch_x, batch_y = Variable(batch_x, volatile=True), Variable(batch_y, volatile=True)
if(use_gpu):
batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
out = model(batch_x)
pred = torch.max(out, 1)[1]
for i in range(len(pred)):
if(pred[i].data[0] == batch_y[i].data[0]):
correct[pred[i].data[0]] += 1
for i in range(len(correct)):
correct[i] /= test_num_list[i]
print(correct)
def test_5_possibility(root):
test_data = MyDataset(txt=root+'label/simple_test.txt', transform=transforms.ToTensor())
test_loader = DataLoader(dataset=test_data, batch_size=64)
model = Net()
model.load_state_dict(torch.load('weight/result.pkl'))
if(use_gpu):
model = model.cuda()
model.eval()
for batch_x, batch_y in test_loader:
batch_x, batch_y = Variable(batch_x, volatile=True), Variable(batch_y, volatile=True)
if(use_gpu):
batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
out = model(batch_x)
if(use_gpu):
out = out.cpu()
out_array = out.data.numpy()
for i in range(len(out_array)):
for j in range(len(out_array[i])):
if out_array[i][j] < 0:
out_array[i][j] = 0
total = sum(out_array[i])
for j in range(len(out_array[i])):
out_array[i][j] /= total
print(out_array[i])
def main():
parser = argparse.ArgumentParser(
description='PyTorch Video Frame Interpolation via Residue Refinement')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
transform = transforms.ToTensor()
model = Net()
state = torch.load('pretrained_model.pth.tar')
model.load_state_dict(state, strict=True)
model = model.cuda()
model.eval()
im1_path = 'data/im1.png'
im2_path = 'data/im2.png'
with torch.no_grad():
img1 = Image.open(im1_path)
img2 = Image.open(im2_path)
img1 = transform(img1)
img2 = transform(img2)
img1, img2 = img1.cuda(), img2.cuda()
output = model(img1.unsqueeze(0), img2.unsqueeze(0))
output = output.squeeze(0).cpu()
output = transforms.functional.to_pil_image(output)
output.save('data/im_interp.png')
def datagen_compress(f):
# Initialize smaller network that uses 20% fewer nodes in the hidden layers
hidden_dim = f.module.hidden_dim if args.cuda else f.hidden_dim
g = Net(hidden_dim=int(ceil(hidden_dim * 0.8)))
g = torch.nn.DataParallel(g.cuda()) if args.cuda else g
return g
def test(datafile, modeldir):
fp = open(datafile, 'rb')
data = pickle.load(fp)
x_test_pos = data['X_test_pos']
x_test_neg = data['X_test_neg']
y_test_pos = np.ones((len(x_test_pos), 1), dtype=np.float32)
y_test_neg = np.zeros((len(x_test_neg), 1), dtype=np.float32)
x_test = np.vstack((x_test_pos, x_test_neg))
y_test = np.vstack((y_test_pos, y_test_neg))
test_dataset = TensorDataset(
torch.from_numpy(x_test).cuda(),
torch.from_numpy(y_test).cuda())
test_loader = DataLoader(dataset=test_dataset,
batch_size=32,
shuffle=False)
net = Net()
net.cuda()
label_loss = nn.BCELoss(reduction='none')
modelfile = modeldir + 'checkpoint.pkl'
net.load_state_dict(torch.load(modelfile))
net.eval()
prob = []
attscore = []
for i, (x, y) in enumerate(test_loader):
x = Variable(x)
y = Variable(y)
with torch.no_grad():
output = net(x)
prob.append(output.cpu().data.numpy())
loss = torch.mean(label_loss(output, y))
attscore.append(net.att_x.cpu().data.numpy())
prob = np.concatenate(prob)
attscore = np.concatenate(attscore)
vfpr, vtpr, vthresholds = metrics.roc_curve(y_test, prob, pos_label=1)
score = np.concatenate([y_test, prob], 1)
vauc = metrics.auc(vfpr, vtpr)
print(vauc)
np.save(modeldir + 'attscore.npy', attscore)
np.save(modeldir + 'score.npy', score)
class DQN():
def __init__(self):
self.evaluate_net = Net(N_ACTIONS)
self.target_net = Net(N_ACTIONS)
self.optimizer = torch.optim.RMSprop(self.evaluate_net.parameters(), lr=LR, alpha=0.95, eps=0.01) ## 10.24 fix alpha and eps
self.loss_func = torch.nn.MSELoss()
self.replay_memory = ReplayMemory(REPLAY_MEMORY_SIZE)
self.learn_step_counter = 0
self.evaluate_net.cuda()
self.target_net.cuda()
def select_action(self, s, epsilon):
if random.random() > epsilon:
q_eval = self.evaluate_net.forward(s)
action = q_eval[0].max(0)[1].cpu().data.numpy() ## 10.21 to cpu
else:
action = np.asarray(random.randrange(N_ACTIONS))
return action
def store_transition(self, s, a, r, s_):
self.replay_memory.store(s, a, r, s_)
def learn(self, ):
if self.learn_step_counter % TARGET_REPLACE_FREQ == 0:
self.target_net.load_state_dict(self.evaluate_net.state_dict())
self.learn_step_counter += 1
s_s, a_s, r_s, s__s = self.replay_memory.sample(BATCH_SIZE)
q_eval = self.evaluate_net(s_s).gather(1, a_s)
q_next = self.target_net(s__s).detach()
q_target = r_s + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)
loss = self.loss_func(q_eval, q_target)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss
def main():
train_loader = torch.utils.data.DataLoader(
vimeo_dataloader("/root/toflow_keras/dataset/vimeo_triplet"),
batch_size=8,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True)
model = Net()
# model._check_gradients(SeparableConvolutionModule())
model = model.cuda()
criterion = L1_Charbonnier_loss().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
for epoch in range(50):
train(train_loader, model, optimizer, criterion, epoch)
check_point_filename = 'weights/checkpoint_epoch_%02d.pth' % epoch
torch.save(model.state_dict(), check_point_filename)
def main(args):
# ===============================================
# Build Data Loader
# ===============================================
train_transformer = transformer.Compose([
transformer.RandomHorizontalFlip(),
transformer.RandomScaleCrop(),
transformer.ArrayToTensor(),
transformer.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
train_dataset = Data('train',
transformer = train_transformer
seed = args.seed,
train = True,
seq_length = 3)
valid_transformer = transformer.Compose([transformer.ArrayToTensor(), normalize])
val_dataset = Data('val',
transformer = valid_transformer
seed = args.seed,
train = True,
seq_length = 3)
train_loader = DataLoader(train_dataset,
batch_size = args.batch_size,
shuffle = True,
num_workers = cfg.workers,
pin_memory = True)
val_loader = DataLoader(val_dataset,
batch_size = args.batch_size,
shuffle = True,
num_workers = cfg.workers,
pin_memory = True)
# ===============================================
# Build Model
# ===============================================
model = Net()
model = model.cuda()
if args.load:
pass
else:
model.init_weights()
model = DataParallel(model)
def main():
parser = argparse.ArgumentParser(
description='PyTorch Video Frame Interpolation via Residue Refinement')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
transform = transforms.ToTensor()
model = Net()
state = torch.load('pretrained_model.pth.tar')
model.load_state_dict(state, strict=True)
model = model.cuda()
model.eval()
im1_path = 'data/im1.png'
im2_path = 'data/im2.png'
with torch.no_grad():
img1 = Image.open(im1_path)
img2 = Image.open(im2_path)
img1 = transform(img1).unsqueeze(0).cuda()
img2 = transform(img2).unsqueeze(0).cuda()
if img1.size(1) == 1:
img1 = img1.expand(-1, 3, -1, -1)
img2 = img2.expand(-1, 3, -1, -1)
_, _, H, W = img1.size()
H_, W_ = int(np.ceil(H / 32) * 32), int(np.ceil(W / 32) * 32)
pader = torch.nn.ReplicationPad2d([0, W_ - W, 0, H_ - H])
img1, img2 = pader(img1), pader(img2)
output = model(img1, img2)
output = output[0, :, 0:H, 0:W].squeeze(0).cpu()
output = transforms.functional.to_pil_image(output)
output.save('data/im_interp.png')
def main():
data_set = DataSet("data/train5/")
train_data_loader = DataLoader(dataset=data_set,
num_workers=4,
batch_size=64,
shuffle=True)
model = Net()
criterion = nn.MSELoss(size_average=False)
if torch.cuda.is_available():
model = model.cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
for epoch in range(20):
train(train_data_loader, model, criterion, optimizer, epoch)
save_checkpoint(model, epoch)
print("train done--------->")
def mask_compress(f, THRESHOLD=0.05):
# Faster than `deepcopy` and we only need the weights
g = Net()
g.load_state_dict(f.state_dict())
g = torch.nn.DataParallel(g.cuda()) if args.cuda else g
# Mask weights/biases
for l in g.parameters():
# 1 if larger than `THRESHOLD`, 0 otherwise
mask = (torch.abs(l.data) > THRESHOLD).float()
l.data.mul_(mask)
# TODO make this work on individual connection level rather than whole layer
# # ensure zero-weight parameters are never updated to mimic pruning
# for param in g.parameters():
# if param == 0:
# param.requires_grad = False
return g
def main():
warnings.filterwarnings('ignore')
setproctitle('jmds.Almonds')
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
model = Net()
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.SGD(
model.parameters(),
lr=model.learning_rate,
momentum=model.momentum
)
if torch.cuda.is_available():
print(" * CUDA AVAILABLE, USING IT")
model = model.cuda()
criterion = criterion.cuda()
print(" * LOADING DATASET")
jal = JMDSAmendoasLoader()
train_batches = jal.get_train_batches(40)
print(" * RUNNING TEACHING PROCESS")
for epoch in range(1, 30):
train(
model,
optimizer,
criterion,
epoch,
train_batches
)
test_print_save(
model,
criterion,
epoch,
train_batches,
jal
)
def load_model(model_path, data, scale_factor, cuda):
net = Net(upscale_factor=scale_factor)
if cuda:
net = net.cuda()
# if not cuda:
# net.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
# else:
# net.load_state_dict(torch.load(model_path))
if os.name == 'posix':
net = torch.load(model_path)
else:
from functools import partial
import pickle
pickle.load = partial(pickle.load, encoding="latin1")
pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
net = torch.load(model_path,
map_location=lambda storage, loc: storage,
pickle_module=pickle)
transform = ToTensor()
ori_tensor = transform(data)
if cuda:
ori_tensor = Variable(ori_tensor.cuda())
else:
ori_tensor = Variable(ori_tensor)
ori_tensor = torch.unsqueeze(ori_tensor, 0)
output = net(ori_tensor)
if cuda:
result = (output.data).cpu().numpy()
else:
result = (output.data).numpy()
result = result[0, 0, :, :]
return result
class Extractor(object):
def __init__(self, model_path, use_cuda=True):
self.net = Net(reid=True)
# self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
state_dict = torch.load(model_path)['net_dict']
self.net.load_state_dict(state_dict)
print("Loading weights from {}... Done!".format(model_path))
self.net = self.net.cuda()
# self.net.to(self.device)
self.norm = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
def __call__(self, img):
assert isinstance(img, np.ndarray), "type error"
img = img.astype(np.float) #/255.
img = cv2.resize(img, (64, 128))
img = torch.from_numpy(img).float().permute(2, 0, 1)
img = self.norm(img).unsqueeze(0)
with torch.no_grad():
img = img.cuda() #to(self.device)
feature = self.net(img)
return feature.cpu().numpy()
def main():
print("Initialising model")
model = Net(config, emb)
print("Model initialised")
if torch.cuda.is_available():
model = model.cuda()
#load older model if it exists
if os.path.isfile('./squad_model.pth'):
model.load_state_dict(torch.load('./squad_model.pth'))
for epoch in range(args.epochs):
print("Starting epoch " + str(epoch + 1))
trn_loss, trn_accuracy = train(epoch, model)
print("trn_accuarcy: ", str(trn_accuracy))
print("trn_loss: ", str(trn_loss))
dev_loss, dev_accuracy = validation(model)
print("dev_accuarcy: ", str(dev_accuracy))
print("dev_loss: ", str(dev_loss))
if (epoch % 5 == 4):
model_file = 'model_' + str(epoch) + '.pth'
torch.save(model.state_dict(), model_file)
print('\nSaved model to ' + model_file + '.')
def main():
"""The main function
Entry point.
"""
global args
# Setting the hyper parameters
parser = argparse.ArgumentParser(description='Example of Capsule Network')
parser.add_argument('--epochs',
type=int,
default=10,
help='number of training epochs. default=10')
parser.add_argument('--lr',
type=float,
default=0.01,
help='learning rate. default=0.01')
parser.add_argument('--batch-size',
type=int,
default=128,
help='training batch size. default=128')
parser.add_argument('--test-batch-size',
type=int,
default=128,
help='testing batch size. default=128')
parser.add_argument(
'--log-interval',
type=int,
default=10,
help=
'how many batches to wait before logging training status. default=10')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training. default=false')
parser.add_argument(
'--threads',
type=int,
default=4,
help='number of threads for data loader to use. default=4')
parser.add_argument('--seed',
type=int,
default=42,
help='random seed for training. default=42')
parser.add_argument(
'--num-conv-out-channel',
type=int,
default=256,
help='number of channels produced by the convolution. default=256')
parser.add_argument(
'--num-conv-in-channel',
type=int,
default=1,
help='number of input channels to the convolution. default=1')
parser.add_argument('--num-primary-unit',
type=int,
default=8,
help='number of primary unit. default=8')
parser.add_argument('--primary-unit-size',
type=int,
default=1152,
help='primary unit size is 32 * 6 * 6. default=1152')
parser.add_argument(
'--num-classes',
type=int,
default=2,
help='number of digit classes. 1 unit for one MNIST digit. default=10')
parser.add_argument('--output-unit-size',
type=int,
default=1,
help='output unit size. default=16')
parser.add_argument('--num-routing',
type=int,
default=3,
help='number of routing iteration. default=3')
parser.add_argument(
'--use-reconstruction-loss',
type=utils.str2bool,
nargs='?',
default=True,
help='use an additional reconstruction loss. default=True')
parser.add_argument(
'--regularization-scale',
type=float,
default=0.0005,
help=
'regularization coefficient for reconstruction loss. default=0.0005')
parser.add_argument('--dataset',
help='the name of dataset (mnist, cifar10)',
default='mnist')
parser.add_argument(
'--input-width',
type=int,
default=28,
help='input image width to the convolution. default=28 for MNIST')
parser.add_argument(
'--input-height',
type=int,
default=28,
help='input image height to the convolution. default=28 for MNIST')
parser.add_argument('--is-training',
type=int,
default=1,
help='Whether or not is training, default is yes')
parser.add_argument('--weights',
type=str,
default=None,
help='Load pretrained weights, default is none')
args = parser.parse_args()
print(args)
# Check GPU or CUDA is available
args.cuda = not args.no_cuda and torch.cuda.is_available()
# Get reproducible results by manually seed the random number generator
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# Load data
train_loader, test_loader = utils.load_data(args)
# Build Capsule Network
print('===> Building model')
model = Net(num_conv_in_channel=args.num_conv_in_channel,
num_conv_out_channel=args.num_conv_out_channel,
num_primary_unit=args.num_primary_unit,
primary_unit_size=args.primary_unit_size,
num_classes=args.num_classes,
output_unit_size=args.output_unit_size,
num_routing=args.num_routing,
use_reconstruction_loss=args.use_reconstruction_loss,
regularization_scale=args.regularization_scale,
input_width=args.input_width,
input_height=args.input_height,
cuda_enabled=args.cuda)
if args.cuda:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
model.cuda()
cudnn.benchmark = True
model = torch.nn.DataParallel(model)
# Print the model architecture and parameters
print('Model architectures:\n{}\n'.format(model))
print('Parameters and size:')
for name, param in model.named_parameters():
print('{}: {}'.format(name, list(param.size())))
# CapsNet has:
# - 8.2M parameters and 6.8M parameters without the reconstruction subnet on MNIST.
# - 11.8M parameters and 8.0M parameters without the reconstruction subnet on CIFAR10.
num_params = sum([param.nelement() for param in model.parameters()])
# The coupling coefficients c_ij are not included in the parameter list,
# we need to add them manually, which is 1152 * 10 = 11520 (on MNIST) or 2048 * 10 (on CIFAR10)
print('\nTotal number of parameters: {}\n'.format(
num_params + (11520 if args.dataset == 'mnist' else 20480)))
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# Make model checkpoint directory
if not os.path.exists('results/trained_model'):
os.makedirs('results/trained_model')
# Set the logger
writer = SummaryWriter()
if not args.is_training:
print("Loading checkpoint")
checkpoint = torch.load('results/trained_model/model_epoch_50.pth',
map_location=lambda storage, loc: storage)
from collections import OrderedDict
state_dict = checkpoint['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
if name[:2] == 'fc':
name = 'decoder.' + name
new_state_dict[name] = v
epoch = checkpoint['epoch']
model.load_state_dict(new_state_dict)
optimizer.load_state_dict(checkpoint['optimizer'])
# Train and test
for epoch in range(1, args.epochs + 1):
if args.is_training:
train(model, train_loader, optimizer, epoch, writer)
test(model, test_loader, len(train_loader), epoch, writer)
# Save model checkpoint
utils.checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, epoch)
writer.close()
overlay_y = parameters.get_digit_parameters("", "train_pixel_overlay_y", None,
'int')
crop_height = parameters.get_digit_parameters("", "crop_height", None, 'int')
crop_width = parameters.get_digit_parameters("", "crop_width", None, 'int')
dataset = RemoteSensingImg(args.dataroot, args.list, patch_w, patch_h,
overlay_x, overlay_y, args.edge)
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batchSize,
num_workers=args.workers,
shuffle=True)
# torch.backends.cudnn.enabled = True
############## create model
model = Net(args.useBN)
if args.cuda:
model.cuda()
cudnn.benchmark = True
############## resume
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.cuda == False:
checkpoint = torch.load(args.resume,
map_location={'cuda:0': 'cpu'})
else:
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
class Demo(BaseDemo):
def __init__(self, args):
super(Demo, self).__init__(args)
self.model, self.model_gt = self.init_model(self.data.m_kernel)
self.visualizer = Visualizer(args, self.data.reverse_m_dict)
def init_model(self, m_kernel):
self.model = Net(self.im_size, self.im_size, self.im_channel,
self.num_frame - 1, m_kernel.shape[1], self.m_range,
m_kernel)
self.model_gt = GtNet(self.im_size, self.im_size, self.im_channel,
self.num_frame - 1, m_kernel.shape[1],
self.m_range, m_kernel)
if torch.cuda.is_available():
# model = torch.nn.DataParallel(model).cuda()
self.model = self.model.cuda()
self.model_gt = self.model_gt.cuda()
if self.init_model_path is not '':
self.model.load_state_dict(torch.load(self.init_model_path))
return self.model, self.model_gt
def train(self):
optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
base_loss, train_loss = [], []
for epoch in range(self.train_epoch):
optimizer.zero_grad()
if self.data.name in ['box', 'mnist', 'box_complex']:
im, _, _, _ = self.data.get_next_batch(self.data.train_images)
elif self.data.name in ['box2', 'mnist2']:
im, _, _ = self.data.get_next_batch(self.data.train_images)
elif self.data.name in ['robot64', 'mpii64', 'nyuv2']:
im = self.data.get_next_batch(self.data.train_meta)
else:
logging.error('%s data not supported' % self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float())
im_output = Variable(torch.from_numpy(im_output).float())
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
im_pred, m_mask, d_mask = self.model(im_input)
im_diff = im_pred - im_output
loss = torch.abs(im_diff).sum()
loss.backward()
optimizer.step()
train_loss.append(loss.data[0])
if len(train_loss) > 100:
train_loss.pop(0)
ave_train_loss = sum(train_loss) / float(len(train_loss))
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
if len(base_loss) > 100:
base_loss.pop(0)
ave_base_loss = sum(base_loss) / float(len(base_loss))
logging.info(
'epoch %d, train loss: %.2f, average train loss: %.2f, base loss: %.2f',
epoch, loss.data[0], ave_train_loss, ave_base_loss)
if (epoch + 1) % self.test_interval == 0:
logging.info('epoch %d, testing', epoch)
self.validate()
def test(self):
base_loss, test_loss = [], []
test_epe = []
for epoch in range(self.test_epoch):
if self.data.name in ['box', 'mnist', 'box_complex']:
im, motion, _, _ = self.data.get_next_batch(
self.data.test_images)
elif self.data.name in ['box2', 'mnist2']:
im, motion, _ = self.data.get_next_batch(self.data.test_images)
elif self.data.name in ['robot64', 'mpii64', 'nyuv2']:
im, motion = self.data.get_next_batch(
self.data.test_meta), None
elif self.data.name in ['mpii64_sample']:
im, motion = self.data.get_next_batch(
self.data.test_meta), None
im = im[:, -self.num_frame:, :, :, :]
else:
logging.error('%s data not supported' % self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float())
im_output = Variable(torch.from_numpy(im_output).float())
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
im_pred, m_mask, d_mask = self.model(im_input)
im_diff = im_pred - im_output
loss = torch.abs(im_diff).sum()
test_loss.append(loss.data[0])
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
flow = self.motion2flow(m_mask)
depth = self.mask2depth(d_mask)
if motion is None:
gt_motion = None
else:
gt_motion = motion[:, -2, :, :, :]
gt_motion = Variable(torch.from_numpy(gt_motion).float())
if torch.cuda.is_available():
gt_motion = gt_motion.cuda()
epe = (flow - gt_motion) * (flow - gt_motion)
epe = torch.sqrt(epe.sum(1))
epe = epe.sum() / epe.numel()
test_epe.append(epe.cpu().data[0])
if self.display:
self.visualizer.visualize_result(im_input, im_output, im_pred,
flow, gt_motion, depth,
'test_%d.png' % epoch)
if self.display_all:
for i in range(self.batch_size):
self.visualizer.visualize_result(im_input, im_output,
im_pred, flow, gt_motion,
depth, 'test_%d.png' % i,
i)
test_loss = numpy.mean(numpy.asarray(test_loss))
base_loss = numpy.mean(numpy.asarray(base_loss))
improve_loss = base_loss - test_loss
improve_percent = improve_loss / (base_loss + 1e-5)
logging.info('average test loss: %.2f, base loss: %.2f', test_loss,
base_loss)
logging.info('improve_loss: %.2f, improve_percent: %.2f', improve_loss,
improve_percent)
if gt_motion is not None:
test_epe = numpy.mean(numpy.asarray(test_epe))
logging.info('average test endpoint error: %.2f', test_epe)
return improve_percent
def test_gt(self):
base_loss, test_loss = [], []
test_epe = []
for epoch in range(self.test_epoch):
if self.data.name in ['box', 'mnist', 'box_complex']:
im, motion, motion_label, depth = self.data.get_next_batch(
self.data.test_images)
gt_motion_label = motion_label[:, -2, :, :, :]
gt_motion_label = Variable(torch.from_numpy(gt_motion_label))
if torch.cuda.is_available():
gt_motion_label = gt_motion_label.cuda()
elif self.data.name in ['box2', 'mnist2']:
im, motion, depth = self.data.get_next_batch(
self.data.test_images)
else:
logging.error('%s data not supported in test_gt' %
self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
gt_motion = motion[:, -2, :, :, :]
gt_depth = depth[:, -2, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float())
im_output = Variable(torch.from_numpy(im_output).float())
gt_motion = Variable(torch.from_numpy(gt_motion).float())
gt_depth = Variable(torch.from_numpy(gt_depth).float())
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
gt_motion = gt_motion.cuda()
gt_depth = gt_depth.cuda()
if self.data.name in ['box', 'mnist', 'box_complex']:
im_pred, m_mask, d_mask = self.model_gt(
im_input, gt_motion_label, gt_depth, 'label')
elif self.data.name in ['box2', 'mnist2']:
im_pred, m_mask, d_mask = self.model_gt(
im_input, gt_motion, gt_depth)
im_diff = im_pred - im_output
loss = torch.abs(im_diff).sum()
test_loss.append(loss.data[0])
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
flow = self.motion2flow(m_mask)
depth = self.mask2depth(d_mask)
epe = (flow - gt_motion) * (flow - gt_motion)
epe = torch.sqrt(epe.sum(1))
epe = epe.sum() / epe.numel()
test_epe.append(epe.cpu().data[0])
if self.display:
self.visualizer.visualize_result(im_input, im_output, im_pred,
flow, gt_motion, depth,
'test_gt.png')
if self.display_all:
for i in range(self.batch_size):
self.visualizer.visualize_result(im_input, im_output,
im_pred, flow, gt_motion,
depth,
'test_gt_%d.png' % i, i)
test_loss = numpy.mean(numpy.asarray(test_loss))
base_loss = numpy.mean(numpy.asarray(base_loss))
improve_loss = base_loss - test_loss
improve_percent = improve_loss / (base_loss + 1e-5)
logging.info('average ground truth test loss: %.2f, base loss: %.2f',
test_loss, base_loss)
logging.info('improve_loss: %.2f, improve_percent: %.2f', improve_loss,
improve_percent)
test_epe = numpy.mean(numpy.asarray(test_epe))
logging.info('average ground truth test endpoint error: %.2f',
test_epe)
return improve_percent
def mask2depth(self, d_mask):
[batch_size, num_depth, height, width] = d_mask.size()
depth_number = Variable(
torch.zeros(batch_size, num_depth, height, width))
if torch.cuda.is_available():
depth_number = depth_number.cuda()
for i in range(num_depth):
depth_number[:, i, :, :] = i
depth = Variable(torch.zeros(batch_size, 1, height, width))
if torch.cuda.is_available():
depth = depth.cuda()
depth[:, 0, :, :] = (d_mask * depth_number).sum(1)
return depth
def label_data(args, hard_model_name, easy_model_name, bounding_boxes,
id_to_class):
# Loading models
hard_model_dict = torch.load(args.experiment + '/' + hard_model_name)
easy_model_dict = torch.load(args.experiment + '/' + easy_model_name)
use_cuda = torch.cuda.is_available()
hard_model = Net()
easy_model = Net()
if use_cuda:
hard_model.cuda()
easy_model.cuda()
hard_model.load_state_dict(hard_model_dict)
hard_model.eval()
easy_model.load_state_dict(easy_model_dict)
easy_model.eval()
from data import data_transforms
image_dir = args.data + '/self_supervised_images/'
hard_output_dir = args.data + '/self_supervised/'
easy_output_dir = args.data + '/easy_self_supervised/'
def pil_loader(path):
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
print("Self-supervised learning running...")
for f in tqdm(os.listdir(image_dir), position=0, leave=True):
if 'jpg' in f:
data = data_transforms['val'](pil_loader(image_dir + '/' + f))
data = data.view(1, data.size(0), data.size(1),
data.size(2)).cuda()
hard_output = hard_model(data)
hard_output = F.softmax(hard_output)
hard_prob = hard_output.data.max(1, keepdim=True)[0]
hard_pred = hard_output.data.max(1, keepdim=True)[1]
hard_pred = hard_pred.item()
data = data_transforms['detect'](pil_loader(image_dir + '/' + f))
data = data.view(1, data.size(0), data.size(1),
data.size(2)).cuda()
box = bounding_boxes[bounding_boxes.index == f[:32]]
a, b, c, d = int(box['x1']), int(box['y1']), int(box['x2']), int(
box['y2'])
cropped = data[:, :, b:d + b, a:c + a]
cropped = F.interpolate(cropped, size=384)
easy_output = easy_model(cropped)
easy_output = F.softmax(easy_output)
easy_prob = easy_output.data.max(1, keepdim=True)[0]
easy_pred = easy_output.data.max(1, keepdim=True)[1]
easy_pred = easy_pred.item()
# Only label new data if very confident (--> few mis-labels)
if easy_pred == hard_pred:
if easy_prob > 0.9:
shutil.copy(image_dir + '/' + f,
hard_output_dir + '/' + id_to_class[easy_pred])
if hard_prob > 0.9:
save_image(
cropped, easy_output_dir + '/' +
id_to_class[easy_pred] + '/' + f)
# Free memory
del easy_model
del hard_model
torch.cuda.empty_cache()
print("Done!")
audio_retrieval_val = audio_retrieval
face_ind_test = np.load(os.path.join(root_feat,'no_face_ind_retrieval.npy'))
face_ind_test = 1 - face_ind_test
print 'Done.'
# Model
video_modality_dim = {'face': (128,128), 'audio': (128*16,128),
'visual': (2048,2048), 'motion': (1024,1024)}
net = Net(video_modality_dim,300,
audio_cluster=16,text_cluster=args.text_cluster_size)
net.train()
if args.GPU:
net.cuda()
# Optimizers + Loss
max_margin = MaxMarginRankingLoss(margin=args.margin)
if args.optimizer == 'adam':
optimizer = optim.Adam(net.parameters(), lr=args.lr)
elif args.optimizer == 'sgd':
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum)
if args.GPU:
max_margin.cuda()
n_display = args.n_display
dataset_size = len(dataset)
we_dim=args.we_dim,
n_pair=args.n_pair,
max_words=args.max_words,
sentence_dim=args.sentence_dim,
)
net.train()
# Optimizers + Loss
loss_op = MaxMarginRankingLoss(
margin=args.margin,
negative_weighting=args.negative_weighting,
batch_size=args.batch_size,
n_pair=args.n_pair,
hard_negative_rate=args.hard_negative_rate,
)
net.cuda()
loss_op.cuda()
if args.pretrain_path != '':
net.load_checkpoint(args.pretrain_path)
optimizer = optim.Adam(net.parameters(), lr=args.lr)
if args.verbose:
print('Starting training loop ...')
def TrainOneBatch(model, opt, data, loss_fun):
text = data['text'].cuda()
video = data['video'].cuda()
video = video.view(-1, video.shape[-1])
class Demo(BaseDemo):
def __init__(self, args):
super(Demo, self).__init__(args)
self.model, self.model_gt = self.init_model(self.data.m_kernel)
self.visualizer = Visualizer(args, self.data.reverse_m_dict)
def init_model(self, m_kernel):
self.model = Net(self.im_size, self.im_size, self.im_channel,
self.num_frame, m_kernel.shape[1], self.m_range,
self.net_depth)
self.model_gt = GtNet(self.im_size, self.im_size, self.im_channel,
self.num_frame, m_kernel.shape[1], self.m_range)
if torch.cuda.is_available():
# model = torch.nn.DataParallel(model).cuda()
self.model = self.model.cuda()
self.model_gt = self.model_gt.cuda()
if self.init_model_path is not '':
self.model.load_state_dict(torch.load(self.init_model_path))
return self.model, self.model_gt
def train(self):
optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
base_loss, train_loss = [], []
for epoch in range(self.train_epoch):
optimizer.zero_grad()
if self.data.name in ['box', 'mnist', 'box_complex']:
im, _, _, _ = self.data.get_next_batch(self.data.train_images)
elif self.data.name in ['box2', 'mnist2']:
im, _, _ = self.data.get_next_batch(self.data.train_images)
elif self.data.name in [
'robot64', 'mpii64', 'mpi128', 'nyuv2', 'robot128',
'viper64', 'viper128', 'robot128c'
]:
im = self.data.get_next_batch(self.data.train_images)
else:
logging.error('%s data not supported' % self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float())
im_output = Variable(torch.from_numpy(im_output).float())
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
im_pred, flow = self.model(im_input, im_output)
flow = flow * self.im_size / 2 # resize flow from [-1, 1] back to image scale
im_diff = im_pred - im_input[:, -self.
im_channel:, :, :] # inverse warping loss
loss = torch.abs(im_diff).sum()
# loss = (im_diff * im_diff).sum()
flow_diff1 = flow[:, :, 1:, :] - flow[:, :, :-1, :]
flow_diff2 = flow[:, :, :, 1:] - flow[:, :, :, :-1]
loss = loss + 0.1 * (torch.abs(flow_diff1).sum() +
torch.abs(flow_diff2).sum())
# loss = loss + 0.1 * ((flow_diff1 * flow_diff1).sum() + (flow_diff2 * flow_diff2).sum())
loss.backward()
optimizer.step()
train_loss.append(loss.data[0])
if len(train_loss) > 100:
train_loss.pop(0)
ave_train_loss = sum(train_loss) / float(len(train_loss))
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
if len(base_loss) > 100:
base_loss.pop(0)
ave_base_loss = sum(base_loss) / float(len(base_loss))
logging.info(
'epoch %d, train loss: %.2f, average train loss: %.2f, base loss: %.2f',
epoch, loss.data[0], ave_train_loss, ave_base_loss)
if (epoch + 1) % self.test_interval == 0:
logging.info('epoch %d, testing', epoch)
self.validate()
def test(self):
base_loss, test_loss = [], []
test_epe = []
motion = None
for epoch in range(self.test_epoch):
if self.data.name in ['box', 'mnist', 'box_complex']:
im, motion, _, _ = self.data.get_next_batch(
self.data.test_images)
elif self.data.name in ['box2', 'mnist2']:
im, motion, _ = self.data.get_next_batch(self.data.test_images)
elif self.data.name in [
'robot64', 'mpii64', 'mpi128', 'nyuv2', 'robot128',
'viper64', 'viper128', 'robot128c'
]:
im, motion = self.data.get_next_batch(
self.data.test_images), None
elif self.data.name in ['mpii64_sample']:
im, motion = self.data.get_next_batch(
self.data.test_images), None
im = im[:, -self.num_frame:, :, :, :]
else:
logging.error('%s data not supported' % self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float(),
volatile=True)
im_output = Variable(torch.from_numpy(im_output).float(),
volatile=True)
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
im_pred, flow = self.model(im_input, im_output)
flow = flow * self.im_size / 2 # resize flow from [-1, 1] back to image scale
im_diff = im_pred - im_input[:, -self.
im_channel:, :, :] # inverse warping loss
loss = torch.abs(im_diff).sum()
# loss = (im_diff * im_diff).sum()
flow_diff1 = flow[:, :, 1:, :] - flow[:, :, :-1, :]
flow_diff2 = flow[:, :, :, 1:] - flow[:, :, :, :-1]
loss = loss + 0.1 * (torch.abs(flow_diff1).sum() +
torch.abs(flow_diff2).sum())
# loss = loss + 0.01 * ((flow_diff1 * flow_diff1).sum() + (flow_diff2 * flow_diff2).sum())
test_loss.append(loss.data[0])
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
if motion is None:
gt_motion = None
else:
gt_motion = motion[:, -2, :, :, :]
gt_motion = Variable(torch.from_numpy(gt_motion).float())
if torch.cuda.is_available():
gt_motion = gt_motion.cuda()
epe = (flow - gt_motion) * (flow - gt_motion)
epe = torch.sqrt(epe.sum(1))
epe = epe.sum() / epe.numel()
test_epe.append(epe.cpu().data[0])
if self.display:
self.visualizer.visualize_result(im_input, im_output, im_pred,
flow, gt_motion,
'test_%d.png' % epoch)
if self.display_all:
for i in range(self.batch_size):
self.visualizer.visualize_result(im_input, im_output,
im_pred, flow, gt_motion,
'test_%d.png' % i, i)
test_loss = numpy.mean(numpy.asarray(test_loss))
base_loss = numpy.mean(numpy.asarray(base_loss))
improve_loss = base_loss - test_loss
improve_percent = improve_loss / (base_loss + 1e-5)
logging.info('average test loss: %.2f, base loss: %.2f', test_loss,
base_loss)
logging.info('improve_loss: %.2f, improve_percent: %.2f', improve_loss,
improve_percent)
if motion is not None:
test_epe = numpy.mean(numpy.asarray(test_epe))
logging.info('average test endpoint error: %.2f', test_epe)
return improve_percent
def test_gt(self):
base_loss, test_loss = [], []
test_epe = []
for epoch in range(self.test_epoch):
if self.data.name in ['box', 'mnist', 'box_complex']:
im, motion, _, _ = self.data.get_next_batch(
self.data.test_images)
elif self.data.name in ['box2', 'mnist2']:
im, motion, _ = self.data.get_next_batch(self.data.test_images)
else:
logging.error('%s data not supported in test_gt' %
self.data.name)
sys.exit()
im_input = im[:, :-1, :, :, :].reshape(self.batch_size, -1,
self.im_size, self.im_size)
im_output = im[:, -1, :, :, :]
gt_motion = motion[:, -2, :, :, :]
im_input = Variable(torch.from_numpy(im_input).float(),
volatile=True)
im_output = Variable(torch.from_numpy(im_output).float(),
volatile=True)
gt_motion = Variable(torch.from_numpy(gt_motion).float())
if torch.cuda.is_available():
im_input, im_output = im_input.cuda(), im_output.cuda()
gt_motion = gt_motion.cuda()
if self.data.name in ['box', 'mnist', 'box_complex']:
im_pred, flow = self.model_gt(im_input, im_output, gt_motion)
elif self.data.name in ['box2', 'mnist2']:
im_pred, flow = self.model_gt(im_input, im_output, gt_motion)
flow = flow * self.im_size / 2 # resize flow from [-1, 1] back to image scale
im_diff = im_pred - im_input[:, -self.
im_channel:, :, :] # inverse warping loss
loss = torch.abs(im_diff).sum()
flow_diff1 = flow[:, :, 1:, :] - flow[:, :, :-1, :]
flow_diff2 = flow[:, :, :, 1:] - flow[:, :, :, :-1]
loss = loss + 0.1 * (torch.abs(flow_diff1).sum() +
torch.abs(flow_diff2).sum())
test_loss.append(loss.data[0])
base_loss.append(
torch.abs(im_input[:, -self.im_channel:, :, :] -
im_output).sum().data[0])
epe = (flow - gt_motion) * (flow - gt_motion)
epe = torch.sqrt(epe.sum(1))
epe = epe.sum() / epe.numel()
test_epe.append(epe.cpu().data[0])
if self.display:
self.visualizer.visualize_result(im_input, im_output, im_pred,
flow, gt_motion,
'test_gt.png')
if self.display_all:
for i in range(self.batch_size):
self.visualizer.visualize_result(im_input, im_output,
im_pred, flow, gt_motion,
'test_gt_%d.png' % i, i)
test_loss = numpy.mean(numpy.asarray(test_loss))
base_loss = numpy.mean(numpy.asarray(base_loss))
improve_loss = base_loss - test_loss
improve_percent = improve_loss / (base_loss + 1e-5)
logging.info('average ground truth test loss: %.2f, base loss: %.2f',
test_loss, base_loss)
logging.info('improve_loss: %.2f, improve_percent: %.2f', improve_loss,
improve_percent)
test_epe = numpy.mean(numpy.asarray(test_epe))
logging.info('average ground truth test endpoint error: %.2f',
test_epe)
return improve_percent
def train(train_samples,
valid_samples,
word2num,
lr = 0.001,
epoch = 5,
use_cuda = False):
print('Training...')
# Prepare training data
print(' Preparing training data...')
statement_word2num = word2num[0]
subject_word2num = word2num[1]
speaker_word2num = word2num[2]
speaker_pos_word2num = word2num[3]
state_word2num = word2num[4]
party_word2num = word2num[5]
context_word2num = word2num[6]
train_data = train_samples
dataset_to_variable(train_data, use_cuda)
valid_data = valid_samples
dataset_to_variable(valid_data, use_cuda)
# Construct model instance
print(' Constructing network model...')
model = Net(len(statement_word2num),
len(subject_word2num),
len(speaker_word2num),
len(speaker_pos_word2num),
len(state_word2num),
len(party_word2num),
len(context_word2num))
if use_cuda: model.cuda()
# Start training
print(' Start training')
optimizer = optim.Adam(model.parameters(), lr = lr)
model.train()
step = 0
display_interval = 2000
for epoch_ in range(epoch):
print(' ==> Epoch '+str(epoch_)+' started.')
random.shuffle(train_data)
total_loss = 0
for sample in train_data:
optimizer.zero_grad()
prediction = model(sample)
label = Variable(torch.LongTensor([sample.label]))
loss = F.cross_entropy(prediction, label)
loss.backward()
optimizer.step()
step += 1
if step % display_interval == 0:
print(' ==> Iter: '+str(step)+' Loss: '+str(loss))
total_loss += loss.data.numpy()
print(' ==> Epoch '+str(epoch_)+' finished. Avg Loss: '+str(total_loss/len(train_data)))
valid(valid_data, word2num, model)
return model
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
train_set = get_training_set(opt.upscale_factor)
test_set = get_test_set(opt.upscale_factor)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model')
model = Net(upscale_factor=opt.upscale_factor)
criterion = nn.MSELoss()
if cuda:
model = model.cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
def train(epoch):
epoch_loss = 0
for iteration, batch in enumerate(training_data_loader, 1):
input, target = Variable(batch[0]), Variable(batch[1])
if cuda:
input = input.cuda()
target = target.cuda()
optimizer.zero_grad()
loss = criterion(model(input), target)
