joonmyung 1.5.2__tar.gz → 1.5.5__tar.gz

{joonmyung-1.5.2 → joonmyung-1.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: joonmyung
-Version: 1.5.2
+Version: 1.5.5
 Summary: JoonMyung's Library
 Home-page: https://github.com/pizard/JoonMyung.git
 Author: JoonMyung Choi

{joonmyung-1.5.2 → joonmyung-1.5.5}/joonmyung/analysis/analysis.py

@@ -5,7 +5,7 @@ from joonmyung.analysis.model import JModel
 from joonmyung.draw import saliency, overlay, drawImgPlot, drawHeatmap, unNormalize
 from joonmyung.meta_data import data2path
 from joonmyung.data import getTransform
-from joonmyung.metric import targetPred
+from joonmyung.metric import targetPred, accuracy
 from joonmyung.log import AverageMeter
 from joonmyung.utils import to_leaf, to_np
 from tqdm import tqdm
@@ -20,13 +20,15 @@ import cv2
 
 def anaModel(transformer_class):
     class VisionTransformer(transformer_class):
+        def has_parameter(self, parameter_name):
+            return parameter_name in self.__init__.__code__.co_varnames
+
         def forward_features(self, x):
             x = self.patch_embed(x)
-            cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
-            if self.dist_token is None:
+            if self.has_parameter("cls_token"):
+                cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
                 x = torch.cat((cls_token, x), dim=1)
-            else:
-                x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
+
 
 
             if self.analysis[0] == 1:   # PATCH
@@ -43,17 +45,19 @@ def anaModel(transformer_class):
 
             x = self.blocks(x)
             x = self.norm(x)
-            if self.dist_token is None:
+            if self.has_parameter("cls_token") and self.has_parameter("dist_token"):
+                return x[:, 0], x[:, 1]
+            elif self.has_parameter("cls_token"):
                 return self.pre_logits(x[:, 0])
             else:
-                return x[:, 0], x[:, 1]
+                return self.pre_logits(x.mean(dim=1))
+
 
     return VisionTransformer
 
 class Analysis:
-    def __init__(self, model, analysis = [0], activate = [True, False, False], detach=True, key_name=None, num_classes = 1000
+    def __init__(self, model, analysis = [0], activate = [True, False, False, False], detach=True, key_name=None, num_classes = 1000
                  , cls_start=0, cls_end=1, patch_start=1, patch_end=None
-                 , ks = 5
                  , amp_autocast=suppress, device="cuda"):
         # Section A. Model
         self.num_classes = num_classes
@@ -63,18 +67,17 @@ class Analysis:
         model.__class__ = model_
         model.analysis = analysis
         self.model = model
-        self.ks = ks
         self.detach = detach
 
         # Section B. Attention
         self.kwargs_roll = {"cls_start" : cls_start, "cls_end" : cls_end,
                             "patch_start" : patch_start, "patch_end" : patch_end}
 
-
         # Section C. Setting
         hooks = [{"name_i": 'attn_drop', "name_o": 'decoder', "fn_f": self.attn_forward, "fn_b": self.attn_backward},
                  {"name_i": 'qkv', "name_o": 'decoder', "fn_f": self.qkv_forward, "fn_b": self.qkv_backward},
-                 {"name_i": 'head', "name_o": 'decoder', "fn_f": self.head_forward, "fn_b": self.head_backward}]
+                 {"name_i": 'head', "name_o": 'decoder', "fn_f": self.head_forward, "fn_b": self.head_backward},
+                 {"name_i": 'patch_embed.norm', "name_o": 'decoder', "fn_f": self.input_forward, "fn_b": self.input_backward}]
         hooks = [h for h, a in zip(hooks, activate) if a]
 
 
@@ -86,20 +89,18 @@ class Analysis:
                 if hook["name_i"] in name and hook["name_o"] not in name:
                     module.register_forward_hook(hook["fn_f"])
                     module.register_backward_hook(hook["fn_b"])
+        self.resetInfo()
 
     def attn_forward(self, module, input, output):
-        # input  : 1 * (8, 3, 197, 197)
-        # output : (8, 3, 197, 197)
-        self.info["attn"]["f"].append(output.detach() if self.detach else output)
+        # input/output : 1 * (8, 3, 197, 197) / (8, 3, 197, 197)
+        self.info["attn"]["f"] = output.detach() if self.detach else output
 
     def attn_backward(self, module, grad_input, grad_output):
-        # input  : 1 * (8, 3, 197, 192)
-        # output : (8, 3, 197, 576)
-        self.info["attn"]["b"].insert(0, grad_input[0].detach() if self.detach else grad_input[0])
+        # input/output : 1 * (8, 3, 197, 192) / (8, 3, 197, 576)
+        self.info["attn"]["b"] = grad_input[0].detach() if self.detach else grad_input[0]
 
     def qkv_forward(self, module, input, output):
-        # input  : 1 * (8, 197, 192)
-        # output : (8, 197, 576)
+        # input/output : 1 * (8, 197, 192) / (8, 197, 576)
         self.info["qkv"]["f"].append(output.detach())
 
     def qkv_backward(self, module, grad_input, grad_output):
@@ -108,38 +109,50 @@ class Analysis:
 
     def head_forward(self, module, input, output):
         # input : 1 * (8(B), 192(D)), output : (8(B), 1000(C))
-        # TP = targetPred(output.detach(), self.targets.detach(), topk=self.ks)
-        TP = targetPred(to_leaf(output), to_leaf(self.targets), topk=self.ks)
-        self.info["head"]["TP"] = torch.cat([self.info["head"]["TP"], TP], dim=0) if "TP" in self.info["head"].keys() else TP
+        B = output.shape[0]
+        pred = targetPred(output, self.targets, topk=5)
+        self.info["head"]["TF"] += (pred[:, 0] == pred[:, 1])
+
+        acc1, acc5 = accuracy(output, self.targets, topk=(1,5))
+        self.info["head"]["acc1"].update(acc1.item(), n=B)
+        self.info["head"]["acc5"].update(acc5.item(), n=B)
 
     def head_backward(self, module, grad_input, grad_output):
         pass
 
-    def resetInfo(self):
-        self.info = {"attn": {"f": [], "b": []}, "qkv": {"f": [], "b": []},
-                     "head": {"acc1" : AverageMeter(),
-                              "acc5" : AverageMeter(),
-                              "pred" : None
+    def input_forward(self, module, input, output):
+        norm = F.normalize(output, dim=-1)
+        self.info["input"]["sim"] += (norm @ norm.transpose(-1, -2)).mean(dim=(-1, -2))
 
-                              }}
+    def input_backward(self, module, grad_input, grad_output):
+        pass
 
-    def __call__(self, samples, index=None, **kwargs):
-        self.resetInfo()
+    def resetInfo(self):
+        self.info = {"attn" : {"f": None, "b": None},
+                     "qkv"  : {"f": None, "b": None},
+                     "head" : {"acc1" : AverageMeter(),
+                               "acc5" : AverageMeter(),
+                               "TF"   : [], "pred" : []},
+                     "input": {"sim" : []}
+                     }
+
+    def __call__(self, samples, targets = None, **kwargs):
         self.model.zero_grad()
         self.model.eval()
 
         if type(samples) == torch.Tensor:
+            self.targets = targets
             outputs = self.model(samples, **kwargs)
             return outputs
         else:
-            for sample, self.targets in tqdm(samples):
+            for sample, targets in tqdm(samples):
+                self.targets = targets
                 _ = self.model(sample)
             return False
 
     def anaSaliency(self, attn=True, grad=False, output=None, index=None,
                     head_fusion="mean", discard_ratios=[0.], data_from="cls",
-                    ls_attentive=[], ls_rollout=[],
-                    reshape=False, device="cuda"):
+                    reshape=False, activate= [True, True, False], device="cuda"):
 
         if attn:
             attn = self.info["attn"]["f"]
