def demo_main(char_set, weight, name):
_, valid_transform = get_transform()
demo_data = DemoDataset('cleaned_data', name, valid_transform)
test_loader = DataLoader(
dataset=demo_data,
batch_size=3,
shuffle=False,
num_workers=1,
pin_memory=True,
)
model = ConvNet(1, len(char_set))
if torch.cuda.is_available():
model = model.cuda()
print('load weights from {}'.format(weight))
model.load_state_dict(torch.load(weight))
model.eval()
def map_indexlist_char(ind_list, char_set):
return ''.join([char_set[i] for i in ind_list])
with torch.no_grad():
for batch_idx, (x, imgpath) in enumerate(test_loader):
if batch_idx > 0:
break
x = x.cuda()
out = model(x)
_, pred_label = torch.max(out, 1)
pred_name = map_indexlist_char(pred_label.tolist(), char_set)
print('name {} pred name {}'.format(name, pred_name))
def get_concat(im1, im2):
dst = Image.new('RGB', (im1.width + im2.width, im1.height))
dst.paste(im1, (0, 0))
dst.paste(im2, (im1.width, 0))
return dst
concat_im = None
for img in demo_data.images():
im = Image.open(img)
if concat_im is None:
concat_im = im
else:
concat_im = get_concat(concat_im, im)
#concat_im.show()
concat_im.save('demo.jpg')
def initialise():
game_controller = GameController(game_cfg.start_bbox, game_cfg.end_bbox,
game_cfg.start_thres)
player_controller = PlayerController(general_cfg.app)
rl_recorder = RlRecorder()
# TODO, read replay from disk
player_controller.activate_chrome() # switch to chrome
timer = Timer(game_cfg.space_time_gap)
performances = {'iter': [], 'score': []}
if cnn_cfg.load_model and os.path.isfile(cnn_cfg.chkpnt_path):
cnn = torch.load(cnn_cfg.chkpnt_path)
cnn.cnn_cfg = cnn_cfg
print("Load cnn model from ", cnn_cfg.chkpnt_path)
else:
cnn = ConvNet(cnn_cfg, num_classes=cnn_cfg.num_classes, lr=cnn_cfg.lr)
print("Create new CNN done!")
if torch.cuda.is_available():
cnn = cnn.cuda()
print("Cuda is available!")
return game_controller, player_controller, rl_recorder, timer, performances, cnn
shuffle=True)
test_dataset = datasets.MNIST(root='../data/',
train=False,
download=True,
transform=transforms.ToTensor())
loader_test = torch.utils.data.DataLoader(test_dataset,
batch_size=param['test_batch_size'],
shuffle=True)
# Load the pretrained model
net = ConvNet()
net.load_state_dict(torch.load('models/convnet_pretrained.pkl'))
if torch.cuda.is_available():
print('CUDA ensabled.')
net.cuda()
print("--- Pretrained network loaded ---")
test(net, loader_test)
# prune the weights
masks = filter_prune(net, param['pruning_perc'])
net.set_masks(masks)
print("--- {}% parameters pruned ---".format(param['pruning_perc']))
test(net, loader_test)
# Retraining
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.RMSprop(net.parameters(),
lr=param['learning_rate'],
weight_decay=param['weight_decay'])
def main():
# data normalization
input_size = 224
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# data loaders
kwargs = {'num_workers': 8, 'pin_memory': True} if args.cuda else {}
if args.da:
train_transforms = transforms.Compose([
random_transform,
transforms.ToPILImage(),
transforms.Resize((input_size, input_size)),
transforms.ToTensor(), normalize
])
else:
train_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((input_size, input_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
])
test_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((input_size, input_size)),
transforms.ToTensor(), normalize
])
train_loader = torch.utils.data.DataLoader(DataLoader(df_train,
train_transforms,
root=args.data_dir,
mode=args.mode),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(DataLoader(df_gal,
test_transforms,
root=args.data_dir,
mode=args.mode),
batch_size=args.batch_size,
shuffle=False,
**kwargs)
# instanciate the models
output_shape, backbone = get_backbone(args)
embed = LinearProjection(output_shape, args.dim_embed)
model = ConvNet(backbone, embed)
# instanciate the proxies
fsem = get_semantic_fname(args.word)
path_semantic = os.path.join('aux', 'Semantic', args.dataset, fsem)
train_proxies = get_proxies(path_semantic, df_train['cat'].cat.categories)
test_proxies = get_proxies(path_semantic, df_gal['cat'].cat.categories)
train_proxynet = ProxyNet(args.n_classes,
args.dim_embed,
proxies=torch.from_numpy(train_proxies))
test_proxynet = ProxyNet(args.n_classes_gal,
args.dim_embed,
proxies=torch.from_numpy(test_proxies))
# criterion
criterion = ProxyLoss(args.temperature)
if args.multi_gpu:
model = nn.DataParallel(model)
if args.cuda:
backbone.cuda()
embed.cuda()
model.cuda()
train_proxynet.cuda()
test_proxynet.cuda()
parameters_set = []
low_layers = []
upper_layers = []
for c in backbone.children():
low_layers.extend(list(c.parameters()))
for c in embed.children():
upper_layers.extend(list(c.parameters()))
parameters_set.append({
'params': low_layers,
'lr': args.lr * args.factor_lower
})
parameters_set.append({'params': upper_layers, 'lr': args.lr * 1.})
optimizer = optim.SGD(parameters_set,
lr=args.lr,
momentum=0.9,
nesterov=True,
weight_decay=args.wd)
n_parameters = sum([p.data.nelement() for p in model.parameters()])
print(' + Number of params: {}'.format(n_parameters))
scheduler = CosineAnnealingLR(optimizer,
args.epochs * len(train_loader),
eta_min=3e-6)
print('Starting training...')
for epoch in range(args.start_epoch, args.epochs + 1):
# update learning rate
scheduler.step()
# train for one epoch
train(train_loader, model, train_proxynet.proxies.weight, criterion,
optimizer, epoch, scheduler)
val_acc = evaluate(test_loader, model, test_proxynet.proxies.weight,
criterion)
# saving
if epoch == args.epochs:
save_checkpoint({'epoch': epoch, 'state_dict': model.state_dict()})
print('\nResults on test set (end of training)')
write_logs('\nResults on test set (end of training)')
test_acc = evaluate(test_loader, model, test_proxynet.proxies.weight,
criterion)
def main():
# data normalization
input_size = 224
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# data loaders
kwargs = {'num_workers': 8, 'pin_memory': True} if args.cuda else {}
test_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((input_size, input_size)),
transforms.ToTensor(), normalize
])
feats = {}
labels = {}
for domain in ['im', 'sk']:
key = '_'.join([domain, 'model_path'])
dirname = os.path.dirname(args.__dict__[key])
fpath = os.path.join(dirname, 'features.npz')
results_path = os.path.join(dirname, 'results.txt')
if os.path.isfile(fpath) and args.rewrite is False:
data = np.load(fpath)
feats[domain] = data['features']
labels[domain] = data['labels']
txt = ('Domain (%s): Acc %.2f' % (domain, data['acc'] * 100.))
print(txt)
write_logs(txt, results_path)
df_gal = splits[domain]['gal']
fsem = get_semantic_fname(args.word)
path_semantic = os.path.join('aux', 'Semantic', args.dataset, fsem)
test_proxies = get_proxies(path_semantic,
df_gal['cat'].cat.categories)
else:
df_gal = splits[domain]['gal']
test_loader = torch.utils.data.DataLoader(
DataLoader(df_gal,
test_transforms,
root=args.data_dir,
mode=domain),
batch_size=args.batch_size * 10,
shuffle=False,
**kwargs)
# instanciate the models
output_shape, backbone = get_backbone(args)
embed = LinearProjection(output_shape, args.dim_embed)
model = ConvNet(backbone, embed)
# instanciate the proxies
fsem = get_semantic_fname(args.word)
path_semantic = os.path.join('aux', 'Semantic', args.dataset, fsem)
test_proxies = get_proxies(path_semantic,
df_gal['cat'].cat.categories)
test_proxynet = ProxyNet(args.n_classes_gal,
args.dim_embed,
proxies=torch.from_numpy(test_proxies))
# criterion
criterion = ProxyLoss(args.temperature)
if args.multi_gpu:
model = nn.DataParallel(model)
# loading
checkpoint = torch.load(args.__dict__[key])
model.load_state_dict(checkpoint['state_dict'])
txt = ("\n=> loaded checkpoint '{}' (epoch {})".format(
args.__dict__[key], checkpoint['epoch']))
print(txt)
write_logs(txt, results_path)
if args.cuda:
backbone.cuda()
embed.cuda()
model.cuda()
test_proxynet.cuda()
txt = 'Extracting testing set (%s)...' % (domain)
print(txt)
x, y, acc = extract_predict(test_loader, model,
test_proxynet.proxies.weight,
criterion)
feats[domain] = x
labels[domain] = y
np.savez(fpath,
features=feats[domain],
labels=labels[domain],
acc=acc)
fpath_train = os.path.join(dirname, 'features_train.npz')
if args.train and not os.path.isfile(fpath_train):
df_train = splits[domain]['train']
train_loader = torch.utils.data.DataLoader(
DataLoader(df_train,
test_transforms,
root=args.data_dir,
mode=domain),
batch_size=args.batch_size * 10,
shuffle=False,
**kwargs)
train_proxies = get_proxies(path_semantic,
df_train['cat'].cat.categories)
train_proxynet = ProxyNet(
args.n_classes_gal,
args.dim_embed,
proxies=torch.from_numpy(train_proxies))
train_proxynet.cuda()
txt = 'Extracting training set (%s)...' % (domain)
print(txt)
x, y, _ = extract_predict(train_loader, model,
train_proxynet.proxies.weight,
criterion)
fpath = os.path.join(dirname, 'features_train.npz')
np.savez(fpath,
features=feats[domain],
features_train=x,
labels=labels[domain],
labels_train=y,
acc=acc)
txt = ('Domain (%s): Acc %.2f' % (domain, acc * 100.))
print(txt)
write_logs(txt, results_path)
if args.shape:
print('\nRetrieval per model')
new_feat_im, new_labels_im = average_views(splits['im']['test'],
feats['im'], labels['im'])
idx = retrieve(feats['sk'], new_feat_im)
metrics = score_shape(labels['sk'], new_labels_im, idx)
names = ['NN', 'FT', 'ST', 'E', 'nDCG', 'mAP']
txt = [('%s %.3f' % (name, value))
for name, value in zip(names, metrics)]
txt = '\t'.join(txt)
print(txt)
write_logs(txt, results_path)
print('\nRetrieval per model with refinement')
alpha = 0.4
g_sk_x = KNN(feats['sk'], new_feat_im, K=1, mode='ones')
new_sk_x = slerp(alpha, L2norm(feats['sk']), L2norm(g_sk_x))
idx = retrieve(new_sk_x, new_feat_im)
metrics = score_shape(labels['sk'], new_labels_im, idx)
names = ['NN', 'FT', 'ST', 'E', 'nDCG', 'mAP']
txt = [('%s %.3f' % (name, value))
for name, value in zip(names, metrics)]
txt = '\t'.join(txt)
print(txt)
write_logs(txt, results_path)
else:
print('\nRetrieval')
txt = evaluate(feats['im'], labels['im'], feats['sk'], labels['sk'])
print(txt)
write_logs(txt, results_path)
print('\nRetrieval with refinement')
if args.overwrite:
alpha = 0.7
else:
alpha = 0.4
g_sk_x = KNN(feats['sk'], feats['im'], K=1, mode='ones')
new_sk_x = slerp(alpha, L2norm(feats['sk']), L2norm(g_sk_x))
txt = evaluate(feats['im'], labels['im'], new_sk_x, labels['sk'])
print(txt)
write_logs(txt, results_path)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=1000)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=1000)
################ initialize the model ################
if args.model == 'convnet':
model = ConvNet()
elif args.model == 'mymodel':
model = MyModel()
else:
raise Exception('Incorrect model name')
if args.cuda:
model.cuda()
######## Define loss function and optimizer ##########
############## Write your code here ##################
params = model.parameters()
optimizer = optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
criterion = nn.CrossEntropyLoss()
######################################################
def train(epoch):
""" Runs training for 1 epoch
epoch: int, denotes the epoch number for printing
"""
############# Write train function ###############
mean_training_loss = 0.0
def main():
# data normalization
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# data loaders
kwargs = {'num_workers': 8, 'pin_memory': True} if args.cuda else {}
test_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((224, 224)),
transforms.ToTensor(), normalize
])
feats = {}
labels = {}
for domain in ['im', 'sk']:
key = '_'.join([domain, 'model_path'])
dirname = os.path.dirname(args.__dict__[key])
fpath = os.path.join(dirname, 'features.npz')
results_path = os.path.join(dirname, 'results.txt')
if os.path.isfile(fpath) and args.rewrite is False:
data = np.load(fpath)
feats[domain] = data['features']
labels[domain] = data['labels']
txt = ('Domain (%s): Acc %.2f' % (domain, data['acc'] * 100.))
print(txt)
write_logs(txt, results_path)
df_gal = splits[domain]['test']
fsem = get_semantic_fname(args.word)
path_semantic = os.path.join('aux', 'Semantic', args.dataset, fsem)
test_proxies = get_proxies(path_semantic,
df_gal['cat'].cat.categories)
else:
df_gal = splits[domain]['test']
test_loader = torch.utils.data.DataLoader(
DataLoader(df_gal,
test_transforms,
root=args.data_dir,
mode=domain),
batch_size=args.batch_size * 1,
shuffle=False,
**kwargs)
# instanciate the models
output_shape, backbone = get_backbone(args)
embed = LinearProjection(output_shape, args.dim_embed)
model = ConvNet(backbone, embed)
# instanciate the proxies
fsem = get_semantic_fname(args.word)
path_semantic = os.path.join('aux', 'Semantic', args.dataset, fsem)
test_proxies = get_proxies(path_semantic,
df_gal['cat'].cat.categories)
test_proxynet = ProxyNet(args.n_classes_gal,
args.dim_embed,
proxies=torch.from_numpy(test_proxies))
# criterion
criterion = ProxyLoss(args.temperature)
if args.multi_gpu:
model = nn.DataParallel(model)
# loading
checkpoint = torch.load(args.__dict__[key])
model.load_state_dict(checkpoint['state_dict'])
txt = ("\n=> loaded checkpoint '{}' (epoch {})".format(
args.__dict__[key], checkpoint['epoch']))
print(txt)
if args.cuda:
backbone.cuda()
embed.cuda()
model.cuda()
test_proxynet.cuda()
txt = 'Extracting testing set (%s)...' % (domain)
print(txt)
x, y, acc = extract_predict(test_loader, model,
test_proxynet.proxies.weight,
criterion)
feats[domain] = x
labels[domain] = y
np.savez(fpath,
features=feats[domain],
labels=labels[domain],
acc=acc)
txt = ('Domain (%s): Acc %.2f' % (domain, acc * 100.))
print(txt)
print('\nFew-Shot')
fs(feats, labels, test_proxies)
def main(args):
init_process_group(backend='nccl')
with open(args.config) as file:
config = json.load(file)
config.update(vars(args))
config = apply_dict(Dict, config)
backends.cudnn.benchmark = True
backends.cudnn.fastest = True
cuda.set_device(distributed.get_rank() % cuda.device_count())
train_dataset = ImageDataset(root=config.train_root,
meta=config.train_meta,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, ) * 3,
(0.5, ) * 3)
]))
val_dataset = ImageDataset(root=config.val_root,
meta=config.val_meta,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, ) * 3,
(0.5, ) * 3)
]))
train_sampler = utils.data.distributed.DistributedSampler(train_dataset)
val_sampler = utils.data.distributed.DistributedSampler(val_dataset)
train_data_loader = utils.data.DataLoader(
dataset=train_dataset,
batch_size=config.local_batch_size,
sampler=train_sampler,
num_workers=config.num_workers,
pin_memory=True)
val_data_loader = utils.data.DataLoader(dataset=val_dataset,
batch_size=config.local_batch_size,
sampler=val_sampler,
num_workers=config.num_workers,
pin_memory=True)
model = ConvNet(conv_params=[
Dict(in_channels=3,
out_channels=32,
kernel_size=5,
padding=2,
stride=2,
bias=False),
Dict(in_channels=32,
out_channels=64,
kernel_size=5,
padding=2,
stride=2,
bias=False),
],
linear_params=[
Dict(in_channels=3136,
out_channels=1024,
kernel_size=1,
bias=False),
Dict(in_channels=1024,
out_channels=10,
kernel_size=1,
bias=True),
])
config.global_batch_size = config.local_batch_size * distributed.get_world_size(
)
config.optimizer.lr *= config.global_batch_size / config.global_batch_denom
optimizer = optim.Adam(model.parameters(), **config.optimizer)
epoch = 0
global_step = 0
if config.checkpoint:
checkpoint = Dict(torch.load(config.checkpoint))
model.load_state_dict(checkpoint.model_state_dict)
optimizer.load_state_dict(checkpoint.optimizer_state_dict)
epoch = checkpoint.last_epoch + 1
global_step = checkpoint.global_step
def train(data_loader):
nonlocal global_step
model.train()
for images, labels in data_loader:
images = images.cuda()
labels = labels.cuda()
optimizer.zero_grad()
logits = model(images)
loss = nn.functional.cross_entropy(logits, labels)
loss.backward(retain_graph=True)
average_gradients(model.parameters())
optimizer.step()
predictions = logits.topk(k=1, dim=1)[1].squeeze()
accuracy = torch.mean((predictions == labels).float())
average_tensors([loss, accuracy])
global_step += 1
dprint(f'[training] epoch: {epoch} global_step: {global_step} '
f'loss: {loss:.4f} accuracy: {accuracy:.4f}')
@torch.no_grad()
def validate(data_loader):
model.eval()
losses = []
accuracies = []
for images, labels in data_loader:
images = images.cuda()
labels = labels.cuda()
logits = model(images)
loss = nn.functional.cross_entropy(logits, labels)
predictions = logits.topk(k=1, dim=1)[1].squeeze()
accuracy = torch.mean((predictions == labels).float())
average_tensors([loss, accuracy])
losses.append(loss)
accuracies.append(accuracy)
loss = torch.mean(torch.stack(losses)).item()
accuracy = torch.mean(torch.stack(accuracies)).item()
dprint(f'[validation] epoch: {epoch} global_step: {global_step} '
f'loss: {loss:.4f} accuracy: {accuracy:.4f}')
@torch.no_grad()
def feed(data_loader):
model.eval()
for images, _ in data_loader:
images = images.cuda()
logits = model(images)
def save():
if not distributed.get_rank():
os.makedirs('checkpoints', exist_ok=True)
torch.save(
dict(model_state_dict=model.state_dict(),
optimizer_state_dict=optimizer.state_dict(),
last_epoch=epoch,
global_step=global_step),
os.path.join('checkpoints', f'epoch_{epoch}'))
if config.training:
model.cuda()
broadcast_tensors(model.state_dict().values())
for epoch in range(epoch, config.num_training_epochs):
train_sampler.set_epoch(epoch)
train(train_data_loader)
validate(val_data_loader)
save()
if config.validation:
model.cuda()
broadcast_tensors(model.state_dict().values())
validate(val_data_loader)
if config.quantization:
model.cuda()
broadcast_tensors(model.state_dict().values())
with QuantizationEnabler(model):
with BatchStatsUser(model):
for epoch in range(epoch, config.num_quantization_epochs):
train_sampler.set_epoch(epoch)
train(train_data_loader)
validate(val_data_loader)
save()
with AverageStatsUser(model):
for epoch in range(epoch, config.num_quantization_epochs):
train_sampler.set_epoch(epoch)
train(train_data_loader)
validate(val_data_loader)
save()
