def __init__(self,
device: Union[str, torch.device] = "cpu",
load_from_path: str = None,
augmentations: nn.Sequential = None) -> None:
"""Initialise model and learner setup."""
self.device = device
self.model = models.wide_resnet101_2(pretrained=False).to(self.device)
self.timestamp = datetime.now().strftime("%m-%d-%Y_%H-%M-%S")
if load_from_path is not None:
print("Loading model...")
state_dict = torch.load(load_from_path)
self.model.load_state_dict(state_dict)
if augmentations is None:
self.learner = BYOL(self.model,
image_size=64,
hidden_layer="avgpool")
if augmentations is not None:
self.learner = BYOL(self.model,
image_size=64,
hidden_layer="avgpool",
augment_fn=augmentations)
self.opt = torch.optim.Adam(self.learner.parameters(),
lr=0.0001,
betas=(0.9, 0.999))
self.loss_history: list[float] = []
def _init_self_supervised(self) -> None:
"""
Create BYOL network and optimizer.
"""
backbone_metadata = BACKBONE_METADATA_BY_NAME[self.hparams.backbone_name]
net = backbone_metadata.build_fn()
# Transform random_crop_size determines final image size.
assert (
self.data_config.train_transform1.random_crop_size
== self.data_config.train_transform2.random_crop_size
), "Crop size must be the same for all transforms."
assert (
self.data_config.eval_transform.center_crop_size
== self.data_config.train_transform1.random_crop_size
), "Crop size must be the same for all transforms."
image_size = self.data_config.train_transform1.random_crop_size
self.byol_model = cast(
BYOL,
self.context.wrap_model(BYOL(net, image_size)),
)
# Trick to support passing in pair of augmented images to byol_pytorch.
# Normally, byol_pytorch does augmentations inside its forward pass, which is slow.
self.byol_model.augment1 = LambdaModule(lambda x: x.first)
self.byol_model.augment2 = LambdaModule(lambda x: x.second)
self.byol_opt = self.context.wrap_optimizer(
build_byol_optimizer(self.hparams, self.byol_model)
)
def byol_vit(num_classes, loss='softmax', pretrained=True, **kwargs):
vit = vittimmdiet(num_classes, loss=loss, pretrained=pretrained, **kwargs)
torchreid.utils.load_pretrained_weights(
vit,
'/home/danish/deep-person-reid/scripts/log/model/model.pth.tar-11')
augment_fn = nn.Sequential(kornia.augmentation.RandomHorizontalFlip())
augment_fn2 = nn.Sequential(
kornia.augmentation.RandomHorizontalFlip(),
kornia.filters.GaussianBlur2d((3, 3), (1.5, 1.5)))
learner = BYOL(
vit,
image_size=224,
hidden_layer='head',
projection_size=num_classes, # the projection size
projection_hidden_size=
2048, # the hidden dimension of the MLP for both the projection and prediction
moving_average_decay=0.99,
# the moving average decay factor for the target encoder, already set at what paper recommends
augment_fn=augment_fn,
augment_fn2=augment_fn2,
)
#print(learner)
return learner
def train_resnet_by_byol(args, train_loader):
# define cnn(resnet) and byol
resnet = resnet50(pretrained=True).to(args.device)
learner = BYOL(resnet, image_size = args.height, hidden_layer = 'avgpool')
opt = torch.optim.Adam(learner.parameters(), lr=args.lr)
# train resnet via BYOL
print("BYOL training start -- ")
for epoch in range(args.epoch):
loss_history=[]
for data, _ in train_loader:
images = data.to(args.device)
loss = learner(images)
loss_float = loss.detach().cpu().numpy().tolist()
loss_history.append(loss_float)
opt.zero_grad()
loss.backward()
opt.step()
learner.update_moving_average() # update moving average of target encoder
if epoch % 5 == 0:
print(f"EPOCH: {epoch} / loss: {sum(loss_history)/len(train_loader)}")
return resnet
def get_learner(config):
resnet = torchvision.models.resnet18(pretrained=True)
learner = BYOL(resnet, image_size=32, hidden_layer='avgpool')
opt = torch.optim.Adam(learner.parameters(), lr=config.lr)
learner = learner.cuda()
return learner
def __init__(
self, net, train_dataset: Dataset, valid_dataset: Dataset,
epochs: int, learning_rate: float,
batch_size: int = 32, num_gpus: int = 1, **kwargs):
super().__init__()
self.batch_size = batch_size
self.train_dataset = train_dataset
self.valid_dataset = valid_dataset
self.learning_rate = learning_rate
self.epochs = epochs
self.learner = BYOL(net, **kwargs)
self.num_gpus = num_gpus
def run_train(dataloader):
if torch.cuda.is_available():
dev = "cuda:0"
else:
dev = "cpu"
device = torch.device(dev)
model = AlexNet().to(device)
learner = BYOL(
model,
image_size=32,
)
opt = torch.optim.Adam(learner.parameters(), lr=3e-4)
for epoch in range(EPOCH):
for step, (images, _) in tqdm(
enumerate(dataloader),
total=int(60000 / BATCH_SIZE),
desc="Epoch {d}: Training on batch".format(d=epoch)):
images = images.to(device)
loss = learner(images)
opt.zero_grad()
loss.backward()
opt.step()
learner.update_moving_average(
) # update moving average of target encoder
wandb.log({"train_loss": loss.cpu(), "step": step, "epoch": epoch})
# save your improved network
torch.save(model.state_dict(), './improved-net.pt')
def __init__(self, net, **kwargs):
super().__init__()
self.learner = BYOL(net, **kwargs)
import torch
from byol_pytorch import BYOL
from torchvision import models
resnet = models.resnet50(pretrained=True)
learner = BYOL(resnet, image_size=256, hidden_layer='avgpool')
opt = torch.optim.Adam(learner.parameters(), lr=3e-4)
def sample_unlabelled_images():
return torch.randn(20, 3, 256, 256)
for _ in range(100):
images = sample_unlabelled_images()
loss = learner(images)
opt.zero_grad()
loss.backward()
opt.step()
learner.update_moving_average() # update moving average of target encoder
# save your improved network
torch.save(resnet.state_dict(), './improved-net.pt')
def __init__(self, net, **kwargs):
super().__init__()
self.learner = BYOL(net, **kwargs)
self.last = nn.Sequential(nn.Linear(1000, 2), nn.Softmax(dim=1))
if args.depth == 18:
model = models.resnet18(pretrained=False).cuda()
elif args.depth == 34:
model = models.resnet34(pretrained=False).cuda()
elif args.depth == 50:
model = models.resnet50(pretrained=False).cuda()
elif args.depth == 101:
model = models.resnet101(pretrained=False).cuda()
else:
assert ("Not supported Depth")
learner = BYOL(
model,
image_size=args.image_size,
hidden_layer='avgpool',
projection_size=256,
projection_hidden_size=4096,
moving_average_decay=0.99,
use_momentum=False # turn off momentum in the target encoder
)
opt = torch.optim.Adam(learner.parameters(), lr=args.lr)
ds = ImagesDataset(args.image_folder, args.image_size)
trainloader = DataLoader(ds, batch_size=args.batch_size, num_workers=NUM_WORKERS, shuffle=True)
losses = AverageMeter()
for epoch in range(args.epoch):
bar = Bar('Processing', max=len(trainloader))
for batch_idx, inputs in enumerate(trainloader):
loss = learner(inputs.cuda())
#
# model = CNN(n_classes, 'mobilenet_v2')
#
# model = ViT(
# image_size=224,
# patch_size=32,
# num_classes=3,
# dim=512,
# depth=6,
# heads=16,
# mlp_dim=1024,
# dropout=0.1,
# emb_dropout=0.1
# )
print(model)
learner = BYOL(model, image_size=224, hidden_layer='fc1')
opt = torch.optim.Adam(learner.parameters(), lr=1e-4)
def sample_unlabelled_images():
return torch.randn(32, 3, 224, 224)
for i in range(200):
images = sample_unlabelled_images()
loss = learner(images.cuda())
opt.zero_grad()
loss.backward()
if i % 10 == 0:
print(i, loss.item())
#define the multihead classifier, that has one output layer per task id
multi_head = Multi_head(cil_step)
multi_head.to(device)
# add list of ssl models to pass for training
ssl_models = []
ssl_dict = {}
for method in ssl_methods:
if 'byol' in ssl_methods:
byol_model = BYOL(
encoder,
image_size=32,
hidden_layer= -1
)
byol_model.to(device)
ssl_dict['byol'] = byol_model
if 'xxx' in ssl_methods: # add extra ssl model here
pass
if 'ae' in ssl_methods: # add extra ssl model here
ae = AE(encoder=encoder)
ae.to(device)
ssl_dict['ae'] = ae
os.makedirs(model_dir, exist_ok=True)
data_dir = os.path.join(root_dir, 'train')
csv_path = os.path.join(root_dir, 'train.csv')
if __name__ == '__main__':
train_set = MiniDataset(csv_path, data_dir)
train_loader = DataLoader(train_set,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
shuffle=True)
resnet = models.resnet50(pretrained=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
learner = BYOL(resnet, image_size=IMAGE_SIZE, hidden_layer='avgpool')
learner.to(device)
opt = torch.optim.Adam(learner.parameters(), lr=LR)
learner.train()
for epoch in range(EPOCHS):
for i, data in enumerate(train_loader):
data = data.to(device)
loss = learner(data)
opt.zero_grad()
loss.backward()
opt.step()
learner.update_moving_average(
) # update moving average of target encoder
import torch
from byol_pytorch import BYOL
from byol_pytorch.model import FCN, Predictor
from byol_pytorch.dataset import SegDataset
encoder = FCN(input_dim=3, output_classes=24)
predictor = Predictor(dim=24)
learner = BYOL(encoder=encoder,
predictor=predictor,
image_size=1024,
hidden_layer='avgpool',
use_momentum=False)
opt = torch.optim.Adam(learner.parameters(), lr=3e-4)
imagenet = SegDataset(root_dir="/images/HER2/images")
dataloader = torch.utils.data.DataLoader(imagenet,
batch_size=2,
shuffle=True,
num_workers=0)
device = torch.device(2)
for _ in range(100):
print("start epoch, ", _)
for local_batch in dataloader:
local_batch = local_batch.to(device)
loss = learner(local_batch)
opt.zero_grad()
loss.backward()
opt.step()
def __init__(self, net, dataset: Dataset, batch_size: int = 32, **kwargs):
super().__init__()
self.batch_size = batch_size
self.dataset = dataset
self.learner = BYOL(net, **kwargs)