@@ -152,31 +165,29 @@ class Analysis:
             loss.backward(retain_graph=True)
             grad = self.info["attn"]["b"]
 
-        return saliency(attn, grad
-                        , head_fusion=head_fusion, discard_ratios=discard_ratios, data_from=data_from
-                        , ls_rollout=ls_rollout, ls_attentive=ls_attentive, reshape=reshape, device=device)
+        return saliency(attn, grad, activate=activate,
+                        head_fusion=head_fusion, discard_ratios=discard_ratios, data_from=data_from,
+                        reshape=reshape, device=device)
 
 
 if __name__ == '__main__':
     # Section A. Data
     dataset_name, device, amp_autocast, debug = "imagenet", 'cuda', torch.cuda.amp.autocast, True
     data_path, num_classes, _, _ = data2path(dataset_name)
-    view, activate = [False, True, False, False, False], [True, False, False]
+    view, activate = [False, True, False, False, True], [True, False, False]
         # VIEW     : IMG, SALIENCY:ATTN, SALIENCY:OPENCV, SALIENCY:GRAD, ATTN. MOVEMENT
         # ACTIVATE : ATTN, QKV, HEAD
     analysis = [0] # [0] : INPUT TYPE, [0 : SAMPLE + POS, 1 : SAMPLE, 2 : POS]
 
     dataset = JDataset(data_path, dataset_name, device=device)
