class BaseModel:
def __init__(self,
model=None,
in_channels=3,
n_classes=5,
batch_size=128,
*args,
**kwargs):
if model is None:
# define the model
self.model = ResNetSimCLR(in_channels, n_classes, *args, **kwargs)
else:
# use a pretrained model
self.model = model
self.model = self.model.to(self.device)
self.num_params = sum(p.numel() for p in self.model.parameters())
self.batch_size = batch_size
@property
def device(self):
return torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
def load_model(self, path):
self.model.load_state_dict(torch.load(path)['model_state_dict'])
def load_data(self, dataset):
# create data loader
data_loader = DataLoader(dataset,
batch_size=self.batch_size,
drop_last=True,
shuffle=True,
num_workers=2)
return data_loader
def end_to_end():
s=1
input_shape = (128,128,3)
# set up training data
train_trans = TrainDataAugmentation(s, input_shape)
train_aug = train_trans.augment()
train_data = datasets.ImageFolder('unlabelled', transform=SimCLRDataTransform(train_aug))
# set up tuning data
tune_trans = TuneDataAugmentation(s, input_shape)
tune_aug = tune_trans.augment()
tune_data = datasets.ImageFolder('labelled', transform=TuneDataTransform(tune_aug))
# # set up val data
# val_data = tv.datasets.CIFAR10(
# root='Data', train=False, download=True,
# transform=TuneDataTransform(tune_aug)
# )
val_data = None
# create the base model
model = ResNetSimCLR(in_channels=3, n_classes=10, mlp_layers=2)
simclr = SimCLR(model)
print(f'Our model has {simclr.trainer.num_params:,} parameters')
results = simclr.full_model_maker(
train_data, tune_data, val_data, n_cycles=1, train_epochs=5, tune_epochs=2,
train_path='chkpt/train.tar', tune_path='chkpt/tune.tar'
)
return results
def __init__(self,
model=None,
in_channels=3,
n_classes=5,
batch_size=128,
*args,
**kwargs):
if model is None:
# define the model
self.model = ResNetSimCLR(in_channels, n_classes, *args, **kwargs)
else:
# use a pretrained model
self.model = model
self.model = self.model.to(self.device)
self.num_params = sum(p.numel() for p in self.model.parameters())
self.batch_size = batch_size
def finish_the_job():
s = 1
input_shape = (128, 128, 3)
# set up training data
train_trans = TrainDataAugmentation(s, input_shape)
train_aug = train_trans.augment()
train_data = tv.datasets.ImageFolder(
root='Data/unlabeled', transform=SimCLRDataTransform(train_aug))
# set up tuning and validation data
tune_trans = TuneDataAugmentation(s, input_shape)
tune_aug = tune_trans.augment()
tune_val_data = tv.datasets.ImageFolder(
root='Data/supervised', transform=TuneDataTransform(tune_aug))
tune_data, val_data = random_split(tune_val_data, [7112, 16595])
saved = torch.load('chkpt/race_train.tar70')
trained_model = model = ResNetSimCLR(in_channels=3,
n_classes=5,
mlp_layers=2,
blocks_layers=[3, 4, 6, 3])
trained_model.load_state_dict(saved['model_state_dict'])
simclr = SimCLR(trained_model)
model = simclr.trainer.return_model()
simclr.make_tuner(model)
tune_epochs = 50
tune_path = 'chkpt/race_tune.tar'
model, tune_losses = simclr.tune(tune_data, tune_epochs, tune_path)
simclr.make_validator(simclr.tuner.model)
acc, actual, predicted = simclr.validate(val_data)
results = (model, tune_losses, acc, actual, predicted)
return results
def make_tuner(self, model):
'''
Parameters:
model (tuple(ResNet, nn.Linear)): the state dicts of the resnet
and first layer of the projection head of for a simclr model
'''
resnet_dict = model[0]
head_dict = model[1]
in_channels = self.trainer.model.in_channels
n_classes = self.trainer.model.n_classes
model = ResNetSimCLR(in_channels,
self.n_classes,
mlp_layers=3,
blocks_layers=[3, 4, 6, 3])
model.resnet.load_state_dict(resnet_dict)
model.projection_head.layers[0].load_state_dict(head_dict)
self.tuner = FineTune(model, batch_size=self.tune_batch_size)
def train_race_fr():
s = 1
input_shape = (128, 128, 3)
# set up training data
train_trans = TrainDataAugmentation(s, input_shape)
train_aug = train_trans.augment()
train_data = tv.datasets.ImageFolder(
root='Data/unlabeled', transform=SimCLRDataTransform(train_aug))
# set up tuning and validation data
tune_trans = TuneDataAugmentation(s, input_shape)
tune_aug = tune_trans.augment()
tune_val_data = tv.datasets.ImageFolder(
root='Data/supervised', transform=TuneDataTransform(tune_aug))
tune_data, val_data = random_split(tune_val_data, [7112, 16595])
# create the base model
model = ResNetSimCLR(in_channels=3,
n_classes=5,
mlp_layers=2,
blocks_layers=[3, 4, 6, 3])
simclr = SimCLR(model)
print(f'Our model has {simclr.trainer.num_params:,} parameters')
results = simclr.full_model_maker(train_data,
tune_data,
val_data,
n_cycles=1,
train_epochs=71,
tune_epochs=31,
train_path='chkpt/race_train.tar',
tune_path='chkpt/race_tune.tar')
return results
loss.item()))
progress.update()
if clr:
clr.step()
# collect the running loss so that i can see how it's doing
running_loss += loss.item()
if i % 40 == 39:
# print the man loss
losses.append(running_loss / 40)
running_loss = 0.0
return net, losses
# make the net
net = ResNetSimCLR(3,
10,
mlp_layers=3,
blocks_sizes=[2**i for i in [5, 6, 7, 8]],
blocks_layers=[2, 2, 2, 2])
net = net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)
net, loss = train_model(net, trainloader, criterion, optimizer, 10, device)
plt.plot(loss)
plt.savefig('vis/test_test')