np.random.seed(opt.seed)
if opt.device.type != "cpu":
torch.backends.cudnn.benchmark = True
num_patches = opt.grid_dims * opt.grid_dims
# load model
models, optimizers = [None for i in range(num_patches)], [None for i in range(num_patches)]
for i in range(num_patches):
models[i], optimizers[i] = load_vision_model.load_model_and_optimizer(opt)
logs = logger.Logger(opt)
train_loader, _, supervised_loader, _, test_loader, _ = get_dataloader.get_dataloader(
opt
)
if opt.loss == 1:
train_loader = supervised_loader
try:
# Train the model
train(opt, models)
except KeyboardInterrupt:
print("Training got interrupted, saving log-files now.")
logs.create_log(models)
# random seeds
torch.manual_seed(opt.seed)
torch.cuda.manual_seed(opt.seed)
np.random.seed(opt.seed)
# load pretrained model
context_model, _ = load_vision_model.load_model_and_optimizer(
opt, reload_model=True, calc_loss=False)
context_model.module.switch_calc_loss(False)
## model_type=2 is supervised model which trains entire architecture; otherwise just extract features
if opt.model_type != 2:
context_model.eval()
_, _, train_loader, _, test_loader, _ = get_dataloader.get_dataloader(opt)
classification_model = load_vision_model.load_classification_model(opt)
if opt.model_type == 2:
params = list(context_model.parameters()) + list(
classification_model.parameters())
else:
params = classification_model.parameters()
optimizer = torch.optim.Adam(params)
criterion = torch.nn.CrossEntropyLoss()
logs = logger.Logger(opt)
try: