test_transforms = transforms.Compose([transforms.Resize((385, 505)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
train_transforms.transforms.insert(0, RandAugment(2, 14))
class2idx = config['class2idx']
batch = config['batch_size']
num_workers = config['num_workers']
train_data = datasets.ImageFolder(root = config['Train'], transform = train_transforms)
valid_data = myset(config['Val'], config['val_txt'], class2idx, test_transforms)
test_data = myset(config['Test'], config['test_txt'], class2idx, test_transforms)
train_loader = torch.utils.data.DataLoader(train_data, batch_size = batch , num_workers = num_workers,shuffle = True)
valid_loader = torch.utils.data.DataLoader(valid_data, batch_size = batch, num_workers = num_workers,shuffle = False)
test_loader = torch.utils.data.DataLoader(test_data, batch_size = batch, num_workers = num_workers,shuffle = False)
ms = []
model1 = res(3)
ms.append(model1)
model2 = mob(3)
ms.append(model2)
model3 = alex(3)
ms.append(model3)
optimizer = optim.Adam([{'params': model1.parameters()}, {'params': model2.parameters()}, {'params': model3.parameters()}], lr = 0.00001, weight_decay=1e-3)
dml = DML(ms, optimizer, parallel = True)
dml.train(300, train_loader, valid_loader)
