print("[INFO] Device type:", str(device))
from datamanager import DataManager
manager = DataManager(args.seed)
num_classes = manager.get_num_classes(args.dataset)
train_transform = manager.get_train_transforms("lineval", args.dataset)
train_loader, _ = manager.get_train_loader(dataset=args.dataset,
data_type="single",
data_size=args.data_size,
train_transform=train_transform,
repeat_augmentations=None,
num_workers=args.num_workers,
drop_last=False)
test_loader = manager.get_test_loader(args.dataset, args.data_size)
if (args.backbone == "conv4"):
from backbones.conv4 import Conv4
feature_extractor = Conv4(flatten=True)
elif (args.backbone == "resnet8"):
from backbones.resnet_small import ResNet, BasicBlock
feature_extractor = ResNet(BasicBlock, [1, 1, 1],
channels=[16, 32, 64],
flatten=True)
elif (args.backbone == "resnet32"):
from backbones.resnet_small import ResNet, BasicBlock
feature_extractor = ResNet(BasicBlock, [5, 5, 5],
channels=[16, 32, 64],
flatten=True)
elif (args.backbone == "resnet56"):
repeat_augmentations=args.K,
num_workers=args.num_workers,
drop_last=False)
elif (args.method == "standard"):
from methods.standard import StandardModel
model = StandardModel(feature_extractor,
num_classes,
tot_epochs=args.epochs)
if (args.dataset == "stl10"):
train_transform = manager.get_train_transforms("finetune",
args.dataset)
else:
train_transform = manager.get_train_transforms("standard",
args.dataset)
test_loader = manager.get_test_loader(args.dataset,
data_size=args.data_size,
num_workers=args.num_workers)
train_loader, _ = manager.get_train_loader(dataset=args.dataset,
data_type="single",
data_size=args.data_size,
train_transform=train_transform,
repeat_augmentations=None,
num_workers=args.num_workers,
drop_last=False)
elif (args.method == "rotationnet"):
from methods.rotationnet import Model
model = Model(feature_extractor)
train_transform = manager.get_train_transforms(args.method, args.dataset)
if (args.dataset == "stl10"): data_type = "unsupervised"
else: data_type = "single"
train_loader, _ = manager.get_train_loader(dataset=args.dataset,
