#################################################
# Script runtime options
model_name = 'vit_base_patch16_224'
model_func = timm.create_model
root_dir = '../..'
data_dir = os.path.join(root_dir,'data/images')
working_dir = os.path.join(root_dir,'models/classification', model_name)
batch_size = 16
num_workers = 4
num_epochs = 40
#################################################
os.makedirs(working_dir, exist_ok=True)
# Get data transforms
data_transforms = makeDefaultTransforms()
# Setup data loaders with augmentation transforms
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
for x in ['train', 'test']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
shuffle=True, num_workers=num_workers)
for x in ['train', 'test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
class_names = image_datasets['train'].classes
print('Number of data')
print('========================================')
for dataset in dataset_sizes.keys():
print(dataset,' size:: ', dataset_sizes[dataset],' images')
#################################################
# Script runtime options
model_name = 'swin_small_patch4_window7_224'
model_func = timm.create_model
root_dir = '../..'
data_dir = os.path.join(root_dir, 'data/images')
working_dir = os.path.join(root_dir, 'models/classification', model_name)
batch_size = 16
num_workers = 4
num_epochs = 40
#################################################
os.makedirs(working_dir, exist_ok=True)
# Get data transforms
data_transforms = makeDefaultTransforms(img_crop_size=224, img_resize=256)
# Setup data loaders with augmentation transforms
image_datasets = {
x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
for x in ['train', 'test']
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
for x in ['train', 'test']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
}
model_func = ptcv_prepare_model
is_model_torchvision = False
root_dir = '../..'
data_dir = os.path.join(root_dir, 'data/images')
working_dir = os.path.join(root_dir, 'models/classification', model_name)
batch_size = 16
num_workers = 4
num_epochs = 40
#################################################
os.makedirs(working_dir, exist_ok=True)
# Get data transforms
data_transforms = makeDefaultTransforms(img_crop_size=331, img_resize=512)
# Setup data loaders with augmentation transforms
image_datasets = {
x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
for x in ['train', 'test']
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
for x in ['train', 'test']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
class_names = image_datasets['train'].classes
## SETUP DIRS
# Clean up the output directory by removing it if desired
if clean_up and (os.path.exists(working_dir)):
shutil.rmtree(working_dir)
# Create the output directory for results
os.makedirs(working_dir, exist_ok=True)
## SETUP DEVICE
# Setup the device to run the computations
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
## SETUP DATALOADERS
print('[INFO] Begining setting up training of model {}'.format(model_name))
print('[INFO] Setting up dataloaders for train and test sets.')
# Get data transforms
data_transforms = makeDefaultTransforms(img_crop_size=img_crop_size, img_resize=img_resize)
# Make train and test loaders
train_loader, val_loader = create_dataloaders(data_transforms=data_transforms, data_dir=data_dir, batch_size=batch_size, num_workers=num_workers)
## SETUP MODEL OBJECTS
# Get the model, optimizer, loss criterion and learning rate scheduler objects
print('[INFO] Getting model {} from library and sending to device, setting up loss criterion, optimizer and learning rate scheduler...'.format(model_name), end='')
model, optimizer, criterion, lr_scheduler = initialize(
device=device,
model_func=model_func,
model_args=model_args,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
is_torchvision=is_torchvision,
num_classes=200