-    data_idxs = [[c, i] for i in range(1000) for c in range(50)]
-    # data_idxs = [[21, 0], [22, 0], [2, 0], [0, 0], [0, 1], [1, 1], [2, 1], [3, 1]]
+    # data_idxs = [[c, i] for i in range(1000) for c in range(50)]
+    data_idxs = [[1, 0]]
 
     # Section B. Model
     model_number, model_name = 0, "deit_tiny_patch16_224" # deit, vit | tiny, small, base
     # model_number, model_name = 1, "deit_tiny_patch16_224"
 
     # Section C. Setting
-    ls_rollout, ls_attentive, col = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 12
-
     modelMaker = JModel(num_classes, device=device)
     model = modelMaker.getModel(model_number, model_name)
     model = Analysis(model, analysis = analysis, activate = activate, device=device)
@@ -190,24 +201,36 @@ if __name__ == '__main__':
             drawImgPlot(unNormalize(sample, "imagenet"))
 
         if view[1]: # SALIENCY W/ MODEL
-            # ls_rollout, ls_attentive, col = [], [0,2,4,6,8,10], 6
-            # discard_ratios, v_ratio, head_fusion, data_from = [0.0, 0.4, 0.8], 0.1, "mean", "cls"
-            discard_ratios, v_ratio, head_fusion, data_from = [0.0], 0.0, "mean", "patch"  # Attention, Gradient
-            rollout, attentive = model.anaSaliency(True, False, output, discard_ratios=discard_ratios,
-                                                   ls_attentive = ls_attentive, ls_rollout=ls_rollout,
+            col, discard_ratios, v_ratio, head_fusion, data_from = 12, [0.0], 0.0, "mean", "patch"  # Attention, Gradient
+            results = model.anaSaliency(True, False, output, discard_ratios=discard_ratios,
                                                    head_fusion  = head_fusion, index=target, data_from=data_from,
-                                                   reshape      = True) # (12(L), 8(B), 14(H), 14(W))
-            datas_attn    = overlay(sample, attentive, dataset_name)
-            drawImgPlot(datas_attn, col=col)
+                                                   reshape      = True, activate=[True, True, True]) # (12(L), 8(B), 14(H), 14(W))
+            data_roll = overlay(sample, results["rollout"], dataset_name)
+            drawImgPlot(data_roll, col=col)
 
-            discard_ratios, v_ratio, head_fusion, data_from = [0.0], 0.0, "mean", "cls"  # Attention, Gradient
-            rollout, attentive = model.anaSaliency(True, False, output, discard_ratios=discard_ratios,
-                                                   ls_attentive = ls_attentive, ls_rollout=ls_rollout,
-                                                   head_fusion  = head_fusion, index=target, data_from=data_from,
-                                                   reshape      = True) # (12(L), 8(B), 14(H), 14(W))
-            datas_attn = overlay(sample, attentive, dataset_name)
-            drawImgPlot(datas_attn, col=col)
+            data_attn    = overlay(sample, results["attentive"], dataset_name)
+            drawImgPlot(data_attn, col=col)
+
+            data_vidTLDR = overlay(sample, results["vidTLDR"], dataset_name)
+            drawImgPlot(data_vidTLDR, col=col)
+
+            discard_ratios, v_ratio, head_fusion, data_from = [0.0], 0.1, "mean", "cls"
+            results = model.anaSaliency(True, False, output, discard_ratios=discard_ratios,
+                                        head_fusion=head_fusion, index=target, data_from=data_from,
+                                        reshape=True, activate=[True, True, True])  # (12(L), 8(B), 14(H), 14(W))
+
+            data_roll = overlay(sample, results["rollout"], dataset_name)
+            drawImgPlot(data_roll, col=col)
+
+            data_attn = overlay(sample, results["attentive"], dataset_name)
+            drawImgPlot(data_attn, col=col)
 
