def get_data_loaders(self):
root_dir = self.config['dataset_root_dir']
# get labeled dataset
labeled_set = self.config['labeled_set']
self.train_lab_dataset = PickleDataset(os.path.join(
root_dir, f'{labeled_set}.pkl'),
config=self.config,
sort=True)
self.train_lab_loader = get_data_loader(
self.train_lab_dataset,
batch_size=self.config['batch_size'],
shuffle=self.config['shuffle'],
drop_last=False)
# get unlabeled dataset
unlabeled_set = self.config['unlabeled_set']
self.train_unlab_dataset = PickleDataset(os.path.join(
root_dir, f'{unlabeled_set}.pkl'),
config=self.config,
sort=True)
self.train_unlab_loader = get_data_loader(
self.train_unlab_dataset,
batch_size=self.config['batch_size'],
shuffle=self.config['shuffle'],
drop_last=False,
speech_only=True)
# get dev dataset
clean_dev_set = self.config['clean_dev_set']
# do not sort dev set
self.clean_dev_dataset = PickleDataset(os.path.join(
root_dir, f'{clean_dev_set}.pkl'),
sort=True)
self.clean_dev_loader = get_data_loader(
self.clean_dev_dataset,
batch_size=self.config['batch_size'] // 2,
shuffle=False,
drop_last=False)
# get dev dataset
noisy_dev_set = self.config['noisy_dev_set']
# do not sort dev set
self.noisy_dev_dataset = PickleDataset(os.path.join(
root_dir, f'{noisy_dev_set}.pkl'),
sort=True)
self.noisy_dev_loader = get_data_loader(
self.noisy_dev_dataset,
batch_size=self.config['batch_size'] // 2,
shuffle=False,
drop_last=False)
return
def train(args):
full_data = get_data_loader(args)
VAE_human = VAE(512).to(device)
VAE_cartoon = VAE(512).to(device)
optimiser_human = optim.Adam(VAE_human.parameters(), lr=0.0002)
optimiser_cartoon = optim.Adam(VAE_cartoon.parameters(), lr=0.0002)
VAE_human.train()
VAE_cartoon.train()
print("Start Training....")
for epoch in trange(args.num_epochs):
total_VAE_human_loss = 0.0
total_VAE_cartoon_loss = 0.0
total_data = 0
for batch_num, data in enumerate(full_data):
human, cartoon = data[0].to(device), data[1].to(
device) # x is cartoon, y is human
total_data += human.shape[0]
total_VAE_human_loss += train_VAE_1_step(VAE_human, VAE_cartoon,
optimiser_human, human)
total_VAE_cartoon_loss += train_VAE_1_step(VAE_cartoon, VAE_human,
optimiser_cartoon,
cartoon)
avg_VAE_human_loss = total_VAE_human_loss / total_data
avg_VAE_cartoon_loss = total_VAE_cartoon_loss / total_data
print("Avg VAE Cartoon Loss: {}".format(avg_VAE_cartoon_loss))
print("Avg VAE Human Loss: {}".format(avg_VAE_human_loss))
def main(opt):
logger.info('Loading model: %s', opt.model_file)
checkpoint = torch.load(opt.model_file)
checkpoint_opt = checkpoint['opt']
# Update/Overwrite some test options like batch size, location to metadata
# file
vars(checkpoint_opt).update(vars(opt))
logger.info('Updated input arguments: %s',
json.dumps(vars(checkpoint_opt), sort_keys=True, indent=4))
logger.info('Building model...')
model = get_model(checkpoint_opt, num_classes=checkpoint_opt.num_classes)
test_loader = get_data_loader(checkpoint_opt,
training=False,
return_org_image=True,
data_list=opt.test_data_list)
logger.info('Loading model parameters...')
model = DataParallelModel(model)
model.load_state_dict(checkpoint['model'])
if torch.cuda.is_available():
model.cuda()
logger.info('Start testing...')
test(checkpoint_opt, model, test_loader)
def test(args, device):
full_data = get_data_loader(args)
if args.model_type == "CNN":
from CNN import CNN
model = CNN(args).to(device)
elif args.model_type == "MLP":
from MLP import MLP
model = MLP(args).to(device)
elif args.model_type == "LSTM":
from LSTM import LSTM
model = LSTM(args).to(device)
optimiser = optim.Adam(
model.parameters(), lr=args.learning_rate)
state = torch.load(args.model_path, map_location=device)
model.load_state_dict(state['model'])
optimiser.load_state_dict(state['optimiser'])
total_difference = 0
n = 0
for batch_num, data in enumerate(full_data):
x, y = data[0].float().to(device), data[1].float().to(device)
num_of_predictions = x.shape[0]
pred = model(x)
pred = pred.reshape(y.shape)
total_difference += sum((abs(pred - y)/y) * 100)
n += num_of_predictions
return total_difference/n
def main_global(args):
data_dir = args.data_dir
params = {'batch_size': args.batch, 'shuffle': False}
if args.bert_fts:
type_dir = "all_bertemb/"
else:
type_dir = "all/"
data_dir_back = ""
if (args.trainon == 'bothway') or (args.trainon == 'bothWselect'):
if args.bert_fts:
data_dir_back = args.data_dir + "all_backward_bertemb/"
else:
data_dir_back = args.data_dir + "all_backward/"
train_data = EventDataset(args.data_dir + type_dir, "train",
args.glove2vocab, data_dir_back, args.bert_fts)
print('train_data: %s in total' % len(train_data))
train_generator = get_data_loader(train_data, **params)
dev_data = EventDataset(args.data_dir + type_dir, "dev", args.glove2vocab,
data_dir_back, args.bert_fts)
print('dev_data: %s in total' % len(dev_data))
dev_generator = get_data_loader(dev_data, **params)
if args.bert_fts:
data_dir_back = args.data_dir + "all_backward_bertemb/"
else:
data_dir_back = args.data_dir + "all_backward/"
test_data = EventDataset(args.data_dir + type_dir, "test",
args.glove2vocab, data_dir_back, args.bert_fts)
test_generator = get_data_loader(test_data, **params)
s_time = time.time()
models = [NNClassifier()]
score = 0
for model in models:
dev_f1 = model.train_epoch(train_generator, dev_generator, args)
print('total time escape', time.time() - s_time)
evaluator = Evaluator(model)
#print(evaluator.evaluate(test_generator, args))
score = evaluator.get_score(test_generator, args)
#evaluator.collect_result(test_generator, args)
print('final test f1: %.4f' % (score))
return float(dev_f1), float(score)
def parallel_cv(self, split, emb=np.array([]), pos_emb=[], args=None):
params = {'batch_size': args.batch, 'shuffle': False}
if args.bert_fts:
type_dir = "cv_bertemb"
else:
type_dir = "cv_shuffle" if args.cv_shuffle else 'cv'
backward_dir = ""
if (args.trainon == 'bothway') or (args.trainon == 'bothWselect'):
if args.bert_fts:
backward_dir = "%scv_backward_bertemb/fold%s/" % (
args.data_dir, split)
else:
backward_dir = "%scv_backward/fold%s/" % (args.data_dir, split)
train_data = EventDataset(
args.data_dir + '%s/fold%s/' % (type_dir, split), "train",
args.glove2vocab, backward_dir, args.bert_fts)
train_generator = get_data_loader(train_data, **params)
dev_data = EventDataset(
args.data_dir + '%s/fold%s/' % (type_dir, split), "dev",
args.glove2vocab, backward_dir, args.bert_fts)
dev_generator = get_data_loader(dev_data, **params)
seeds = [0, 10, 20]
accumu_f1 = 0.
accumu_epoch = 0.
for seed in seeds:
exec("args.%s=%s" % ('seed', seed))
f1, epoch = self._train(train_generator,
dev_generator,
emb,
pos_emb,
args,
in_cv=True)
accumu_f1 += f1
accumu_epoch += epoch
avg_f1 = accumu_f1 / float(len(seeds))
avg_epoch = accumu_epoch / float(len(seeds))
return avg_f1, avg_epoch
def __init__(self, args):
super(Trainer, self).__init__()
self.epoch = args.epoch
self.batch_size = args.batch_size
self.data_dir = args.data_dir
self.save_dir = args.save_dir
self.result_dir = args.result_dir
self.log_dir = args.log_dir
self.gpu_mode = args.gpu_mode
self.verbose = args.verbose
if args.model == 'fcn16s':
self.model = FCN16s()
elif args.model == 'fcn32s':
self.model = FCN32s()
elif args.model == 'fcn8s':
self.model = FCN8s()
elif args.model == 'pspnet':
self.model = PSPnet()
else:
print("No this model type")
exit(-1)
if self.gpu_mode:
self.model = self.model.cuda()
self.parameter = self.model.parameters()
self.optimizer = optim.Adam(self.parameter, lr=args.learning_rate)
self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer,
gamma=0.5)
self.train_dataloader = get_data_loader(self.data_dir,
self.batch_size,
split='train')
self.test_dataloader = get_data_loader(self.data_dir, 1, split='val')
# experiment_id = args.model + time.strftime('%m%d%H%m')
# self.writer = SummaryWriter(log_dir=self.log_dir + '/tboard_' + experiment_id)
self.loss = FocalLoss(gamma=1.25)
if args.pretrain != '':
self._load_pretrain(args.pretrain)
def main(opt):
logger.info('Loading model: %s', opt.model_file)
checkpoint = torch.load(opt.model_file)
checkpoint_opt = checkpoint['opt']
# Load model location
model = LaneNet(cnn_type=checkpoint_opt.cnn_type)
model = DataParallelModel(model)
# Update/Overwrite some test options like batch size, location to metadata
# file
vars(checkpoint_opt).update(vars(opt))
test_loader = get_data_loader(checkpoint_opt,
split='test',
return_org_image=True)
logger.info('Building model...')
model.load_state_dict(checkpoint['model'])
if torch.cuda.is_available():
model = model.cuda()
postprocessor = PostProcessor()
clustering = LaneClustering()
logger.info('Start testing...')
if opt.loader_type == 'tusimpletest':
x_lanes, _, times, _ = test(model,
test_loader,
postprocessor,
clustering,
genline_method=opt.genline_method)
output_tuprediction(opt.meta_file, x_lanes, times, opt.output_file)
if opt.loader_type == 'culanetest':
x_lanes, y_list, _, image_files = test(
model,
test_loader,
postprocessor,
clustering,
genline_method=opt.genline_method)
output_culaneprediction(opt.output_dir, x_lanes, y_list, image_files)
if opt.loader_type == 'dirloader':
visualize(model,
test_loader,
postprocessor,
clustering,
show_demo=opt.show_demo,
output_dir=opt.output_dir,
genline_method=opt.genline_method)
def interpolate(args):
model, z_dim = build_model(args.model)
wrapper = Wrapper(args, model, z_dim)
# load the target and mask
loader = get_data_loader(args.input)
criterion = Criterion(args)
for idx, (image, _) in enumerate(loader):
save_images(image, 'output/interpolate/%d' % (idx))
target = image.to(device)
param = sample_noise(z_dim, device, args.latent, model, from_mean=True)
param, recon = optimize_para(wrapper, param, target, criterion, args.n_iters)
save_images(recon, 'output/interpolate/%d_%s_%s' % (idx, args.model, args.latent))
if idx % 2 == 0:
src = param
continue
dst = param
image_list = []
with torch.no_grad():
# todo: interpolation code
save_gifs(image_list, 'output/interpolate/%d_%s_%s' % (idx, args.model, args.latent))
if idx >= 3:
break
return
def parse_arg():
"""Creates a parser for command-line arguments.
"""
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='stylegan', choices=['vanilla', 'stylegan'])
parser.add_argument('--mode', type=str, default='sample', choices=['sample', 'project', 'draw', 'interpolate'])
parser.add_argument('--latent', type=str, default='z', choices=['z', 'w', 'w+'])
parser.add_argument('--n_iters', type=int, default=1000, help="number of optimization steps in the image projection")
parser.add_argument('--perc_wgt', type=float, default=0., help="perc loss lambda")
parser.add_argument('--input', type=str, default='data/cat/*.png', help="path to the input image")
return parser.parse_args()
if __name__ == '__main__':
args = parse_arg()
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
if args.mode == 'sample':
sample(args)
elif args.mode == 'project':
project(args)
elif args.mode == 'draw':
draw(args)
elif args.mode == 'interpolate':
interpolate(args)
def draw(args):
# define and load the pre-trained model
model, z_dim = build_model(args.model)
wrapper = Wrapper(args, model, z_dim)
# load the target and mask
loader = get_data_loader(args.input, alpha=True)
criterion = Criterion(args, True)
for idx, (rgb, mask) in enumerate(loader):
rgb, mask = rgb.to(device), mask.to(device)
save_images(rgb, 'output/draw/%d_data' % idx, 1)
save_images(mask, 'output/draw/%d_mask' % idx, 1)
def get_z_accuracy_data(self, state_dict=None):
# this function will collect the representation of the testing dataset,
# for z_accuracy_test, you need to put the "training data" in the "
# test_set" config so that able to go further test
if not state_dict:
self.load_model(self.config['load_model_path'],
self.config['load_optimizer'])
else:
self.model.load_state_dict(state_dict)
# get test dataset
root_dir = self.config['dataset_root_dir']
test_set = self.config['test_set']
test_file_name = self.config['test_file_name']
test_dataset = PickleDataset(os.path.join(root_dir, f'{test_set}.p'),
config=None,
sort=False)
test_loader = get_data_loader(test_dataset,
batch_size=self.config['batch_size'],
shuffle=False)
self.model.eval()
stored_data = dict()
for step, data in enumerate(test_loader):
bos = self.vocab['<BOS>']
eos = self.vocab['<EOS>']
pad = self.vocab['<PAD>']
xs, ilens, ys, ys_in, ys_out, spks, envs, trans = to_gpu(
data, bos, eos, pad)
# feed previous
_, prediction, _, enc_outputs, enc_lens = self.model(
xs,
ilens,
None,
max_dec_timesteps=self.config['max_dec_timesteps'])
representations = enc_outputs.cpu().detach()
speakers = spks.cpu().tolist()
environments = envs.cpu().tolist()
for instance in range(enc_outputs.size(0)):
stored_data[((step * self.config['batch_size']) +
instance)] = {
'embedding':
trim_representation(representations[instance],
enc_lens[instance]),
'speaker':
speakers[instance],
'env':
environments[instance],
'transcripts':
trans[instance]
}
pickle.dump(stored_data, open(self.config['z_data_path'], 'wb'))
self.model.train()
return
def main():
global args, best_prec1
args = parser.parse_args()
# Check the save_dir exists or not
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
#Some hard-code setting for fast experiments
if args.train_dataset == 'cifar10':
num_classes = 10
elif args.train_dataset == 'cifar100':
num_classes = 100
else:
print("undefined num_classes")
if args.arch.startswith('dense'):
size = args.size_dense
elif args.arch.startswith('wide'):
size = args.size_wide
elif args.arch.startswith('vgg'):
size = 0
else:
raise NotImplementedError
print(num_classes)
model = torch.nn.DataParallel(models.__dict__[args.arch](num_classes,
size))
model.cuda()
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
#Get data
train_loader, val_loader = get_data_loader(args.infer_dataset)
# define loss function (criterion) and optimizer
criterionList = get_criterion_list(args.arch)
validate_and_save(val_loader, model, criterionList, args, num_classes)
def test(self, state_dict=None):
# load model
if not state_dict:
self.load_model(self.config['load_model_path'],
self.config['load_optimizer'])
else:
self.model.load_state_dict(state_dict)
# get test dataset
root_dir = self.config['dataset_root_dir']
test_set = self.config['test_set']
test_dataset = PickleDataset(os.path.join(root_dir, f'{test_set}.pkl'),
config=None,
sort=False)
test_loader = get_data_loader(test_dataset,
batch_size=1,
shuffle=False,
drop_last=False)
self.model.eval()
all_prediction, all_ys = [], []
for step, data in enumerate(test_loader):
xs, ilens, ys = to_gpu(data)
# feed previous
(_, _, prediction, _), _ = self.model(
xs,
ilens,
ys=None,
max_dec_timesteps=self.config['max_dec_timesteps'])
all_prediction = all_prediction + prediction.cpu().numpy().tolist()
all_ys = all_ys + [y.cpu().numpy().tolist() for y in ys]
self.model.train()
cer, prediction_sents, ground_truth_sents = self.ind2sent(
all_prediction, all_ys)
with open(f'{test_set}.txt', 'w') as f:
for p in prediction_sents:
f.write(f'{p}\n')
print(f'{test_set}: {len(prediction_sents)} utterances, CER={cer:.4f}')
return cer
def get_data_loaders(self):
root_dir = self.config['dataset_root_dir']
# get labeled dataset
labeled_set = self.config['labeled_set']
self.train_lab_dataset = PickleDataset(os.path.join(
root_dir, f'{labeled_set}.p'),
config=self.config,
sort=True)
self.train_lab_loader = get_data_loader(
self.train_lab_dataset,
batch_size=self.config['batch_size'],
shuffle=self.config['shuffle'])
# get dev dataset
dev_set = self.config['dev_set']
# do not sort dev set
self.dev_dataset = PickleDataset(os.path.join(root_dir,
f'{dev_set}.p'),
sort=True)
self.dev_loader = get_data_loader(self.dev_dataset,
batch_size=self.config['batch_size'],
shuffle=False)
return
def main_local(args):
data_dir = args.data_dir
params = {'batch_size': args.batch, 'shuffle': False}
if args.bert_fts:
type_dir = "all_bertemb/"
else:
type_dir = "all/"
data_dir_back = ""
if (args.trainon == 'bothway') or (args.trainon == 'bothWselect'):
if args.bert_fts:
data_dir_back = args.data_dir + "all_backward_bertemb/"
else:
data_dir_back = args.data_dir + "all_backward/"
train_data = EventDataset(args.data_dir + type_dir, "train",
args.glove2vocab, data_dir_back, args.bert_fts)
print('total train_data %s samples' % len(train_data))
train_generator = get_data_loader(train_data, **params)
dev_data = EventDataset(args.data_dir + type_dir, "dev", args.glove2vocab,
data_dir_back, args.bert_fts)
print('total dev_data %s samples' % len(dev_data))
dev_generator = get_data_loader(dev_data, **params)
if args.bert_fts:
data_dir_back = args.data_dir + "all_backward_bertemb/"
else:
data_dir_back = args.data_dir + "all_backward/"
test_data = EventDataset(args.data_dir + type_dir, "test",
args.glove2vocab, data_dir_back, args.bert_fts)
test_generator = get_data_loader(test_data, **params)
models = [NNClassifier()]
for model in models:
dev_f1 = model.train_epoch(train_generator, dev_generator, args)
evaluator = Evaluator(model)
#evaluator.for_analysis(test_generator, args)
score = evaluator.get_score(test_generator, args)
return float(dev_f1), float(score)
def __init__(self, model, conf):
stuple = get_shape(conf.dataset)
self.input_shape = (conf.batch_size, stuple[0], stuple[1], stuple[1])
self.model = model
if torch.cuda.is_available():
self.model = self.model.cuda()
self.opt = conf
self.default_root = os.getcwd()
self.sample_path =os.path.join(self.default_root, conf.sample_path)
if self.opt.mode == 'train' or self.opt.mode == 'mix':
train_tfs = get_transformer(self.opt.dataset, True, crop_size=self.opt.crop_size,
image_size=self.opt.image_size)
self.train_dl = get_data_loader(self.opt.dataset, os.path.join(self.default_root, self.opt.root),
self.opt.batch_size, train_tfs, True)
test_tfs = get_transformer(self.opt.dataset, False, crop_size=None, image_size=self.opt.image_size)
self.test_dl = get_data_loader(self.opt.dataset, os.path.join(self.default_root, self.opt.root), self.opt.batch_size, test_tfs, False)
self.criterion = nn.CrossEntropyLoss()
self.optimizer = Adam(self.model.parameters(), lr=self.opt.lr, betas=(self.opt.beta1, self.opt.beta2))
self.num_classes = 1000 if self.opt.dataset == 'img_folder' else 10
def run_one_epoch(self, epoch):
'''
Train one epoch of the whole dataset.
Args:
- dataset: class Dataset.
- otuput_probs: if True, model output log probs, in which case we also calculate accuracy
'''
self._model.train()
self.logger.info('=> Training epoch %d' % epoch)
data_iter = get_data_loader(
self.dataset,
"train",
self.config['Trainer']['train_batch_sz'],
max_len=self.config['Model']
['max_len'], # upper bound of input sentence length
max_sample=self.config['Trainer']
['total_samples'], # total number of samples to train
n_workers=self.config['Trainer']['n_workers'])
kwargs = self.forward_args()
losses, accuracies = 0, 0
for step, (data, data_lens) in enumerate(data_iter):
self._opt.zero_grad()
loss, acc, kwargs = self.forward_pass(data, data_lens, kwargs)
self.logger.debug('loss per batch = %f' % loss)
losses += loss.detach().item()
accuracies += acc
nn.utils.clip_grad_norm_(self._model.parameters(),
2) # gradient clipping
loss.backward()
# debug: print gradients
grad_of_param = {}
for name, parameter in self._model.named_parameters():
grad_of_param[name] = parameter.grad
#self.logger.debug('gradient of %s: \n%s'%(name, str(parameter.grad)))
self._opt.step()
loss_per_epoch = losses / (step + 1)
acc_per_epoch = accuracies / (step + 1)
if math.isnan(loss_per_epoch):
self.logger.error("Get NaN loss for epoch %d-- exiting" % epoch)
sys.exit(1)
return loss_per_epoch, acc_per_epoch
def train(configs, args):
if args.load:
model = torch.load(args.load)
else:
model = importlib.import_module("model.{}".format(args.model))
model = model.CAE()
bind_model(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.DataParallel(model).to(device)
criterion = nn.BCEWithLogitsLoss().to(device)
optimizer = optim.Adam(model.parameters(),
lr=configs['LEARNING_RATE'], weight_decay=configs['WEIGHT_DECAY'])
train_dataset, val_dataset = get_data_loader(DATASET_PATH, configs['BS'], configs['RATIO'], configs['VMAX'])
train_size = len(train_dataset)
val_size = len(val_dataset)
for epoch in range(int(args.start) if args.start else 0, configs['ITER']):
avg_loss = 0.0
avg_eval_loss = 0.0
for i, data in enumerate(train_dataset):
loss = train_iter(model, data, criterion, optimizer, args.target, args.input)
avg_loss += loss
print('Epoch :', epoch, ', Batch : ', i + 1, '/', train_size, ', ERROR in this minibatch: ', loss)
if epoch % configs['SAVE_EVERY'] == 0 and epoch != 0:
nsml.save(epoch)
with torch.no_grad():
for i, data in enumerate(val_dataset):
loss = eval_iter(model, data, criterion, args.target, args.input)
avg_eval_loss += loss
print('Batch : ', i + 1, '/', val_size, ', Validation ERROR in this minibatch: ', loss)
print('epoch:', epoch, ' train_loss:', float(avg_loss / train_size))
# print('epoch:', epoch, ' Average CER:', float(avg_cer/total_batch))
print('epoch:', epoch, ' Validation_loss:', float(avg_eval_loss / val_size))
def main():
use_cuda = torch.cuda.is_available() and args.cuda
device = torch.device('cuda' if use_cuda else 'cpu')
print('Loading model ...')
model = Net(device)
if args.load is not None:
print('Loading checkpoint ...')
model.load_state_dict(torch.load(args.load))
if use_cuda:
model.cuda()
print('Loading data ...')
data_loader = get_data_loader(args.data_root, args.stage, args.batch_size)
print('Preparation done')
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999))
criterion = nn.CrossEntropyLoss()
if args.stage == 'train':
train(model, data_loader, optimizer, criterion)
def project(args):
# load images
loader = get_data_loader(args.input, is_train=False)
# define and load the pre-trained model
model, z_dim = build_model(args.model)
wrapper = Wrapper(args, model, z_dim)
print('model {} loaded'.format(args.model))
# todo: implement your criterion here.
criterion = Criterion(args)
# project each image
for idx, (data, _) in enumerate(loader):
target = data.to(device)
save_images(data, 'output/project/%d_data' % idx, 1)
param = sample_noise(z_dim, device, args.latent, model)
optimize_para(wrapper, param, target, criterion, args.n_iters,
'output/project/%d_%s_%s_%g' % (idx, args.model, args.latent, args.perc_wgt))
if idx >= 0:
break
def draw(args):
# define and load the pre-trained model
model, z_dim = build_model(args.model)
wrapper = Wrapper(args, model, z_dim)
# load the target and mask
loader = get_data_loader(args.input, alpha=True)
criterion = Criterion(args, apply_mask=True)
for idx, (rgb, mask) in enumerate(loader):
rgb, mask = rgb.to(device), mask.to(device)
save_images(rgb, 'output/draw/%d_data' % idx, 1)
save_images(mask, 'output/draw/%d_mask' % idx, 1)
# todo: optimize sketch 2 image
param = sample_noise(z_dim, device, args.latent, args.model, model)
optimize_para(
wrapper, param, (rgb, mask), criterion, args.n_iters,
'output/draw/%d_%s_%s_%g' %
(idx, args.model, args.latent, args.perc_wgt))
if idx >= 0:
break
def run_eval(config):
val_meta = read_meta(data_path=config['data_dir'], meta_type='valid')
cat = 'all'
if 'cat' in config:
cat = config['cat']
val_loader = get_data_loader(config['data_dir'], val_meta,
shuffle=False, rgb=config['rgb'],
transform=config['transform'], cat=cat)
model = config['model']
state_dict = torch.load(config['model_path'],
map_location=(None if config['use_gpu'] else 'cpu'))
model.load_state_dict(state_dict)
if config['use_gpu']:
model = model.cuda()
evaluate(model, val_loader, config['n_round'],
config['use_gpu'], config['im_type'])
def validate(self, ds_name='dev'):
self._model.eval()
losses, accuracies = 0, 0
data_iter_eval = get_data_loader(
self.dataset,
ds_name,
self.config['Trainer']['eval_batch_sz'],
max_len=self.config['Model']
['max_len'], # upper bound of input sentence length
n_workers=self.config['Trainer']['n_workers'])
kwargs = self.forward_args()
kwargs['output_probs'] = True
with torch.no_grad():
for step, (data, data_lens) in enumerate(data_iter_eval):
loss, acc, kwargs = self.forward_pass(data, data_lens, kwargs)
accuracies += acc
losses += loss.detach().item()
loss_per_epoch = losses / (step + 1)
return loss_per_epoch, accuracies / (step + 1)
def main(opt):
logger.info('Loading model: %s', opt.model_file)
test_opt = {
'label_file': opt.test_label,
'imageinfo_file': opt.test_imageinfo,
'image_dir': opt.test_image_dir,
'batch_size': opt.batch_size,
'num_workers': opt.num_workers,
'train': False
}
checkpoint = torch.load(opt.model_file)
test_loader = get_data_loader(test_opt)
num_labels = test_loader.dataset.get_num_labels()
logger.info('Building model...')
checkpoint_opt = checkpoint['opt']
model = DepNet(
num_labels,
finetune=checkpoint_opt.finetune,
cnn_type=checkpoint_opt.cnn_type,
pretrained=False)
criterion = nn.MultiLabelSoftMarginLoss()
model.load_state_dict(checkpoint['model'])
if torch.cuda.is_available():
model.cuda()
criterion.cuda()
logger.info('Start testing...')
test_loss, test_score = test(checkpoint_opt, model, criterion, test_loader)
logger.info('Test loss: \n%s', test_loss)
logger.info('Test score: \n%s', test_score)
out = {'map': test_score.map()}
logger.info('Writing output to %s', opt.output_file)
with open(opt.output_file, 'w') as f:
json.dump(out, f)
def interpolate(args):
model, z_dim = build_model(args.model)
wrapper = Wrapper(args, model, z_dim)
# load the target and mask
loader = get_data_loader(args.input)
criterion = Criterion(args)
for idx, (image, _) in enumerate(loader):
save_images(image, 'output/interpolate/%d' % (idx))
target = image.to(device)
param = sample_noise(z_dim,
device,
args.latent,
args.model,
model,
from_mean=True)
param, recon = optimize_para(wrapper, param, target, criterion,
args.n_iters)
save_images(
recon,
'output/interpolate/%d_%s_%s' % (idx, args.model, args.latent))
if idx % 2 == 0:
src = param
continue
dst = param
image_list = []
with torch.no_grad():
# todo: interpolation code
for i in range(20):
inter_param = src * (i * 0.05) + dst * (1 - i * 0.05)
image_list.append(wrapper(inter_param))
save_gifs(
image_list,
'output/interpolate/%d_%s_%s' % (idx, args.model, args.latent))
if idx >= 3:
break
return
def main(config):
# use cuda ?
if config.use_cuda:
from torch.backends import cudnn
cudnn.benchmark = True
elif torch.cuda.is_available():
print("WARNING: You have a CUDA device, so you should probably run with --cuda")
# seed
if config.seed == 0:
config.seed = random.randint(1, 10000) # fix seed
print("Random Seed: ", config.seed)
random.seed(config.seed)
torch.manual_seed(config.seed)
if config.use_cuda:
torch.cuda.manual_seed_all(config.seed)
# create directories if not exist
if not os.path.exists(config.out_path):
os.makedirs(config.out_path)
# dataLoader
trainLoader, testLoader = get_data_loader(config)
print('train samples num: ', len(trainLoader), ' test samples num: ', len(testLoader))
# model net
model = RetinaNet(num_classes=config.n_classes)
print(model)
if config.pretrained != '':
model.load_state_dict(torch.load(config.pretrained))
print('load', config.pretrained)
solver = Solver(config, model, trainLoader, testLoader)
for epoch in range(config.n_epochs):
solver.train(epoch)
solver.test(epoch)
def main(opt):
# Set the random seed manually for reproducibility.
if torch.cuda.is_available():
torch.cuda.manual_seed(opt.seed)
else:
torch.manual_seed(opt.seed)
train_loader = get_data_loader(opt,
split='train',
return_org_image=False)
val_loader = get_data_loader(opt,
split='val',
return_org_image=False)
output_dir = os.path.dirname(opt.output_file)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
logger.info('Building model...')
model = LaneNet(cnn_type=opt.cnn_type, embed_dim=opt.embed_dim)
model = DataParallelModel(model)
criterion_disc = DiscriminativeLoss(delta_var=0.5,
delta_dist=1.5,
norm=2,
usegpu=True)
criterion_ce = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate)
if opt.start_from:
logger.info('Restart training from %s', opt.start_from)
checkpoint = torch.load(opt.start_from)
model.load_state_dict(checkpoint['model'])
if torch.cuda.is_available():
criterion_disc.cuda()
criterion_ce.cuda()
model = model.cuda()
logger.info("Start training...")
best_loss = sys.maxsize
best_epoch = 0
for epoch in tqdm(range(opt.num_epochs), desc='Epoch: '):
learning_rate = adjust_learning_rate(opt, optimizer, epoch)
logger.info('===> Learning rate: %f: ', learning_rate)
# train for one epoch
train(
opt,
model,
criterion_disc,
criterion_ce,
optimizer,
train_loader)
# validate at every val_step epoch
if epoch % opt.val_step == 0:
val_loss = test(
opt,
model,
criterion_disc,
criterion_ce,
val_loader)
logger.info('Val loss: %s\n', val_loss)
loss = val_loss.avg
if loss < best_loss:
logger.info(
'Found new best loss: %.7f, previous loss: %.7f',
loss,
best_loss)
best_loss = loss
best_epoch = epoch
logger.info('Saving new checkpoint to: %s', opt.output_file)
torch.save({
'epoch': epoch,
'model': model.state_dict(),
'best_loss': best_loss,
'best_epoch': best_epoch,
'opt': opt
}, opt.output_file)
else:
logger.info(
'Current loss: %.7f, best loss is %.7f @ epoch %d',
loss,
best_loss,
best_epoch)
if epoch - best_epoch > opt.max_patience:
logger.info('Terminated by early stopping!')
break
def main():
global args, best_prec1, device
args = parser.parse_args()
batch_size = args.batch_size * max(1, args.num_gpus)
args.lr = args.lr * (batch_size / 256.)
print(batch_size, args.lr, args.num_gpus)
num_classes = 1000
num_training_samples = 1281167
args.num_batches_per_epoch = num_training_samples // batch_size
assert os.path.isfile(args.load) and args.load.endswith(".pth.tar")
args.save = os.path.dirname(args.load)
training_mode = 'retrain' if args.retrain else 'finetune'
args.save = os.path.join(args.save, training_mode)
if not os.path.exists(args.save):
os.makedirs(args.save)
args.model_save_path = os.path.join(
args.save, "epochs_{}_{}".format(args.epochs,
os.path.basename(args.load)))
args.distributed = args.world_size > 1
##########################################################
## create file handler which logs even debug messages
#import logging
#log = logging.getLogger()
#log.setLevel(logging.INFO)
#ch = logging.StreamHandler()
#fh = logging.FileHandler(args.logging_file_path)
#formatter = logging.Formatter('%(asctime)s - %(message)s')
#ch.setFormatter(formatter)
#fh.setFormatter(formatter)
#log.addHandler(fh)
#log.addHandler(ch)
##########################################################
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# Use CUDA
args.use_cuda = torch.cuda.is_available() and not args.no_cuda
# Random seed
random.seed(0)
torch.manual_seed(0)
if args.use_cuda:
torch.cuda.manual_seed_all(0)
device = 'cuda'
cudnn.benchmark = True
else:
device = 'cpu'
if args.evaluate == 1:
device = 'cuda:0'
assert os.path.isfile(args.load)
print("=> loading checkpoint '{}'".format(args.load))
checkpoint = torch.load(args.load)
model = mobilenetv2(cfg=checkpoint['cfg'])
cfg = model.cfg
total_params = print_model_param_nums(model.cpu())
total_flops = print_model_param_flops(model.cpu(),
224,
multiply_adds=False)
print(total_params, total_flops)
if not args.distributed:
model = torch.nn.DataParallel(model).to(device)
else:
model.to(device)
model = torch.nn.parallel.DistributedDataParallel(model)
##### finetune #####
if not args.retrain:
model.load_state_dict(checkpoint['state_dict'])
# define loss function (criterion) and optimizer
if args.label_smoothing:
criterion = CrossEntropyLabelSmooth(num_classes).to(device)
else:
criterion = nn.CrossEntropyLoss().to(device)
### all parameter ####
no_wd_params, wd_params = [], []
for name, param in model.named_parameters():
if param.requires_grad:
if ".bn" in name or '.bias' in name:
no_wd_params.append(param)
else:
wd_params.append(param)
no_wd_params = nn.ParameterList(no_wd_params)
wd_params = nn.ParameterList(wd_params)
optimizer = torch.optim.SGD([
{
'params': no_wd_params,
'weight_decay': 0.
},
{
'params': wd_params,
'weight_decay': args.weight_decay
},
],
args.lr,
momentum=args.momentum)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.model_save_path):
print("=> loading checkpoint '{}'".format(args.model_save_path))
checkpoint = torch.load(args.model_save_path)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.model_save_path, checkpoint['epoch']))
else:
pass
# Data loading code
train_loader, val_loader = \
get_data_loader(args.data, train_batch_size=batch_size, test_batch_size=32, workers=args.workers)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
#adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'cfg': cfg,
#'m': args.m,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
},
args.model_save_path)
print(' + Number of params: %.3fM' % (total_params / 1e6))
print(' + Number of FLOPs: %.3fG' % (total_flops / 1e9))
def test(self, state_dict=None):
# load model
if not state_dict:
self.load_model(self.config['load_model_path'],
self.config['load_optimizer'])
else:
self.model.load_state_dict(state_dict)
# get test dataset
root_dir = self.config['dataset_root_dir']
test_set = self.config['test_set']
test_file_name = self.config['test_file_name']
test_dataset = PickleDataset(os.path.join(root_dir, f'{test_set}.p'),
config=None,
sort=False)
test_loader = get_data_loader(test_dataset,
batch_size=2,
shuffle=False)
self.model.eval()
all_prediction, all_ys = [], []
gold_transcripts = []
for step, data in enumerate(self.dev_loader):
bos = self.vocab['<BOS>']
eos = self.vocab['<EOS>']
pad = self.vocab['<PAD>']
xs, ilens, ys, ys_in, ys_out, spks, envs, trans = to_gpu(
data, bos, eos, pad)
# feed previous
_, prediction, _, enc_outputs, enc_lens = self.model(
xs,
ilens,
None,
max_dec_timesteps=self.config['max_dec_timesteps'])
all_prediction = all_prediction + prediction.cpu().numpy().tolist()
all_ys = all_ys + [y.cpu().numpy().tolist() for y in ys]
gold_transcripts += trans
cer, prediction_sents, ground_truth_sents = self.ind2sent(
all_prediction, all_ys)
print(f'dev set CER: {cer:.4f}')
all_prediction, all_ys = [], []
gold_transcripts = []
for step, data in enumerate(test_loader):
bos = self.vocab['<BOS>']
eos = self.vocab['<EOS>']
pad = self.vocab['<PAD>']
xs, ilens, ys, ys_in, ys_out, spks, envs, trans = to_gpu(
data, bos, eos, pad)
# feed previous
_, prediction, _, enc_outputs, enc_lens = self.model(
xs,
ilens,
None,
max_dec_timesteps=self.config['max_dec_timesteps'])
all_prediction = all_prediction + prediction.cpu().numpy().tolist()
all_ys = all_ys + [y.cpu().numpy().tolist() for y in ys]
gold_transcripts += trans
self.model.train()
cer, prediction_sents, ground_truth_sents = self.ind2sent(
all_prediction, all_ys)
with open(f'{test_file_name}.txt', 'w') as f:
for p in prediction_sents:
f.write(f'{p}\n')
print(
f'{test_file_name}: {len(prediction_sents)} utterances, CER={cer:.4f}'
)
return cer
def main(config=None):
# set data loader
val_loader = get_data_loader(config, evaluate=True)
model = get_model()
trainer = Trainer(config=config, model=model, validation_loader=val_loader, train_loader_fn=train_loader_fn)
trainer.train()