Image-SSL/simclr.py at main · Samarth26/Image-SSL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
import logging
import os
import sys

import torch
import torch.nn.functional as F
from torch.cuda.amp import GradScaler
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from utils import save_config_file, accuracy, save_checkpoint
import wandb

torch.manual_seed(0)


class SimCLR(object):

    def __init__(self, *args, **kwargs):
        self.args = kwargs['args']
        self.model = kwargs['model'].to(self.args.device)
        self.optimizer = kwargs['optimizer']
        self.scheduler = kwargs['scheduler']

        # Create a timestamped directory inside the log_dir
        import time
        timestamp = time.strftime('%b%d_%H-%M-%S')
        if self.args.log_dir:
            self.writer = SummaryWriter(log_dir=os.path.join(self.args.log_dir, timestamp))
        else:
            self.writer = SummaryWriter() # Default behavior (runs/timestamp_hostname)

        if self.args.enable_wandb:
            wandb.init(project=self.args.wandb_project, entity=self.args.wandb_entity, config=self.args, dir=self.writer.log_dir)
        logging.basicConfig(filename=os.path.join(self.writer.log_dir, 'training.log'), level=logging.DEBUG)
        self.criterion = torch.nn.CrossEntropyLoss().to(self.args.device)

    def info_nce_loss(self, features):

        labels = torch.cat([torch.arange(self.args.batch_size) for i in range(self.args.n_views)], dim=0)
        labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()
        labels = labels.to(self.args.device)

        features = F.normalize(features, dim=1)

        similarity_matrix = torch.matmul(features, features.T)
        # assert similarity_matrix.shape == (
        #     self.args.n_views * self.args.batch_size, self.args.n_views * self.args.batch_size)
        # assert similarity_matrix.shape == labels.shape

        # discard the main diagonal from both: labels and similarities matrix
        mask = torch.eye(labels.shape[0], dtype=torch.bool).to(self.args.device)
        labels = labels[~mask].view(labels.shape[0], -1)
        similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)
        # assert similarity_matrix.shape == labels.shape

        # select and combine multiple positives
        positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)

        # select only the negatives the negatives
        negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)

        logits = torch.cat([positives, negatives], dim=1)
        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(self.args.device)

        logits = logits / self.args.temperature
        return logits, labels

    def train(self, train_loader):

        scaler = GradScaler(enabled=self.args.fp16_precision)

        # save config file
        save_config_file(self.writer.log_dir, self.args)

        n_iter = 0
        logging.info(f"Start SimCLR training for {self.args.epochs} epochs.")
        logging.info(f"Training with gpu: {self.args.disable_cuda}.")

        best_loss = float('inf')

        for epoch_counter in range(self.args.epochs):
            epoch_loss = 0.0
            pbar = tqdm(train_loader, desc=f"Epoch {epoch_counter+1}/{self.args.epochs}")
            for images, _ in pbar:
                images = torch.cat(images, dim=0)

                images = images.to(self.args.device)

                with torch.amp.autocast('cuda', enabled=self.args.fp16_precision):
                    features = self.model(images)
                    logits, labels = self.info_nce_loss(features)
                    loss = self.criterion(logits, labels)

                self.optimizer.zero_grad()

                scaler.scale(loss).backward()

                scaler.step(self.optimizer)
                scaler.update()

                epoch_loss += loss.item()

                # warmup for the first 10 epochs
                if n_iter >= 10 * len(train_loader): # Assuming len(train_loader) is steps per epoch
                    self.scheduler.step()

                if n_iter % self.args.log_every_n_steps == 0:
                    top1, top5 = accuracy(logits, labels, topk=(1, 5))
                    self.writer.add_scalar('loss', loss, global_step=n_iter)
                    self.writer.add_scalar('acc/top1', top1[0], global_step=n_iter)
                    self.writer.add_scalar('acc/top5', top5[0], global_step=n_iter)
                    self.writer.add_scalar('learning_rate', self.scheduler.get_lr()[0], global_step=n_iter)
                    if self.args.enable_wandb:
                        wandb.log({
                            'loss': loss.item(),
                            'acc/top1': top1[0].item(),
                            'acc/top5': top5[0].item(),
                            'learning_rate': self.scheduler.get_lr()[0],
                            'epoch': epoch_counter
                        }, step=n_iter)

                # Update progress bar with current loss
                pbar.set_postfix({'loss': f'{loss.item():.4f}'})

                n_iter += 1

            epoch_loss /= len(train_loader)
            logging.debug(f"Epoch: {epoch_counter}\tLoss: {epoch_loss}")

            # Save checkpoint if loss improved
            if epoch_loss < best_loss:
                best_loss = epoch_loss
                checkpoint_name = 'checkpoint_{:04d}.pth.tar'.format(epoch_counter + 1)
                save_checkpoint({
                    'epoch': self.args.epochs,
                    'arch': self.args.arch,
                    'state_dict': self.model.state_dict(),
                    'optimizer': self.optimizer.state_dict(),
                }, is_best=True, filename=os.path.join(self.writer.log_dir, checkpoint_name))
                logging.info(f"Best model saved at epoch {epoch_counter + 1} with loss {best_loss:.4f}")

            # Save checkpoint every 10 epochs or at the last epoch
            if (epoch_counter + 1) % 10 == 0 or (epoch_counter + 1) == self.args.epochs:
                checkpoint_name = 'checkpoint_{:04d}.pth.tar'.format(epoch_counter + 1)
                save_checkpoint({
                    'epoch': self.args.epochs,
                    'arch': self.args.arch,
                    'state_dict': self.model.state_dict(),
                    'optimizer': self.optimizer.state_dict(),
                }, is_best=False, filename=os.path.join(self.writer.log_dir, checkpoint_name))
                logging.info(f"Model checkpoint saved at {os.path.join(self.writer.log_dir, checkpoint_name)}")

        logging.info("Training has finished.")