+            data_vidTLDR = overlay(sample, results["vidTLDR"], dataset_name)
+            drawImgPlot(data_vidTLDR, col=col)
+
+            print(1)
+
+            # roll = F.normalize(results["rollout"].reshape(12, 196), dim=-1)
 
             # datas_rollout = overlay(sample, rollout,   dataset_name)
             # drawImgPlot(datas_rollout, col=col)
@@ -248,10 +271,13 @@ if __name__ == '__main__':
 
         if view[4]: # ATTENTION MOVEMENT (FROM / TO)
             attn = torch.stack(model.info["attn"]["f"]).mean(dim=2).transpose(0,1) # (8 (B), 12 (L), 197(T_Q), 197(T_K))
+
+            # CLS가 얼마나 참고하는지
             cls2cls     = attn[:, :, :1, 0].mean(dim=2)              # (8(B), 12(L))
             patch2cls   = attn[:, :, :1, 1:].mean(dim=2).sum(dim=-1) # (8(B), 12(L))
-            # PATCH가 받는 정도
+
+            # PATCH가 얼마나 참고하는지
             cls2patch   = attn[:, :, 1:, 0].mean(dim=2)
             patch2patch = attn[:, :, 1:, 1:].mean(dim=2).sum(dim=-1)
             # to_np(torch.stack([cls2cls.mean(dim=0), patch2cls.mean(dim=0), cls2patch.mean(dim=0), patch2patch.mean(dim=0)]))
-            print(1)
+            print(1)

{joonmyung-1.5.2 → joonmyung-1.5.5}/joonmyung/analysis/dataset.py

@@ -1,4 +1,6 @@
+from torchvision.transforms import InterpolationMode
 from joonmyung.meta_data.label import imnet_label, cifar_label
+from torchvision.datasets.folder import default_loader
 from timm.data import create_dataset, create_loader
 from torchvision import transforms
 from joonmyung.utils import getDir
