def main():
# enable mixed-precision computation if desired
if args.amp:
mixed_precision.enable_mixed_precision()
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# get the dataset
dataset = get_dataset(args.dataset)
_, test_loader, _ = build_dataset(
dataset=dataset, batch_size=args.batch_size, input_dir=args.input_dir
)
torch_device = torch.device("cuda")
checkpointer = Checkpointer()
model = checkpointer.restore_model_from_checkpoint(args.checkpoint_path)
model = model.to(torch_device)
model, _ = mixed_precision.initialize(model, None)
test_stats = AverageMeterSet()
test(model, test_loader, torch_device, test_stats)
stat_str = test_stats.pretty_string(ignore=model.tasks)
print(stat_str)
def main():
# create target output dir if it doesn't exist yet
if not os.path.isdir(args.output_dir):
os.mkdir(args.output_dir)
# enable mixed-precision computation if desired
amp = ""
if args.amp:
amp = "torch"
if args.apex:
print("Error: Cannot use both --amp and --apex.")
exit()
if args.apex:
amp = "apex"
mixed_precision.enable_mixed_precision()
# set the RNG seeds (probably more hidden elsewhere...)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# get the dataset
dataset = get_dataset(args.dataset)
encoder_size = get_encoder_size(dataset)
# get a helper object for tensorboard logging
log_dir = os.path.join(args.output_dir, args.run_name)
stat_tracker = StatTracker(log_dir=log_dir)
# get dataloaders for training and testing
train_loader, test_loader, num_classes = \
build_dataset(dataset=dataset,
batch_size=args.batch_size,
input_dir=args.input_dir,
labeled_only=args.classifiers)
torch_device = torch.device('cuda')
checkpointer = Checkpointer(args.output_dir)
if args.cpt_load_path:
model = checkpointer.restore_model_from_checkpoint(
args.cpt_load_path, training_classifier=args.classifiers)
else:
# create new model with random parameters
model = Model(ndf=args.ndf,
n_classes=num_classes,
n_rkhs=args.n_rkhs,
tclip=args.tclip,
n_depth=args.n_depth,
encoder_size=encoder_size,
use_bn=(args.use_bn == 1))
model.init_weights(init_scale=1.0)
checkpointer.track_new_model(model)
model = model.to(torch_device)
# select which type of training to do
task = train_classifiers if args.classifiers else train_self_supervised
task(model, args.learning_rate, dataset, train_loader, test_loader,
stat_tracker, checkpointer, args.output_dir, torch_device, amp)
def main():
# create target output dir if it doesn't exist yet
if not os.path.isdir(args['output_dir']):
os.mkdir(args['output_dir'])
# enable mixed-precision computation if desired
if args['amp']:
mixed_precision.enable_mixed_precision()
# set the RNG seeds (probably more hidden elsewhere...)
torch.manual_seed(args['seed'])
torch.cuda.manual_seed(args['seed'])
# get the dataset
dataset = get_dataset(args['dataset'])
encoder_size = get_encoder_size(dataset)
# get a helper object for tensorboard logging
log_dir = os.path.join(args['output_dir'], args['run_name'])
stat_tracker = StatTracker(log_dir=log_dir)
# get dataloaders for training and testing
train_loader, test_loader, num_classes = \
build_dataset(dataset=dataset,
batch_size=args['batch_size'],
input_dir=args['input_dir'],
labeled_only=args['classifiers'])
torch_device = torch.device('cuda')
checkpointer = Checkpointer(args['output_dir'])
if args['cpt_load_path']:
model = checkpointer.restore_model_from_checkpoint(
args['cpt_load_path'], training_classifier=args['classifiers'])
else:
# create new model with random parameters
model = Model(ndf=args['ndf'],
n_classes=num_classes,
n_rkhs=args['n_rkhs'],
tclip=args['tclip'],
n_depth=args['n_depth'],
encoder_size=encoder_size,
use_bn=(args['use_bn'] == 1))
model.init_weights(init_scale=1.0)
checkpointer.track_new_model(model)
model = model.to(torch_device)
# select which type of training to do
task = train_classifiers if args['classifiers'] else train_self_supervised
if args['classifiers']:
task = train_classifiers
elif args['decoder']:
task = train_decoder
else:
task = train_self_supervised
task(model, args['learning_rate'], dataset, train_loader, test_loader,
stat_tracker, checkpointer, args['output_dir'], torch_device)
def obtain_model(model_type):
if model_type != 'robust':
checkpoint_path = 'runs/amdim_cpt.pth'
checkpointer = Checkpointer()
print('Loading model')
model = checkpointer.restore_model_from_checkpoint(checkpoint_path)
torch_device = torch.device('cuda')
model = model.to(torch_device)
else:
dataset = robustness.datasets.CIFAR()
model_kwargs = {
'arch':
'resnet50',
'dataset':
dataset,
'resume_path':
f'../robust_classif/robustness_applications/models/CIFAR.pt'
}
model, _ = model_utils.make_and_restore_model(**model_kwargs)
model.eval()
model = CommonModel(model, model_type)
return model
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
img = self._load_image(item['image'])
# convert 3D to 2D tensor to store in kaldi-format
uttid = item["uttid"]
return {"uttid": uttid, "image": img}
if __name__ == '__main__':
parser = get_parser()
args = parser.parse_args()
# Restore model from the checkpoint
ckpt = Checkpointer()
ckpt.restore_model_from_checkpoint(
cpt_path="amdim_ndf256_rkhs2048_rd10.pth")
ckpt.model.to('cuda')
img_tmp_ark = os.path.splitext(args.img_as_feats_scp)[0] + '.tmp.ark'
ds = ImageDataset(args.places_json)
with kio.open_or_fd(img_tmp_ark, 'wb') as f:
for i in tqdm(range(len(ds))):
item = ds[i]
feats = item["image"]
batch = torch.zeros(2, 3, 128, 128)
batch[0] = feats
batch = batch.to('cuda')
res_dict = ckpt.model(x1=batch, x2=batch, class_only=True)
global_feats = res_dict["rkhs_glb"][:1]
k = item["uttid"]
kio.write_mat(f, global_feats.cpu().detach().numpy(), key=k)
#update run id in config file
model_conf['run_id'] = run_id
json.dump(model_conf, open('config.json', 'w'))
#Construct DataLoader and checkpointer
train_loader = Image_Loader(img_path,
batch_size=batch_size,
shuffle=True,
drop_last=True,
num_workers=num_workers,
input_shape=input_shape,
stage='train')
checkpointer = Checkpointer(run=run)
# Load checkpoint if given, otherwise construct a new model
encoder, mi_estimator = checkpointer.restore_model_from_checkpoint()
# Compute on multiple GPUs, if there are more than one given
if torch.cuda.device_count() > 1:
print("Let's use %d GPUs" % torch.cuda.device_count())
encoder = torch.nn.DataParallel(encoder).module
mi_estimator = torch.nn.DataParallel(mi_estimator).module
encoder.to(device)
mi_estimator.to(device)
enc_optim = torch.optim.Adam(encoder.parameters(), lr=lr)
mi_optim = torch.optim.Adam(mi_estimator.parameters(), lr=lr)
try:
encoder.train()
mi_estimator.train()
torch.autograd.set_detect_anomaly(True)