transforms.Compose([
#transforms.Resize(256),
#transforms.CenterCrop(224),
transforms.Resize((900, 600)),
transforms.CenterCrop((850, 550)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
test_transforms = transforms.Compose([
#transforms.Resize(256),
transforms.RandomRotation(15),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(),
transformations.CropField(),
#transformations.TileCircle(),
transforms.CenterCrop((224, 224)),
#p.torch_transform(),
transforms.ToTensor(),
#transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
#transforms.Normalize([0.7410, 0.7141, 0.6500], [0.0808, 0.0895, 0.1141])
transforms.Normalize([0.2974, 0.3233, 0.2370], [0.1399, 0.1464, 0.1392])
])
test_data_augmentation = 5
test_transforms_grande = transforms.Compose([
transforms.Resize((900, 600)),
transforms.CenterCrop((850, 550)),
#transforms.ColorJitter(),
#transforms.RandomRotation(30),
#p.torch_transform(),
def main(args):
# Training settings
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
config = wandb.config
config.lr = args.lr
config.batch_size = args.batch_size
config.gamma = args.gamma
config.epochs = args.epochs
config.test_batch_size = args.test_batch_size
config.log_interval = args.log_interval
config.image_size = args.image_size
config.dry_run = args.dry_run
config.num_workers = args.num_workers
config.stage = args.stage
config.run_name = args.run_name
config.data_dir = args.data_dir if args.data_dir is not None else '/dev/shm/dataset'
tags = ['baseline', f'stage {config.stage}']
wandb.init(project='FADA', config=config, tags=tags)
if config.run_name is not None:
wandb.run.name = config.run_name
result_path = os.path.join(args.result_path, wandb.run.name)
config.result_path = result_path
if not os.path.exists(result_path):
os.makedirs(result_path)
device = torch.device("cuda" if use_cuda else "cpu")
train_kwargs = {'batch_size': config.batch_size }
test_kwargs = {'batch_size': config.test_batch_size }
if use_cuda:
cuda_kwargs = {'num_workers': config.num_workers,
'pin_memory': True,
'shuffle': True}
train_kwargs.update(cuda_kwargs)
test_kwargs.update(cuda_kwargs)
data_transforms = {
'train': transforms.Compose([
transforms.RandomRotation(15),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(),
transformations.TileHerb(),
transforms.CenterCrop((config.image_size, config.image_size)),
transforms.ToTensor(),
transforms.Normalize([0.2974, 0.3233, 0.2370], [0.1399, 0.1464, 0.1392])
]),
'val': transforms.Compose([
transformations.CropField(),
transforms.CenterCrop((config.image_size, config.image_size)),
transforms.ToTensor(),
transforms.Normalize([0.2974, 0.3233, 0.2370], [0.1399, 0.1464, 0.1392])
]),
'val_photo': transforms.Compose([
transformations.CropField(),
transforms.CenterCrop((config.image_size, config.image_size)),
transforms.ToTensor(),
transforms.Normalize([0.2974, 0.3233, 0.2370], [0.1399, 0.1464, 0.1392])
])
}
data_dir = config.data_dir
herbarium = os.path.join(data_dir, 'herbarium')
photo = os.path.join(data_dir, 'photo')
classifier = models.ClassifierPro().to(device)
encoder = tmodels.resnet50(pretrained=True).to(device)
encoder.fc = nn.Sequential()
ssnet = models.TaxonNet(64).to(device)
genusnet = models.TaxonNet(510).to(device)
familynet = models.TaxonNet(151).to(device)
discriminator = models.DCDPro().to(device)
discriminator_genus = models.DCDPro().to(device)
discriminator_family = models.DCDPro().to(device)
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
data_transforms[x])
for x in ['train', 'val']}
base_mapping = image_datasets['train'].class_to_idx
class_name_to_id = image_datasets['train'].class_to_idx
id_to_class_name = {v: k for k, v in class_name_to_id.items()}
siamese_dataset = data_loader.FADADatasetSS(data_dir,
photo,
'train',
image_datasets['train'].class_to_idx,
class_name_to_id,
config.image_size
)
if(config.stage == 1):
stage_1(config, device, image_datasets, classifier, encoder, ssnet, genusnet, familynet, train_kwargs, test_kwargs)
elif(config.stage == 2):
encoder.load_state_dict(torch.load('../best/encoder_fada_extra.pth'))
classifier.load_state_dict(torch.load('../best/classifier_fada_extra.pth'))
stage_2(config, device, siamese_dataset, discriminator, encoder, discriminator_genus, discriminator_family, train_kwargs, test_kwargs)