@@ -43,30 +45,32 @@ class JDataset():
         size                = size if size else setting["size"]
 
         self.transform = [
-                transforms.Compose([transforms.Resize(size, interpolation=3), transforms.ToTensor(), transforms.Normalize(self.distribution["mean"], self.distribution["std"])]),
-                transforms.Compose([transforms.Resize(size, interpolation=3), transforms.ToTensor()]),
+                transforms.Compose([transforms.Resize((256, 256), interpolation=InterpolationMode.BICUBIC), transforms.CenterCrop(size), transforms.ToTensor(), transforms.Normalize(self.distribution["mean"], self.distribution["std"])]),
+                transforms.Compose([transforms.Resize((256, 256), interpolation=InterpolationMode.BICUBIC), transforms.CenterCrop(size), transforms.ToTensor()]),
+                transforms.Compose([transforms.Resize(size, interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(self.distribution["mean"], self.distribution["std"])]),
                 transforms.Compose([transforms.ToTensor()])
         ]
 
         self.device = device
         self.data_path = data_path
         self.label_paths = sorted(getDir(os.path.join(self.data_path, self.data_type)))
-        self.img_paths   = [sorted(glob.glob(os.path.join(self.data_path, self.data_type, label_path, "*"))) for label_path in self.label_paths]
+
         # self.img_paths   = [sorted(glob.glob(os.path.join(self.data_path, self.data_type, "*", "*")))]
+        self.img_paths   = [[path, idx] for idx, label_path in enumerate(self.label_paths) for path in sorted(glob.glob(os.path.join(self.data_path, self.data_type, label_path, "*")))]
 
 
     def __getitem__(self, idx):
-        if len(idx) == 2:
-            label_num, img_num = idx
+        if type(idx) == tuple:
+            idx, transform_type = idx
+        else:
             transform_type = self.transform_type
-        elif len(idx) == 3:
-            label_num, img_num, transform_type = idx
 
-        img_path = self.img_paths[label_num][img_num]
-        img = PIL.Image.open(img_path)
-        data = self.transform[transform_type](img)
+        [img_path, targets] = [idx, 0] if type(idx) == str else self.img_paths[idx]
+
+        sample = default_loader(img_path)
+        sample = self.transform[transform_type](sample)
 
-        return data.unsqueeze(0).to(self.device), torch.tensor(label_num).to(self.device), self.label_name[int(label_num)]
+        return sample[None].to(self.device), torch.tensor(targets).to(self.device), self.label_name[targets]
 
     def getItems(self, indexs):
         ds, ls, lns = [], [], []
@@ -77,10 +81,12 @@ class JDataset():
             lns.append(ln)
         return torch.cat(ds, dim=0), torch.stack(ls, dim=0), lns
 
-    def getItemPath(self, img_path):
-        img = PIL.Image.open(img_path)
-        data = self.transform(img)
-        return data.unsqueeze(0).to(self.device)
+    def getIndex(self, c: list = [0, 1000], i: list = [0, 50]):
+        [c_s, c_e], [i_s, i_e] = c, i
+        c = torch.arange(c_s, c_e).reshape(-1, 1).repeat(1, i_e - i_s).reshape(-1)
+        i = torch.arange(i_s, i_e).reshape(1, -1).repeat(c_e - c_s, 1).reshape(-1)
+        c_i = torch.stack([c, i], dim=-1)
+        return c_i
 
     def __len__(self):
         return

{joonmyung-1.5.2 → joonmyung-1.5.5}/joonmyung/analysis/model.py

@@ -1,7 +1,7 @@
 from collections import OrderedDict
-
 from joonmyung.utils import isDir
 from timm import create_model
+import models.deit
 import torch
 import os
 
@@ -11,7 +11,8 @@ class JModel():
         self.num_classes = num_classes
 
         if model_path:
-            self.model_path = os.path.join(model_path, "checkpoint_{}.pth")
+            self.model_path = os.path.join(model_path, "checkpoint.pth")
+
         if p and model_path:
             print("file list : ", sorted(os.listdir(model_path), reverse=True))
         self.device = device
@@ -27,6 +28,27 @@ class JModel():
             model = create_model(model_name, pretrained=True, num_classes=self.num_classes, in_chans=3, global_pool=None, scriptable=False)
         elif model_type == 1:
             model = torch.hub.load('facebookresearch/deit:main', model_name, pretrained=True)
+        elif model_type == 2:
+            checkpoint = torch.load(self.model_path, map_location='cpu')
+            args = checkpoint['args']
+            model = create_model(
+                        args.model,
+                        pretrained=args.pretrained,
+                        num_classes=args.nb_classes,
+                        drop_rate=args.drop,
+                        drop_path_rate=args.drop_path,
+                        drop_block_rate=None,
+                        img_size=args.input_size,
+                        token_nums=args.token_nums,
+                        embed_type=args.embed_type,
+                        model_type=args.model_type
+                    ).to(self.device)
+            state_dict = []
+            for n, p in checkpoint['model'].items():
+                if "total_ops" not in n and "total_params" not in n:
+                    state_dict.append((n, p))
+            state_dict = dict(state_dict)
+            model.load_state_dict(state_dict)
         else:
             raise ValueError
         model.eval()

{joonmyung-1.5.2 → joonmyung-1.5.5}/joonmyung/draw.py

@@ -157,8 +157,9 @@ def drawBarChart(df, x, y, splitColName, col=1, title=[], fmt=1, p=False, c=Fals
 
 @torch.no_grad()
 def saliency(attentions=None, gradients=None, head_fusion="mean",
-             discard_ratios = [0.0], data_from="cls", ls_rollout=[], ls_attentive=[],
-             reshape=False, device="cpu"):
+             discard_ratios = [0.0], data_from="cls", reshape=False,
+             activate = [True, True, True], device="cpu"):
+
     # attentions : L * (B, H, h, w), gradients : L * (B, H, h, w)
     if type(discard_ratios) is not list: discard_ratios = [discard_ratios]
     saliencys = 1.
@@ -180,14 +181,14 @@ def saliency(attentions=None, gradients=None, head_fusion="mean",
 
     saliencys = saliencys.to(device)
 
-    _, B, _, T = saliencys.shape
+    L, B, _, T = saliencys.shape # (L(12), B(1), T(197), T(197))
     H = W = int(T ** 0.5)
 
-    rollouts, attentive = None, None
-    if ls_rollout:
+    result = {}
+    if activate[0]:
         rollouts, I = [], torch.eye(T, device=device).unsqueeze(0).expand(B, -1, -1)  # (B, 197, 197)
         for discard_ratio in discard_ratios:
-            for start in ls_rollout:
+            for start in range(L):
                 rollout = I
                 for attn in copy.deepcopy(saliencys[start:]):  # (L, B, w, h)
                     # TODO NEED TO CORRECT
@@ -207,20 +208,34 @@ def saliency(attentions=None, gradients=None, head_fusion="mean",
                 rollouts.append(rollout)
         rollouts = torch.stack(rollouts, dim=0)
         rollouts = rollouts[:, :, 1:]
+        rollouts = rollouts / rollouts.sum(dim=-1, keepdim=True)
 
         if reshape:
             rollouts = rollouts.reshape(-1, B, H, W) # L, B, H, W
+        result["rollout"] = rollouts
 
-    if ls_attentive:
+    if activate[1]:
         # attentive = saliencys[ls_attentive, :, 0] \
         #         if data_from == "cls" else saliencys[ls_attentive, :, 1:].mean(dim=2) # (L, B, T)
-        attentive = saliencys[ls_attentive, :, 0] \
-            if data_from == "cls" else saliencys[ls_attentive].mean(dim=2)  # (L, B, T)
+        attentive = saliencys[:, :, 0] \
+            if data_from == "cls" else saliencys.mean(dim=2)  # (L, B, T)
         attentive = attentive[:, :, 1:]
+        attentive = attentive / attentive.sum(dim=-1, keepdim=True)
+
         if reshape:
             attentive = attentive.reshape(-1, B, H, W)
+        result["attentive"] = attentive
+
+    if activate[2]:
+        entropy = (saliencys * torch.log(saliencys)).sum(dim=-1)[:, :, 1:] # (L(12), B(1), T(196))
+        entropy = entropy - entropy.amin(dim=-1, keepdim=True)
+        entropy = entropy / entropy.sum(dim=-1, keepdim=True)
+        if reshape:
+            entropy = entropy.reshape(L, B, H, W)
+        result["vidTLDR"] = entropy
+
 
-    return rollouts, attentive
+    return result
 
 
 

{joonmyung-1.5.2 → joonmyung-1.5.5}/joonmyung/log.py

@@ -49,14 +49,14 @@ class AverageMeter:
 
 class Logger():
     loggers = {}
-    def __init__(self, use_wandb=True, wandb_entity=None, wandb_project=None, wandb_name=None
+    def __init__(self, use_wandb=True, wandb_entity=None, wandb_project=None, wandb_name=None, wandb_tags=None
                  , wandb_watch=False, main_process=True, wandb_id=None, wandb_dir='./'
                  , args=None, model=False):
         self.use_wandb = use_wandb
         self.main_process = main_process
 
         if self.use_wandb and self.main_process:
-            wandb.init(entity=wandb_entity, project=wandb_project, name=wandb_name
+            wandb.init(entity=wandb_entity, project=wandb_project, name=wandb_name, tags=wandb_tags
                        , save_code=True, resume="allow", id = wandb_id, dir=wandb_dir
                        , config=args, settings=wandb.Settings(code_dir="."))
 
@@ -120,7 +120,7 @@ class Logger():
         wandb.finish()
 
     def save(self, model, args, name):
-        if self.main_process:
+        if self.main_process and self.use_wandb:
             path = os.path.join(wandb.run.dir, f"{name}.pth")
             torch.save({"model" : model, "args"  : args}, path)
             wandb.save(path, wandb.run.dir)