# Copyright (c) 2022 Graphcore Ltd. All rights reserved. # -------------------------------------------------------- # Swin Transformer # This file has been modified by Graphcore Ltd. # Copyright (c) 2021 Microsoft # Licensed under The MIT License # The LICENSE referenced above is reproduced below: # MIT License # # Copyright (c) Microsoft Corporation. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE # Written by Ze Liu # -------------------------------------------------------- import os import torch from transformers import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup from transformers import get_constant_schedule def load_pretrained(config, model): checkpoint = torch.load(config.PRETRAINED, map_location="cpu") state_dict = checkpoint["model"] # delete relative_position_index since we always re-init it relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k] for k in relative_position_index_keys: del state_dict[k] # delete relative_coords_table since we always re-init it relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k] for k in relative_position_index_keys: del state_dict[k] # delete attn_mask since we always re-init it attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k] for k in attn_mask_keys: del state_dict[k] # bicubic interpolate relative_position_bias_table if not match relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k] for k in relative_position_bias_table_keys: relative_position_bias_table_pretrained = state_dict[k] relative_position_bias_table_current = model.state_dict()[k] L1, nH1 = relative_position_bias_table_pretrained.size() L2, nH2 = relative_position_bias_table_current.size() if nH1 != nH2: print(f"Error in loading {k}, passing......") else: if L1 != L2: # bicubic interpolate relative_position_bias_table if not match S1 = int(L1**0.5) S2 = int(L2**0.5) relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate( relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2), mode="bicubic", ) state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0) # bicubic interpolate absolute_pos_embed if not match absolute_pos_embed_keys = [k for k in state_dict.keys() if "absolute_pos_embed" in k] for k in absolute_pos_embed_keys: # dpe absolute_pos_embed_pretrained = state_dict[k] absolute_pos_embed_current = model.state_dict()[k] _, L1, C1 = absolute_pos_embed_pretrained.size() _, L2, C2 = absolute_pos_embed_current.size() if C1 != C1: print(f"Error in loading {k}, passing......") else: if L1 != L2: S1 = int(L1**0.5) S2 = int(L2**0.5) absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1) absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2) absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate( absolute_pos_embed_pretrained, size=(S2, S2), mode="bicubic" ) absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1) absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(1, 2) state_dict[k] = absolute_pos_embed_pretrained_resized # check classifier, if not match, then re-init classifier to zero head_bias_pretrained = state_dict["head.bias"] Nc1 = head_bias_pretrained.shape[0] Nc2 = model.head.bias.shape[0] if Nc1 != Nc2: if Nc1 == 21841 and Nc2 == 1000: print("loading ImageNet-22K weight to ImageNet-1K ......") map22kto1k_path = f"dataset/map22kto1k.txt" with open(map22kto1k_path) as f: map22kto1k = f.readlines() map22kto1k = [int(id22k.strip()) for id22k in map22kto1k] state_dict["head.weight"] = state_dict["head.weight"][map22kto1k, :] state_dict["head.bias"] = state_dict["head.bias"][map22kto1k] else: torch.nn.init.constant_(model.head.bias, 0.0) torch.nn.init.constant_(model.head.weight, 0.0) del state_dict["head.weight"] del state_dict["head.bias"] print(f"Error in loading classifier head, re-init classifier head to 0") msg = model.load_state_dict(state_dict, strict=False) print(f"=> loaded successfully '{config.PRETRAINED}'") del checkpoint def get_grad_norm(parameters, norm_type=2): if isinstance(parameters, torch.Tensor): parameters = [parameters] parameters = list(filter(lambda p: p.grad is not None, parameters)) norm_type = float(norm_type) total_norm = 0 for p in parameters: param_norm = p.grad.data.norm(norm_type) total_norm += param_norm.item() ** norm_type total_norm = total_norm ** (1.0 / norm_type) return total_norm def get_lr_scheduler(optimizer, scheduler_type, warmup_steps=None, num_steps=None, last_epoch=-1): if scheduler_type == "linear": scheduler = get_linear_schedule_with_warmup(optimizer, warmup_steps, num_steps) elif scheduler_type == "constant": scheduler = get_constant_schedule(optimizer) elif scheduler_type == "cosine": scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, num_steps, last_epoch=last_epoch) else: raise ValueError("Unknown scheduler_type:", scheduler_type) return scheduler class NameScope: """Create a name scope for a code block. All operators originating from this block will have their names prefixed by the given string. >>> with NameScope("CustomString"): ... y = self.bmm(a, b) ... z = torch.relu(y) """ def __init__(self, name: str): assert isinstance(name, str), "Parameter to NameScope must be a string" self.name = name def __enter__(self): torch.ops.poptorch.push_name_scope(self.name) def __exit__(self, type, value, traceback): torch.ops.poptorch.pop_name_scope() class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self, name, fmt=":f"): self.name = name self.fmt = fmt self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count def __str__(self): fmtstr = "{name}\tval:{val" + self.fmt + "}\tavg:({avg" + self.fmt + "})" return fmtstr.format(**self.__dict__) class CustomOpsNotFoundException(Exception): """Raised when the custom ops .so file is not found.""" pass