# add list of ssl models to pass for training
ssl_models = []
ssl_dict = {}
for method in ssl_methods:
if 'byol' in ssl_methods:
byol_model = BYOL(
encoder,
image_size=32,
hidden_layer= -1
)
byol_model.to(device)
ssl_dict['byol'] = byol_model
if 'xxx' in ssl_methods: # add extra ssl model here
pass
if 'ae' in ssl_methods: # add extra ssl model here
ae = AE(encoder=encoder)
ae.to(device)
ssl_dict['ae'] = ae
# tasks = [[0, 1], [2], [3]]
# epochs = 2, lr = 0.0001, report_step = 20
tasks = create_tasks(dataset_name=dataset_name, cil_step=cil_step, cil_start=cil_start)
data_dir = os.path.join(root_dir, 'train')
csv_path = os.path.join(root_dir, 'train.csv')
if __name__ == '__main__':
train_set = MiniDataset(csv_path, data_dir)
train_loader = DataLoader(train_set,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
shuffle=True)
resnet = models.resnet50(pretrained=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
learner = BYOL(resnet, image_size=IMAGE_SIZE, hidden_layer='avgpool')
learner.to(device)
opt = torch.optim.Adam(learner.parameters(), lr=LR)
learner.train()
for epoch in range(EPOCHS):
for i, data in enumerate(train_loader):
data = data.to(device)
loss = learner(data)
opt.zero_grad()
loss.backward()
opt.step()
learner.update_moving_average(
) # update moving average of target encoder
print('epoch {} / {}, loss={:.4f}'.format(epoch + 1, EPOCHS,
