def get_contrastive_model(feature_dim: int) -> ContrastiveModel:
"""Init contrastive model based on parsed parametrs.
Args:
feature_dim: dimensinality of contrative projection
Returns:
ContrstiveModel instance
"""
encoder = nn.Sequential(ResnetEncoder(arch="resnet50", frozen=False),
nn.Flatten())
projection_head = nn.Sequential(
nn.Linear(2048, 512, bias=False),
nn.ReLU(inplace=True),
nn.Linear(512, feature_dim, bias=True),
)
model = ContrastiveModel(projection_head, encoder)
return model
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
from catalyst import utils
utils.set_global_seed(RANDOM_STATE)
# 1. train, valid and test loaders
transforms = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_data = MNIST(os.getcwd(),
train=True,
download=True,
transform=transforms)
train_labels = train_data.targets.cpu().numpy().tolist()
train_sampler = data.BatchBalanceClassSampler(train_labels,
num_classes=10,
num_samples=4)
train_loader = DataLoader(train_data, batch_sampler=train_sampler)
valid_dataset = MNIST(root=os.getcwd(),
transform=transforms,
train=False,
download=True)
valid_loader = DataLoader(dataset=valid_dataset, batch_size=32)
test_dataset = MNIST(root=os.getcwd(),
transform=transforms,
train=False,
download=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=32)
# 2. model and optimizer
model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 16),
nn.LeakyReLU(inplace=True))
optimizer = Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])
# 3. criterion with triplets sampling
sampler_inbatch = data.HardTripletsSampler(norm_required=False)
criterion = nn.TripletMarginLossWithSampler(
margin=0.5, sampler_inbatch=sampler_inbatch)
# 4. training with catalyst Runner
class CustomRunner(dl.SupervisedRunner):
def handle_batch(self, batch) -> None:
images, targets = batch["features"].float(
), batch["targets"].long()
features = self.model(images)
self.batch = {
"embeddings": features,
"targets": targets,
}
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="embeddings",
target_key="targets",
metric_key="loss"),
loaders="train",
),
dl.SklearnModelCallback(
feature_key="embeddings",
target_key="targets",
train_loader="train",
valid_loaders=["valid", "infer"],
model_fn=RandomForestClassifier,
predict_method="predict_proba",
predict_key="sklearn_predict",
random_state=RANDOM_STATE,
n_estimators=50,
),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="targets",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders=["valid", "infer"],
),
]
runner = CustomRunner(input_key="features", output_key="embeddings")
runner.train(
engine=engine or dl.DeviceEngine(device),
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": valid_loader,
"infer": test_loader
},
verbose=False,
valid_loader="valid",
valid_metric="accuracy",
minimize_valid_metric=False,
num_epochs=TRAIN_EPOCH,
logdir=logdir,
)
valid_path = Path(logdir) / "logs/infer.csv"
best_accuracy = max(
float(row["accuracy"]) for row in read_csv(valid_path))
assert best_accuracy > 0.8
def train_experiment(device, engine=None):
with TemporaryDirectory() as logdir:
# 1. data and transforms
transforms = Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.RandomCrop((28, 28)),
torchvision.transforms.RandomVerticalFlip(),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
Normalize((0.1307, ), (0.3081, )),
])
transform_original = Compose([
ToTensor(),
Normalize((0.1307, ), (0.3081, )),
])
mnist = MNIST("./logdir", train=True, download=True, transform=None)
contrastive_mnist = SelfSupervisedDatasetWrapper(
mnist,
transforms=transforms,
transform_original=transform_original)
train_loader = torch.utils.data.DataLoader(contrastive_mnist,
batch_size=BATCH_SIZE)
mnist_valid = MNIST("./logdir",
train=False,
download=True,
transform=None)
contrastive_valid = SelfSupervisedDatasetWrapper(
mnist_valid,
transforms=transforms,
transform_original=transform_original)
valid_loader = torch.utils.data.DataLoader(contrastive_valid,
batch_size=BATCH_SIZE)
# 2. model and optimizer
encoder = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 16),
nn.LeakyReLU(inplace=True))
projection_head = nn.Sequential(
nn.Linear(16, 16, bias=False),
nn.ReLU(inplace=True),
nn.Linear(16, 16, bias=True),
)
class ContrastiveModel(torch.nn.Module):
def __init__(self, model, encoder):
super(ContrastiveModel, self).__init__()
self.model = model
self.encoder = encoder
def forward(self, x):
emb = self.encoder(x)
projection = self.model(emb)
return emb, projection
model = ContrastiveModel(model=projection_head, encoder=encoder)
optimizer = Adam(model.parameters(), lr=LR)
# 3. criterion with triplets sampling
criterion = NTXentLoss(tau=0.1)
callbacks = [
dl.ControlFlowCallback(
dl.CriterionCallback(input_key="projection_left",
target_key="projection_right",
metric_key="loss"),
loaders="train",
),
dl.SklearnModelCallback(
feature_key="embedding_left",
target_key="target",
train_loader="train",
valid_loaders="valid",
model_fn=RandomForestClassifier,
predict_method="predict_proba",
predict_key="sklearn_predict",
random_state=RANDOM_STATE,
n_estimators=50,
),
dl.ControlFlowCallback(
dl.AccuracyCallback(target_key="target",
input_key="sklearn_predict",
topk_args=(1, 3)),
loaders="valid",
),
]
runner = dl.SelfSupervisedRunner()
logdir = "./logdir"
runner.train(
model=model,
engine=engine or dl.DeviceEngine(device),
criterion=criterion,
optimizer=optimizer,
callbacks=callbacks,
loaders={
"train": train_loader,
"valid": valid_loader
},
verbose=False,
logdir=logdir,
valid_loader="train",
valid_metric="loss",
minimize_valid_metric=True,
num_epochs=TRAIN_EPOCH,
)
valid_path = Path(logdir) / "logs/valid.csv"
best_accuracy = max(
float(row["accuracy"]) for row in read_csv(valid_path)
if row["accuracy"] != "accuracy")
assert best_accuracy > 0.6
def __init__(self, num_features: int, num_classes: int) -> None:
super().__init__()
self.model = nn.Sequential(
nn.Flatten(),
nn.Linear(in_features=num_features, out_features=num_classes))
cifar_train = CIFAR10(root="./data", download=True, transform=None)
simCLR_train = SelfSupervisedDatasetWrapper(
cifar_train,
transforms=transforms,
transform_original=transform_original)
train_loader = torch.utils.data.DataLoader(simCLR_train,
batch_size=batch_size,
num_workers=args.num_workers)
# cifar_test = CifarQGDataset(root="./data", download=True)
# valid_loader = torch.utils.data.DataLoader(
# simCLRDatasetWrapper(cifar_test, transforms=transforms), batch_size=batch_size, num_workers=5
# )
encoder = nn.Sequential(ResnetEncoder(arch="resnet18", frozen=False),
nn.Flatten())
projection_head = nn.Sequential(
nn.Linear(2048, 512, bias=False),
nn.ReLU(inplace=True),
nn.Linear(512, args.feature_dim, bias=True),
)
class ContrastiveModel(torch.nn.Module):
def __init__(self, model, encoder):
super(ContrastiveModel, self).__init__()
self.model = model
self.encoder = encoder
def forward(self, x):
emb = self.encoder(x)
projection = self.model(emb)