{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":12888675,"sourceType":"datasetVersion","datasetId":6093998},{"sourceId":12924544,"sourceType":"datasetVersion","datasetId":8178148},{"sourceId":259267077,"sourceType":"kernelVersion"}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":4.659137,"end_time":"2025-06-14T17:56:59.635778","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-06-14T17:56:54.976641","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%writefile main.py\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, Subset\nfrom torch.optim import AdamW\n\nfrom transformers import get_cosine_schedule_with_warmup  # Cosine schedule with warmup\nimport argparse\nimport mlcrate as mlc\nfrom cmi_utils import score, score_from_int\nfrom types import SimpleNamespace\nimport random\nimport gc\nfrom sklearn.metrics import log_loss\nfrom scipy.special import softmax\nfrom scipy.spatial.transform import Rotation as R\nfrom sklearn.utils.class_weight import compute_class_weight\nfrom copy import deepcopy\nfrom joblib import dump, load\nimport json\nfrom collections import defaultdict\nfrom cmi_utils import *\n\n\n\nKAGGLE = 'kaggle' in os.getcwd()\n\nCOMPETITION = 'cmi-detect-behavior-with-sensor-data'\nif KAGGLE:\n    data_path = f'/kaggle/input/{COMPETITION}/'\n    my_data_path = '/kaggle/input/my-cmi-data/'\nelse:\n    data_path = 'data/'\n    my_data_path = 'my_data/'\n\n\ncfg = SimpleNamespace()\ncfg.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ncfg.debug=False # True # \ncfg.train= True # not KAGGLE # True #  # predict\ncfg.apex=True\ncfg.print_freq=1\navailable_cpus = os.cpu_count()\ncfg.num_workers=1 if KAGGLE else max(1, min(available_cpus - 1, 4))\ncfg.prefix = ''\ncfg.target_gestures = ['Above ear - pull hair','Cheek - pinch skin','Eyebrow - pull hair','Eyelash - pull hair','Forehead - pull hairline','Forehead - scratch','Neck - pinch skin','Neck - scratch']\ncfg.non_target_gestures = ['Write name on leg','Wave hello','Glasses on/off','Text on phone','Write name in air','Feel around in tray and pull out an object','Scratch knee/leg skin','Pull air toward your face','Drink from bottle/cup','Pinch knee/leg skin']\ncfg.gestures = sorted(cfg.target_gestures) + sorted(cfg.non_target_gestures)\ncfg.gesture2id = {v:k for k,v in enumerate(cfg.gestures)}\ncfg.gradient_checkpointing=False\ncfg.scheduler='cosine' # ['linear', 'cosine']\ncfg.batch_scheduler=True\ncfg.num_cycles=0.5\ncfg.min_lr=1e-6\ncfg.eps=1e-6\ncfg.betas=(0.9, 0.999)\ncfg.max_grad_norm=100\ncfg.phase_loss = nn.BCEWithLogitsLoss() # nn.BCELoss() \ncfg.n_fold=5\ncfg.trn_folds= list(range(cfg.n_fold))\ncfg.folds=my_data_path + 'team_folds.csv'\ncfg.head = ''\ncfg.no_decay = True\ncfg.use_groups = ['acc','rot','thm','tof']\ncfg.acc_cols = ['acc_x', 'acc_y', 'acc_z']\ncfg.rot_cols = ['rot_w', 'rot_x', 'rot_y', 'rot_z']\ncfg.imu_fe = False # True\nslurm_job_name = os.getenv(\"SLURM_JOB_NAME\", \"Unknown\")\nprint(f\"SLURM Job Name: {slurm_job_name}\")\ncfg.job_id = slurm_job_name[-1]\ncfg.output_dir = '' if KAGGLE else f'models/{cfg.job_id}/'\ncfg.g_loss_weight = 0\ncfg.max_len = 88\ncfg.batch_size = 32\ncfg.model = 'b0' # 0 initial, 1 public_72, 2 - public_75, , 3 - public_80\n# from public_72 import TwoBranchModel\n# from public_75 import ModelVariant_GRU\n# from public_80 import TwoBranchModel80\n# from public_82 import CMIModel\ncfg.save_history = not KAGGLE\ncfg.image_size = 224\ncfg.tof_win = 16\ncfg.sep = 0\n\ndef seed_everything(seed=42):\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    seed = int(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nfrom torchvision import transforms\nfrom torchvision.models import (efficientnet_b0,efficientnet_b1,efficientnet_b3,\n                                efficientnet_b5,efficientnet_b6,resnet18,\n                                efficientnet_v2_s,efficientnet_v2_m,maxvit_t,\n                                convnext_tiny,convnext_small,convnext_base,\n                                resnet50,mobilenet_v2,squeezenet1_0, squeezenet1_1)\n\ndef save_namespace_json(obj, filename, omit=['device','phase_loss','target_gestures',\n                                        'non_target_gestures','gestures','gesture2id']):\n    omit = omit or []  # Handle None case\n    # Get object attributes and filter out omitted ones\n    obj_dict = {k: v for k, v in vars(obj).items() if k not in omit}\n    with open(filename, 'w') as f:\n        json.dump(obj_dict, f, indent=2)\n\ndef load_namespace_json(filename):\n    with open(filename, 'r') as f:\n        data = json.load(f)\n        return SimpleNamespace(**data)\n\n\ndef should_decay(param_name):\n    \"\"\"Determine if a parameter should have weight decay applied.\"\"\"\n    no_decay_keywords = ['bias','bn','batch_norm','layer_norm',\n                         'group_norm','weight_norm']\n    return not any(nd in param_name.lower() for nd in no_decay_keywords)\n\n\ndef separate_decay_params(module, weight_decay_value):\n    \"\"\"Separate parameters into decay and no-decay groups.\"\"\"\n    decay_params = []\n    no_decay_params = []\n\n    for name, param in module.named_parameters():\n        if not param.requires_grad:\n            continue\n\n        # Check if parameter should have weight decay\n        if should_decay(name):\n            decay_params.append(param)\n        else:\n            no_decay_params.append(param)\n\n    param_groups = []\n    if decay_params:\n        param_groups.append({\n            'params': decay_params,\n            'weight_decay': weight_decay_value\n        })\n    if no_decay_params:\n        param_groups.append({\n            'params': no_decay_params,\n            'weight_decay': 0.0\n        })    \n    return param_groups\n\n\nclass CustomMovingAverage(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Create a convolution layer with custom weights\n        self.conv = nn.Conv2d(3, 3, kernel_size=(5, 1), padding=(2, 0), bias=False)\n        kernel = torch.zeros(3, 3, 5, 1)\n        # Channel 0: Identity (keep unchanged)\n        kernel[0, 0, 2, 0] = 1.0\n        # Channel 1: Moving average window=3\n        kernel[1, 1, 1:4, 0] = 1.0/3.0\n        # Channel 2: Moving average window=5\n        kernel[2, 2, 0:5, 0] = 1.0/5.0\n        # Set the weights (requires detaching from gradients if needed)\n        self.conv.weight.data = kernel\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass CustomEfficientNetB0(torch.nn.Module):\n    def __init__(self, num_classes=18, ver = 'b0', n_clf=0):\n        super().__init__()\n        weights_path = (\"/kaggle/input/cmi-models-new/pretrained/\" if KAGGLE \n                        else 'models/pretrained/')\n        self.n_clf = n_clf\n        if ver == 'b0':\n            self.base_model = efficientnet_b0(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_b0_rwightman-7f5810bc.pth'))\n        elif ver == 'b1':\n            self.base_model = efficientnet_b1(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_b1_rwightman-bac287d4.pth'))\n        elif ver == 'b3':\n            self.base_model = efficientnet_b3(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_b3_rwightman-cf984f9c.pth'))\n        elif ver == 'b5':\n            self.base_model = efficientnet_b5(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_b5_lukemelas-b6417697.pth'))\n        elif ver == 'b6':\n            self.base_model = efficientnet_b6(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_b6_lukemelas-24a108a5.pth'))\n        elif ver == 'v2_s':\n            self.base_model = efficientnet_v2_s(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_v2_s-dd5fe13b.pth'))\n        elif ver == 'v2_m':\n            self.base_model = efficientnet_v2_m(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'efficientnet_v2_m-dc08266a.pth'))\n        elif ver == 'cn_tiny':\n            self.base_model = convnext_tiny(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'convnext_tiny-983f1562.pth'))\n        elif ver == 'cn_small':\n            self.base_model = convnext_small(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'convnext_small-0c510722.pth'))\n        elif ver == 'cn_base':\n            self.base_model = convnext_base(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'convnext_base-6075fbad.pth'))\n        elif ver == 'maxvit_t':\n            self.base_model = maxvit_t(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'maxvit_t-bc5ab103.pth'))\n        elif ver == 'mobilenet_v2':\n            self.base_model = mobilenet_v2(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'mobilenet_v2-b0353104.pth'))\n        elif ver == 'squeezenet1_0':\n            self.base_model = squeezenet1_0(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'squeezenet1_0-b66bff10.pth'))\n\n        if ver == 'squeezenet1_0':\n            self.base_model.classifier[1] = torch.nn.Conv2d(\n                        in_channels=512,out_channels=num_classes,kernel_size=1)\n        elif 'cn_' in ver:\n            self.base_model.classifier[2] = torch.nn.Linear(self.base_model.classifier[2].in_features, num_classes)\n        elif ver == 'maxvit_t':\n            self.base_model.classifier[5] = torch.nn.Linear(self.base_model.classifier[5].in_features, num_classes)\n        elif self.n_clf > 0:\n            in_features = self.base_model.classifier[1].in_features\n            self.base_model.classifier = torch.nn.Identity()\n            self.multi_classifier = MultiClassifier(in_features, num_classes,\n                                                    self.n_clf, dropout_rate=0.2)\n        else:\n            self.base_model.classifier[1] = torch.nn.Linear(self.base_model.classifier[1].in_features, num_classes)\n\n    def get_param_groups(self,low_lr,high_lr,weight_decay=1e-4):\n        feature_layers = list(self.base_model.features.children())\n        num_feature_groups = len(feature_layers)\n        lrs = np.linspace(low_lr, high_lr, num_feature_groups + 1)\n        param_groups = []\n        for lr, layer in zip(lrs,feature_layers):\n            layer_groups = separate_decay_params(layer, weight_decay)\n            for group in layer_groups:\n                group['lr'] = float(lr)\n                param_groups.append(group)\n\n        # Classifier layer\n        classifier_groups = separate_decay_params(\n            self.base_model.classifier[1],weight_decay\n        )\n        for group in classifier_groups:\n            group['lr'] = high_lr\n            param_groups.append(group)\n\n        return param_groups\n\n\n    def forward(self, x):\n        if isinstance(x, dict):\n            x = x['imu']\n        out = self.base_model(x)\n        if self.n_clf > 0:\n            out = self.multi_classifier(out)\n            if not self.training:\n                out = out.mean(dim=0)\n        return out\n\nclass CustomResNet(torch.nn.Module):\n    def __init__(self, num_classes=18, ver = '18'):\n        super().__init__()\n        weights_path = (\"/kaggle/input/cmi-models-new/pretrained/\" if KAGGLE \n                        else 'models/pretrained/')\n        if ver == '18':\n            self.base_model = resnet18(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'resnet18-f37072fd.pth'))\n        elif ver == '50':\n            self.base_model = resnet50(weights=None)\n            self.base_model.load_state_dict(torch.load(\n                weights_path + 'resnet50-0676ba61.pth'))\n        self.base_model.fc = nn.Linear(self.base_model.fc.in_features, num_classes)\n\n    def forward(self, x):\n        if isinstance(x, dict):\n            x = x['imu']\n        return self.base_model(x)\n\n\n# Define the dataset class\nclass CustomDataset(Dataset):\n    def __init__(self, data, cfg, mix_alpha=0, demo=None, # data is dataframe\n                 # quat_rel_stat = []  # list[mean (1,-1),std ]\n                ): \n        self.max_len = cfg.max_len\n        self.imu_fe = cfg.imu_fe\n        self.cfg = cfg\n        self.scale_quant = cfg.scale_quant if hasattr(cfg,'scale_quant') else False\n        self.tof_drop = cfg.tof_drop if hasattr(cfg,'tof_drop') else 0\n\n        if self.scale_quant > 0:\n            rel_quat_stats = np.load(my_data_path + f'rel_quat_stats_{self.scale_quant}.npy')\n            self.rq_median = rel_quat_stats[1:2]\n            self.rq_IQR = rel_quat_stats[2:] - rel_quat_stats[:1]\n            self.clip_quat = cfg.clip_quat\n        self.clip_quat = cfg.clip_quat if hasattr(cfg,'clip_quat') else 0\n        self.imu_image_weight = cfg.imu_image_weight if hasattr(cfg,'imu_image_weight') else 0\n        self.euler = cfg.euler if hasattr(cfg,'euler') else False\n        self.gaf = cfg.gaf if hasattr(cfg,'gaf') else ''\n        self.train_mode = 'gesture' in data.columns\n        if self.train_mode:\n            if args.data_aug > 0:\n                data = pd.concat([data,DataAugmentation(data, args.data_aug,\n                    y_grid=cfg.y_aug_grid,z_grid=cfg.z_aug_grid)], ignore_index=True)\n                # data.to_feather(f'data_aug_{args.data_aug}.feather')\n        self.start,self.end = self.get_bounds(data.sequence_id.values)\n        self.col_groups = {prefix: sorted(data.filter(like=prefix))\n                           for prefix in cfg.use_groups}\n        self.X = {}\n        self.imu_image = cfg.model.lower() in ['b0','b1','b3','b5','b6','v2_s','v2_m','maxvit_t',\n            'cn_tiny','cn_small','cn_base','mobilenet_v2','resnet18','resnet50','squeezenet1_0']\n        if self.cfg.flip_left | self.cfg.flip_rot:\n            self.handedness = (demo[['subject','handedness']].set_index('subject')\n                .reindex(data['subject'].values).handedness.to_numpy('int8'))\n        d=-110\n        self.quat_rot_z = np.array([np.cos(np.deg2rad(d)/2),0,0,\n                                    np.sin(np.deg2rad(d)/2)])\n        for key,cols in self.col_groups.items():\n            if len(cols) > 0:\n                self.X[key] = self.preprocess(data[cols], key, data)\n        self.preprocess_b0 = transforms.Compose([\n            transforms.ToPILImage(),  # Convert numpy array or tensor to PIL image\n            transforms.Resize((cfg.image_size, cfg.image_size)),  # Resize to match EfficientNet input size\n            transforms.Grayscale(num_output_channels=3),  # Replicate grayscale to 3 channels\n            transforms.ToTensor(),  # Convert PIL image to tensor\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize for RGB\n        ])\n        # concat individual images\n        if self.imu_image_weight > 0:\n            sep_size = cfg.sep * 3\n            sizes = {'acc':int(0.5 * (cfg.image_size - sep_size) \n                                      *  self.imu_image_weight),\n                     'thm':1                    \n                    }\n            sizes['tof'] = cfg.image_size - sep_size - 2*sizes['acc'] - 1\n            sizes['rot'] = sizes['acc']\n            self.ind_prep = {\n                k: transforms.Compose([ \n                    transforms.ToPILImage(),\n                    transforms.Resize((cfg.image_size,v)),\n                    transforms.Grayscale(num_output_channels=3),transforms.ToTensor(),\n                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])\n                for k,v in sizes.items()\n                }\n        if self.train_mode:\n            gesture2id = {g:i for i,g in enumerate(cfg.gestures)}\n            self.target = data.groupby('sequence_id',sort=False).gesture.last()\\\n                            .map(gesture2id).to_numpy('int8')\n            self.oh_target = np.eye(len(cfg.gestures))[self.target].astype(np.float32) \n            self.gesture_phase = (data.phase == 'Gesture').to_numpy('int8')\n            self.mix_alpha = mix_alpha\n            self.rng = np.random.default_rng(1)\n            if args.rand_aug:\n                self.preprocess_b0 = transforms.Compose([\n                    transforms.ToPILImage(),  # Convert numpy array or tensor to PIL image\n                    transforms.Resize((self.cfg.image_size, self.cfg.image_size)),  # Resize to match EfficientNet input size\n\n                    # Apply random augmentations\n                    # transforms.RandomHorizontalFlip(p=args.hflip),  # Randomly flip the image horizontally\n                    # transforms.RandomVerticalFlip(p=0.5),    # Randomly flip the image vertically\n                    # transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # Adjust color properties    \n\n                    transforms.Grayscale(num_output_channels=3),  # Replicate grayscale to 3 channels\n                    transforms.ToTensor(),  # Convert PIL image to tensor\n                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize for RGB\n                ])\n\n\n    def __len__(self):\n        return len(self.start)\n\n    def __getitem__(self, idx):\n        start,end = self.start[idx],self.end[idx]\n        start = max(start,end - self.max_len)\n        x = {key: torch.tensor(self.X[key][start:end], dtype=torch.float32)\n             for key in self.X}\n        delta = self.max_len - (end - start)\n        if delta > 0:\n            x = {k: F.pad(v,(0,0,delta,0)) for k,v in x.items()}\n\n        if self.train_mode:\n            x['label'] = torch.tensor(self.target[idx], dtype=torch.long)\n            x['oh_label'] = torch.tensor(self.oh_target[idx], dtype=torch.float32)\n            if self.cfg.g_loss_weight > 0:\n                x['gesture_phase'] = torch.tensor(self.gesture_phase[start:end], dtype=torch.float32)\n                if delta > 0:\n                    x['gesture_phase'] = F.pad(x['gesture_phase'],(delta,0))\n        else:\n            self.tof_drop = 0\n\n        if self.imu_image:\n            if self.imu_image_weight > 0:\n                x['imu'] = self.ind_transform(x) \n            sep = torch.zeros_like(x['rot'][...,:self.cfg.sep]) if self.cfg.sep > 0 else []\n            features = [x['acc'],sep,0.5 * (x['rot'] + 1)]\n            if 'tof' in self.X:\n                features.append(sep)\n                features.append(self.tof_smoothing(x['tof']))\n                features.append(sep)\n                features.append(torch.median(x['thm'],dim=-1,keepdims=True)[0])\n\n                del x['tof']  \n                del x['thm']  \n            if len(self.gaf) > 0:\n                # features in [0,1] gaf output in [-1,1]\n                gaf = []\n                for i in np.arange(0,len(features),2):\n                    gaf.append(np.mean(0.5 * (1 + compute_multivariate_gaf(\n                        (2 * features[i] - 1)[np.newaxis,...], \n                        method=self.gaf))[0],axis=-1))\n                gaf = np.column_stack(gaf).astype(np.float32)\n                features.extend([sep,torch.tensor(gaf)]) \n            features = torch.cat([f for f in features if len(f) > 0],dim=-1)\n            imu = self.preprocess_b0(features)\n            del x['acc'],x['rot']\n            x['imu'] = imu\n        return x\n\n# src_order = [\"acc_x\",\"acc_y\",\"acc_z\",\"acc_mag\",\"acc_mag_jerk\",\n#                 \"rot_w\",\"rot_x\",\"rot_y\",\"rot_z\",\"rot_angle\",\"rot_angle_vel\",\n# ]\n# dst_order = ['acc_x',\"rot_x\",\"acc_y\",\"rot_y\",\"acc_z\",\"rot_z\",\n                # \"rot_w\",\"rot_angle\",\"rot_angle_vel\",\"acc_mag\",\"acc_mag_jerk\",]\n\n    def preprocess(self, df, key, data):\n        if key == 'tof':\n            x = ((df + 1.).clip(0,250).fillna(0) / 256.).to_numpy(np.float16)\n            if self.cfg.flip_tof & (self.handedness==0).any():\n                x[self.handedness==0] = flip_tof(x[self.handedness==0])\n        elif key == 'thm':\n            low,high = 22,35\n            x = (df.fillna(0).clip(low,high) - low).to_numpy(np.float16) / (high - low)\n        elif key == 'acc':\n            low,high = np.array([-24.41,-12.70,-20.52]),np.array([20.70,17.25,17.45])\n            x = df.fillna(0).to_numpy(np.float32)\n            if self.cfg.flip_left:\n                x[self.handedness==0,0] = - x[self.handedness==0,0]\n            if self.imu_fe:\n                acc_mag = np.sqrt((x**2).sum(axis=1))\n                acc_mag_jerk = np.append(0,np.diff(acc_mag))\n                acc_mag_jerk[self.start] = 0\n                x = np.hstack([x,acc_mag.reshape(-1,1),acc_mag_jerk.reshape(-1,1)])\n                low,high = np.append(low,[4.2,-13]),np.append(high,[21.3,13])\n            x = (x.clip(low,high) - low) / (high - low)\n        elif key == 'rot': # df:w,x,y,z\n            if self.cfg.fill_rot_mean:\n                x = df.fillna({'rot_x':-0.119916,'rot_y':-0.059953,'rot_z':-0.188298,\n                               'rot_w':0.360375}).clip(-1,1).to_numpy(np.float32)\n            else:\n                x = df.fillna(0).clip(-1,1).to_numpy(np.float32) # ['rot_w','rot_x','rot_y','rot_z']\n            if self.cfg.flip_rot:\n                proper_left = (~df.isnull().any(axis=1)) & (self.handedness == 0)\n                if proper_left.any():\n                    quat_flip = x[proper_left].copy() \n                    quat_flip = quaternion_multiply(self.quat_rot_z, quat_flip)[:,[1,2,3,0]] \n                    euler_x,euler_y,euler_z = quaternion_to_euler(quat_flip)  \n                    quat_flip = euler_to_quaternion(euler_x, -euler_y, -euler_z)\n                    quat_flip = quat_flip*((quat_flip[:,0:1]>0)*2-1) \n                    x[proper_left] = quat_flip.clip(-1,1)\n            if self.imu_fe:\n                rot_angle = 2 * np.arccos(x[:,0])\n                rot_angle_vel = np.append(0,np.diff(rot_angle)).clip(-1,1)\n                rot_angle_vel[self.start] = 0\n                rot_angle = (2 * rot_angle / np.pi - 1).reshape(-1,1)\n                x = np.hstack([x,rot_angle,rot_angle_vel.reshape(-1,1)])\n                n_rot_base = x.shape[1]\n                if self.euler:\n                    proper = ~np.isclose(x[:,:4],0,atol=1e-7).any(axis=1)\n                    euler_angles = np.zeros_like(x[:,:3])\n                    euler_angles[proper] = np.column_stack(quaternion_to_euler(x[proper,:4]))\n                    x = np.hstack([x,euler_angles / np.pi])\n                    n_rot_base += euler_angles.shape[1]\n                # relative quat to next step\n                if hasattr(self.cfg,'rel_quat'):\n                    # note that relative_rotation_quaternion uses only x[:,:4]\n                    for i in range(self.cfg.rel_quat):\n                        quat_shift = np.zeros_like(x)\n                        quat_shift[i+1:] = x[:-i-1]\n                        for j in range(i+1):\n                            quat_shift[self.start + j] = 0\n                        quat_rel = relative_rotation_quaternion(x,quat_shift)\n                        quat_rel = quat_rel*((quat_rel[:,:1]>0)*2-1)\n                        x = np.hstack([x,quat_rel])\n                    if self.scale_quant > 0:\n                        rq_len = x.shape[1] - n_rot_base\n                        x[:,n_rot_base:] = (x[:,n_rot_base:] - self.rq_median[:,:rq_len])/self.rq_IQR[:,:rq_len]\n                    if self.clip_quat > 0:\n                        x[:,n_rot_base:] = x[:,n_rot_base:].clip(-self.clip_quat,self.clip_quat)\n        return x\n\n    def get_bounds(self, x):\n        end = np.append(np.where(x[1:] != x[:-1])[0] +1, len(x))\n        start = np.append(0,end[:-1])\n        return start,end\n\n    def tof_smoothing(self, tof):\n        t = tof.size(0)  # seq length\n        window_size = self.cfg.tof_win  # Window size for averaging\n        n_windows = tof.size(-1) // window_size\n        x = tof.view(t, n_windows, window_size)\n        if self.tof_drop > 0:\n            mask = torch.bernoulli(torch.full_like(x,1 - self.tof_drop))\n            return (x * mask).sum(dim=-1) / (1e-8 + mask.sum(dim=-1))\n        else: \n            return x.mean(dim=-1)\n\n    def ind_transform(self, x):\n        # sep = torch.zeros_like(x['rot'][...,:self.cfg.sep])\n        sep = torch.zeros(3,self.cfg.image_size,self.cfg.sep)\n        features = [self.ind_prep['acc'](x['acc']),sep,\n                    self.ind_prep['rot'](0.5 * (x['rot'] + 1)),sep,\n                    self.ind_prep['thm'](torch.median(x['thm'],dim=-1,keepdims=True)[0]),sep,\n                    self.ind_prep['thm'](self.tof_smoothing(x['tof']))\n                   ]\n        return torch.cat(features,dim=-1)\n\n\nclass MixtureDataset(Dataset):\n    def __init__(self, base_dataset, alpha, exclude_keys = [], delta = 0):\n        self.base_dataset = base_dataset\n        self.alpha = alpha\n        self.delta = delta\n        self.rng = np.random.default_rng()\n        self.length = len(base_dataset)\n        self.exclude_keys = exclude_keys\n        if isinstance(exclude_keys,str):\n            self.exclude_keys = [exclude_keys]\n\n    def __len__(self):\n        return self.length\n\n    def __getitem__(self, idx):\n        x_i = self.base_dataset[idx]\n        j = np.random.randint(0, self.length)\n        x_j = self.base_dataset[j]\n        # lam = self.rng.beta(self.alpha, self.alpha)\n        lam = 1 - self.alpha * self.rng.random()\n        mixture = {k: lam * x_i[k] + (1 - lam) * x_j[k] \n                   for k in x_i if k not in self.exclude_keys}\n        if self.delta > 0:\n            mixture = {k: v * (1 + self.delta * (2 * torch.rand_like(v) - 1))\n                       for k,v in mixture.items() }\n\n        for k in self.exclude_keys:\n            mixture[k] = x_i[k]\n        return mixture\n\n\nclass AWP:\n    def __init__(self, model, optimizer, scaler, criterion, adv_lr=1, adv_eps=0.001,\n                 start_epoch=0, adv_param=\"weight\"):\n        self.model = model\n        self.optimizer = optimizer\n        self.adv_param = adv_param\n        self.adv_lr = adv_lr\n        self.adv_eps = adv_eps\n        self.start_epoch = start_epoch\n        self.backup = {}\n        self.backup_eps = {}\n        self.scaler = scaler\n        self.criterion = criterion\n\n    def attack_backward(self, x, targets, epoch):\n        if (self.adv_lr == 0) or (epoch < self.start_epoch):\n            return None\n\n        self._save() \n        self._attack_step() \n        with torch.cuda.amp.autocast():\n            outputs = self.model(x)\n            adv_loss = self.criterion(outputs.squeeze(), targets)\n        self.optimizer.zero_grad()\n        self.scaler.scale(adv_loss).backward()\n\n        self._restore()\n\n    def _attack_step(self):\n        e = 1e-6\n        for name, param in self.model.named_parameters():\n            if param.requires_grad and param.grad is not None and self.adv_param in name:\n                norm1 = torch.norm(param.grad)\n                norm2 = torch.norm(param.data.detach())\n                if norm1 != 0 and not torch.isnan(norm1):\n                    r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)\n                    param.data.add_(r_at)\n                    param.data = torch.min(torch.max(\n                        param.data, self.backup_eps[name][0]), self.backup_eps[name][1]\n                    )\n                param.data.clamp_(*self.backup_eps[name])\n\n    def _save(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad and param.grad is not None and self.adv_param in name:\n                if name not in self.backup:\n                    self.backup[name] = param.data.clone()\n                    grad_eps = self.adv_eps * param.abs().detach()\n                    self.backup_eps[name] = (\n                        self.backup[name] - grad_eps,\n                        self.backup[name] + grad_eps,\n                    )\n\n    def _restore(self,):\n        for name, param in self.model.named_parameters():\n            if name in self.backup:\n                param.data = self.backup[name]\n        self.backup = {}\n        self.backup_eps = {}\n\n\nclass CustomNetwork(nn.Module):\n    def __init__(self, d_model=32, nhead=4, num_transformer_layers=2, output_size=18):\n        super().__init__()\n        # CNN for processing input4: [32*8, 5, 8, 8] -> [32*8, 4, 2, 2]\n        self.cnn = nn.Sequential(\n            nn.Conv2d(in_channels=5, out_channels=4, kernel_size=3, stride=2, padding=1),\n            nn.ReLU(),\n            nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, stride=1),\n            nn.ReLU(),\n        )\n\n        input_len = 3 + 4\n        if 'thm' in cfg.use_groups: input_len += 5\n        if 'tof' in cfg.use_groups: \n            input_len += 4*2*2\n        # Linear layer to project concatenated inputs to d_model\n        self.input_projection = nn.Linear(input_len, d_model)\n\n        # Transformer encoder\n        transformer_layer = nn.TransformerEncoderLayer(\n            d_model=d_model,\n            nhead=nhead,\n            dim_feedforward=d_model * 4,\n            dropout=0.1,\n            activation='relu',\n            batch_first=True\n        )\n        self.transformer = nn.TransformerEncoder(\n            transformer_layer,\n            num_layers=num_transformer_layers\n        )\n\n        # Output layer\n        self.output_layer = nn.Linear(d_model, output_size)\n        if cfg.g_loss_weight > 0:\n            self.g_weight = nn.Linear(d_model, 1)\n\n\n    def forward(self, inputs):\n        # Input shapes: [32, 8, 3], [32, 8, 4], [32, 8, 5], [32, 8, 320]\n\n        if 'tof' in inputs:\n            batch_size,seq_len = inputs['tof'].size()[:2]\n            # Reshape input4: [32, 8, 320] -> [32, 8, 5, 8, 8] -> [32*8, 5, 8, 8]\n            input4 = inputs['tof'].reshape(batch_size*seq_len, 5, 8, 8)\n\n            # Apply CNN to input4: [32*8, 5, 8, 8] -> [32*8, 4, 2, 2]\n            input4 = self.cnn(input4)\n\n            # reshape and flatten CNN output: [32*8, 4, 2, 2] -> [32, 8, 16]\n            inputs['tof'] = input4.reshape(batch_size, seq_len, -1)\n\n        # Concatenate inputs: [32, 8, 3 + 4 + 5 + 16] = [32, 8, 16]\n        concat_inputs = torch.cat(list(inputs.values()), dim=-1)\n\n        # Project to d_model: [32, 8, 16] -> [32, 8, d_model]\n        x = self.input_projection(concat_inputs)\n\n        # Apply Transformer: [32, 8, d_model] -> [32, 8, d_model]\n        x = self.transformer(x)\n        if cfg.g_loss_weight > 0:\n            log_weight = self.g_weight(x) # [32, 8, 1]\n            log_weight = F.softplus(log_weight).cumsum(dim=1)\n            log_weight = log_weight - log_weight[:,-1:,:]\n            x = x * torch.exp(log_weight)\n        # Pool over sequence dimension (mean) and project to output: [32, 8, d_model] -> [32, d_model] -> [32, 18]\n        x = x.mean(dim=1)\n        output = self.output_layer(x)\n\n        # Apply softmax for classification probabilities\n        # output = F.softmax(output, dim=-1)\n\n        if self.training & (cfg.g_loss_weight > 0):\n            return output,log_weight.squeeze()\n        else:\n            return output\n\n\ndef get_optimizer(model, args):\n    no_decay = [\"bias\", \"norm\"]\n    if cfg.no_decay:\n        optimizer_parameters = [\n            {'params': [p for n, p in model.named_parameters() \n                        if not any(nd in n for nd in no_decay)],\n             'weight_decay': args.wd},\n            {'params': [p for n, p in model.named_parameters() \n                        if any(nd in n for nd in no_decay)],\n             'weight_decay': 0.0},\n        ]\n    elif args.dif_lr & (args.model.lower() in ['b0','b1','b3','b5','b6','v2_s','v2_m','maxvit_t',\n        'cn_tiny','cn_small','cn_base','mobilenet_v2','resnet18','resnet50','squeezenet1_0']):\n        optimizer_parameters = model.get_param_groups(\n            args.low_lr, args.lr, args.wd)\n    else:\n        optimizer_parameters = model.parameters()\n    optimizer = torch.optim.AdamW(optimizer_parameters, lr=args.lr)\n    return optimizer\n\n\n# Training function with mixed precision\ndef train_fn(model, dataloader, optimizer, scaler, criterion, scheduler, epoch, awp=None):\n    model.train()\n    total_loss = 0\n    for batch in dataloader:\n        inputs = {k:batch[k].to(cfg.device) for k in batch \n                  if k not in ['label','oh_label','gesture_phase']}\n        targets = batch['oh_label'].to(cfg.device)\n        if cfg.g_loss_weight > 0:\n             gesture_phase = batch['gesture_phase'].to(cfg.device)\n\n        # Zero gradients\n        optimizer.zero_grad()\n\n        # Mixed precision context\n        # with autocast():  # Enable FP16 computation\n        with torch.amp.autocast(cfg.device.type):    \n            # Forward pass\n            outputs = model(inputs)  # Remove extra dimension\n\n            # Compute loss\n            if cfg.g_loss_weight > 0:\n                inverse_probs = torch.logit(outputs[1])\n                loss = (criterion(outputs[0].squeeze(), targets) \n                        + cfg.g_loss_weight * cfg.phase_loss(inverse_probs, gesture_phase))\n            else:\n                loss = criterion(outputs.squeeze(), targets)\n\n        # Backward pass with scaling\n        scaler.scale(loss).backward()\n\n        # AWP perturbation\n        if awp is not None:\n            if (np.random.random() < args.awp_prob) & (epoch >= args.awp_start):\n                awp.attack_backward(inputs, targets, epoch)\n\n        # Update optimizer with scaled gradients\n        scaler.step(optimizer)\n\n        # Update the scale factor\n        scaler.update()\n\n        # Step the scheduler\n        scheduler.step()\n\n        total_loss += loss.item()\n\n    avg_loss = total_loss / len(dataloader)\n    return avg_loss\n\n# Validation function\ndef valid_fn(model, dataloader, criterion,return_imu_preds=False):\n    model.eval()\n    total_loss = 0\n    preds = []\n    imu_only_preds = []\n    with torch.no_grad():\n        for batch in dataloader:\n            inputs = {k:batch[k].to(cfg.device) for k in batch \n                      if k not in ['label','oh_label','gesture_phase']}\n            # targets = batch['oh_label'].to(cfg.device)\n            targets = batch['label'].to(cfg.device)\n\n            # Forward pass\n            outputs = model(inputs).squeeze()\n            preds.append(outputs.cpu().numpy())\n            if return_imu_preds:\n                for k in ['tof','thm']:\n                    if k in inputs:\n                        inputs[k].zero_()\n                imu_only_preds.append(model(inputs).squeeze().cpu().numpy())\n            # Compute loss\n            # loss = F.cross_entropy(outputs, targets)\n            # if cfg.public == 1:\n            #     # loss = F.cross_entropy(outputs, torch.argmax(targets, dim=-1))\n            #     loss = F.cross_entropy(outputs, targets)\n            # else:\n            loss = criterion(outputs, targets)\n            total_loss += loss.item()\n\n    avg_loss = total_loss / len(dataloader)\n    if len(imu_only_preds) > 0:\n        imu_only_preds = np.vstack(imu_only_preds)\n    return avg_loss,np.vstack(preds),imu_only_preds\n\n# Testing function\ndef test_fn(model, batch, device, return_imu_preds=False): # batch is one sample here\n    model.eval()\n    with torch.no_grad():\n        inputs = {k:batch[k].unsqueeze(dim=0).to(device) for k in batch \n                  if k not in ['label','gesture_phase']}\n        if return_imu_preds:\n            for k in ['tof','thm']:\n                if k in inputs:\n                    inputs[k].zero_()\n        outputs = model(inputs).squeeze()\n    return outputs.cpu().numpy()\n\n\ndef get_model(cfg):\n    if cfg.model == '0':\n        model = CustomNetwork(d_model=cfg.d_model, nhead=cfg.nhead, num_transformer_layers=2, \n                              output_size=len(cfg.gestures))\n    elif cfg.model == '1':\n        model = TwoBranchModel(imu_dim=7,tof_dim=325 if 'tof' in cfg.use_groups else 0,\n                               n_classes=len(cfg.gestures))\n    elif cfg.model == '2':\n        model = ModelVariant_GRU(num_classes=len(cfg.gestures))\n    elif cfg.model == '3':\n        model = TwoBranchModel80(cfg, tof_dim = 325 if 'tof' in cfg.use_groups else 0)\n    elif cfg.model == '4':\n        model = CMIModel(imu_dim = 7, thm_dim=5, tof_dim=320, n_classes=len(cfg.gestures))\n    elif cfg.model.lower() in ['b0','b1','b3','b5','b6','v2_s','v2_m','cn_tiny','cn_small','cn_base',\n                               'mobilenet_v2','squeezenet1_0','maxvit_t']: \n        model = CustomEfficientNetB0(num_classes=len(cfg.gestures),ver=cfg.model.lower(),\n                                     n_clf = cfg.n_clf if hasattr(cfg,'n_clf') else 0)\n    elif 'resnet' in cfg.model.lower(): \n        model = CustomResNet(num_classes=len(cfg.gestures), ver = cfg.model.lower()[-2:])\n    else:\n        raise ValueError(f'No model {cfg.model.lower()}')\n\n    return model\n\n\ndef train_loop(ds, fold, seed):\n    seed_everything(seed)\n    mlc_t.add('train_loop')\n    if args.folds_file in ['ln_folds','minerppdy_folds']:\n        val_fold = pd.read_csv(cfg.folds)[f'5fold_seed{args.sgkf_seed}'].to_numpy('int8')\n    else:\n        val_fold = pd.read_csv(cfg.folds)[f'{cfg.n_fold}fold_seed1'].to_numpy('int8')\n    train_idx = np.where(val_fold != fold)[0]\n    val_idx = np.where(val_fold == fold)[0]\n    if args.data_aug > 0:\n        train_idx = np.hstack([train_idx + i*len(val_fold) for i in range(args.data_aug+1)])\n    train_dataset = MixtureDataset(Subset(ds,train_idx), args.mix_alpha, \n                                   exclude_keys = 'label', delta = args.delta)\n    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,shuffle=True,\n                              num_workers=cfg.num_workers, pin_memory=True,drop_last=True)\n    fw_preds = []\n    imu_fw_preds = []\n    if len(val_idx) > 0:\n        valid_dataset = Subset(ds,val_idx)\n        valid_loader = DataLoader(valid_dataset, batch_size=4*args.batch_size,shuffle=False,\n                                  num_workers=cfg.num_workers, pin_memory=True,drop_last=False)\n    seed_everything(seed)    \n    model = get_model(cfg).to(cfg.device)\n\n    num_train_steps = int(len(train_dataset) / args.batch_size * args.epochs)\n    num_warmup_steps = int(num_train_steps * args.warmup_ratio)\n    print(f'Fold {fold} {num_train_steps} train steps {num_warmup_steps} warmup steps')\n\n    optimizer = get_optimizer(model, args)\n    scaler =(torch.amp.GradScaler(cfg.device,enabled=cfg.apex) if KAGGLE\n             else torch.cuda.amp.GradScaler(enabled=cfg.apex))\n    scheduler = get_cosine_schedule_with_warmup(optimizer,num_warmup_steps=num_warmup_steps,\n                            num_training_steps=num_train_steps, num_cycles=cfg.num_cycles)\n    class_weights = compute_class_weight('balanced', \n                classes=np.arange(len(cfg.gestures)), y=ds.target[train_idx])\n    if args.loss == '0':\n        pass\n    elif args.loss == '1':\n        criterion = PublicLoss(cfg=cfg, class_weights=class_weights)\n    elif args.loss == '2':\n        criterion = FlexLoss(cfg, class_weights, args.bce_weight, args.aux_weight,\n                             args.ce_for_targets)\n    elif args.loss == '3':\n        criterion = nn.CrossEntropyLoss(\n            weight=torch.tensor(class_weights,dtype=torch.float32).to(cfg.device))\n\n    score,avg_val_loss = 0,0\n    predictions = []\n    history = []\n    awp = AWP(model,optimizer,scaler,criterion,args.awp_lr,args.awp_eps,\n              args.awp_start) if args.awp_lr * args.awp_eps > 0 else None\n\n    for epoch in range(args.epochs):\n        avg_loss = train_fn(model,train_loader,optimizer,scaler,criterion,scheduler,epoch,awp)\n        if len(val_idx) > 0:\n            avg_val_loss,predictions,_ = valid_fn(model,valid_loader,criterion)\n            # scoring\n            class_pred = predictions.argmax(-1)\n            val_y = ds.target[val_idx]\n            score = score_from_int(val_y, class_pred)\n            accuracy = (val_y == class_pred).mean()\n            n = len(cfg.target_gestures)\n            # is_target = softmax(predictions,axis=-1)[:,:n].sum(-1) > 0.5\n            # class_pred[is_target] = predictions[:,:n].argmax(-1)[is_target]\n            # score1 = score_from_int(val_y, class_pred)\n            if (epoch % cfg.print_freq == 0) | (epoch == args.epochs - 1):\n                print(f'Epoch {epoch+1:2} train_loss: {avg_loss:.3f}  val_loss: {avg_val_loss:.3f}',\n                      f'score {score:<.3f} acc {accuracy:<.3f} time: {mlc_t.fsince(\"train_loop\")} ', flush=True)\n            history.append((avg_loss,avg_val_loss,score,accuracy))\n\n    if args.save_fw_preds:\n        valid_loader = DataLoader(ds, batch_size=4*args.batch_size,shuffle=False,\n                                  num_workers=cfg.num_workers, pin_memory=True,drop_last=False)\n        fw_preds,imu_fw_preds = valid_fn(model,valid_loader,criterion,return_imu_preds\n                                         = args.save_imu_fw_preds is not None)[1:3]\n\n    if (args.full_data is not None) | (args.save_models is not None):\n        fname = cfg.output_dir + f\"s{seed}.bin\"\n        if args.save_models:\n            fname = fname.replace('.bin',f'f{fold}.bin')\n        torch.save(model.state_dict(),fname)\n        print(fname,'saved')\n\n\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return predictions,score,avg_val_loss,val_idx,history,fw_preds,imu_fw_preds\n\n\ndef four_digit(a):\n    return f\"{a:.4f}\".split(\".\")[-1]\n\n\nclass WeightedCE(nn.Module):\n    def __init__(self, cfg, class_weights):\n        super().__init__()\n        self.class_weights = torch.FloatTensor(class_weights).to(cfg.device)\n\n    def forward(self, logits, batch_y):\n        if args.new_wce:\n            log_probs = F.log_softmax(logits, dim=-1)\n            weighted_log_probs = log_probs * self.class_weights\n            loss = -torch.sum(weighted_log_probs * batch_y, dim=-1)\n        else:\n            loss = -torch.sum(F.log_softmax(logits, dim=-1) * batch_y, dim=-1).mean()\n            sample_weights = torch.sum(batch_y * self.class_weights, dim=-1)\n            loss = loss * sample_weights\n        return loss.mean()\n\n\nclass FlexLoss(nn.Module):\n    def __init__(self, cfg, class_weights, bce_weight=0, aux_weight=0, \n                 ce_for_targets = False):\n        super().__init__()\n        sum_loss_weights = 1 + bce_weight + aux_weight\n        self.ce_weight = 1 / sum_loss_weights\n        self.bce_weight = bce_weight / sum_loss_weights\n        self.aux_weight = aux_weight /sum_loss_weights\n        self.num_targets = len(cfg.target_gestures)\n        self.num_classes = len(cfg.gestures)\n        target_weights = [1/np.sum(1/class_weights[self.num_targets:]),\n                          1/np.sum(1/class_weights[:self.num_targets])]\n        self.ce_for_targets = ce_for_targets\n        if ce_for_targets:\n            ce_weights = np.append(class_weights[:self.num_targets],target_weights[1])\n            self.ce_loss = WeightedCE(cfg, ce_weights)\n        else:\n            self.ce_loss = WeightedCE(cfg, class_weights)\n        self.bce_loss = WeightedCE(cfg, target_weights)\n        self.aux_loss = WeightedCE(cfg, class_weights[self.num_targets:])\n\n\n    def forward(self, logits, batch_y):\n        if len(batch_y.shape) == 1:\n            batch_y = F.one_hot(batch_y, num_classes=self.num_classes).float()\n        if self.ce_for_targets:\n            ce_logits = torch.cat([logits[...,:self.num_targets],\n                                   logits[...,:self.num_targets].logsumexp(dim=-1,keepdims=True)],dim=-1)\n            ce_y = torch.cat([batch_y[:,:self.num_targets],\n                              batch_y[:,self.num_targets:].sum(dim=-1,keepdims=True)],dim=-1)\n            total_loss = self.ce_weight * self.ce_loss(ce_logits, ce_y)\n        else:\n            total_loss = self.ce_weight * self.ce_loss(logits, batch_y)\n        if self.bce_weight > 0:\n            target_logits = torch.cat([logits[...,self.num_targets:].logsumexp(dim=-1,keepdims=True),\n                                       logits[...,:self.num_targets].logsumexp(dim=-1,keepdims=True)],dim=-1)\n            target_y = torch.cat([batch_y[:,self.num_targets:].sum(dim=-1,keepdims=True),\n                                  batch_y[:,:self.num_targets].sum(dim=-1,keepdims=True)],dim=-1)\n            total_loss = total_loss + self.bce_weight * self.bce_loss(target_logits, target_y)\n        if self.aux_weight > 0:\n            aux_logits = logits[...,self.num_targets:]\n            aux_y = batch_y[:,self.num_targets:] / (batch_y[:,self.num_targets:].sum(\n                dim=-1,keepdims=True) + 1e-7)\n            total_loss = total_loss + self.aux_weight * self.aux_loss(aux_logits, aux_y)\n        return total_loss\n\n\ndef label_smoothing_loss(pred, target, smoothing=0.1):\n    \"\"\"Label smoothing loss\"\"\"\n    confidence = 1.0 - smoothing\n    log_probs = F.log_softmax(pred, dim=-1)\n    nll_loss = -log_probs.gather(dim=-1, index=target.unsqueeze(1))\n    nll_loss = nll_loss.squeeze(1)\n    smooth_loss = -log_probs.mean(dim=-1)\n    loss = confidence * nll_loss + smoothing * smooth_loss\n    return loss.mean()\n\n\nclass PublicLoss(nn.Module):\n    def __init__(self, cfg, class_weights):\n        super().__init__()\n        self.class_weights = torch.FloatTensor(class_weights).to(cfg.device)\n\n    def forward(self, logits, batch_y):\n        # Handle mixup targets\n        if len(batch_y.shape) == 2 and batch_y.shape[1] > 1:  # One-hot or mixed\n            loss = -torch.sum(F.log_softmax(logits, dim=1) * batch_y, dim=1).mean()\n            # For accuracy, use argmax of target\n            # targets = batch_y.argmax(dim=1)\n        else:\n            targets = batch_y.long()\n            loss = label_smoothing_loss(logits, targets, smoothing=0.1)\n\n        # Apply class weights\n        if len(batch_y.shape) == 2:\n            sample_weights = torch.sum(batch_y * self.class_weights.unsqueeze(0), dim=1)\n            loss = (loss * sample_weights).mean()\n        return loss\n\n\nclass MultiClassifier(nn.Module):\n    def __init__(self, in_features, n_classes, n_clf, dropout_rate=0.2):\n        super().__init__()\n\n        self.n_clf = n_clf\n        self.dropout_rate = dropout_rate\n\n        self.classifiers = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(in_features, in_features // 2),\n                nn.SELU(),\n                nn.Dropout(dropout_rate), \n                nn.Linear(in_features // 2, n_classes)\n            ) for _ in range(n_clf)\n        ])\n\n    def forward(self, x):\n        return torch.stack([classifier(x) for classifier in self.classifiers])\n\n\ndef parse_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--model\", type=str, default='0', required=False)\n    parser.add_argument(\"--lr\", type=float, default=0.001, required=False)\n    parser.add_argument(\"--low_lr\", type=float, default=1e-5, required=False)\n    parser.add_argument(\"--head_lr\", type=float, default=1, required=False)\n    parser.add_argument(\"--wd\", type=float, default=0.001, required=False)\n    parser.add_argument(\"--max_len\", type=int, default=32, required=False)\n    parser.add_argument(\"--seed\", type=str, default='123', required=False) # list of seeds\n    parser.add_argument(\"--epochs\", type=int, default=1, required=False)\n    parser.add_argument(\"--batch_size\", type=int, default=32, required=False)\n    parser.add_argument(\"--input\", type=str, default=\"./data/\", required=False)\n    parser.add_argument(\"--output\", type=str, default='', required=False)\n    parser.add_argument(\"--my_data_path\", type=str, default=\"./my_data/\", required=False)\n    parser.add_argument(\"--n_cycles\", type=float, default=0.5, required=False)\n    parser.add_argument(\"--warmup_ratio\", type=float, default=0, required=False)\n    parser.add_argument('--full_data', action=argparse.BooleanOptionalAction)\n    parser.add_argument(\"--grad_norm\", type=float, default=1000, required=False)\n    parser.add_argument(\"--gamma\", type=float, default=0,required=False) # for focal loss\n    # path to the model checkpoints\n    parser.add_argument(\"--weights\", type=str, default='', required=False)  \n    parser.add_argument('--predict', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--g_loss_weight', type=float, default=0,required=False)    \n    parser.add_argument('--speed_test', action=argparse.BooleanOptionalAction)\n    # checkpoints indices to be selected for prediction; the order in not specified\n    parser.add_argument(\"--checkpoint_idx\", type=str, default='01234', required=False)\n    parser.add_argument(\"--gpu\", type=str, default='t4', required=False) \n    # 'ln_folds' my with sgkf seeds 'minerppdy_folds' minerppdy with sgkf seeds\n    parser.add_argument(\"--folds_file\", type=str, default='team_folds', required=False) \n    parser.add_argument(\"--d_model\", type=int, default=64, required=False)\n    parser.add_argument(\"--nhead\", type=int, default=4, required=False)\n    parser.add_argument(\"--loss\", type=str, default='0', required=False)\n    parser.add_argument('--bce_weight', type=float, default=0,required=False)    \n    parser.add_argument('--aux_weight', type=float, default=0,required=False)    \n    parser.add_argument('--ce_for_targets', action=argparse.BooleanOptionalAction)    \n    parser.add_argument('--mix_alpha', type=float, default=0,required=False)    \n    parser.add_argument('--delta', type=float, default=0,required=False)    \n    parser.add_argument('--save_fw_preds', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--save_imu_fw_preds', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--save_models', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--imu_only', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--imu_fe', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--awp_lr', type=float, default=0,required=False)    \n    parser.add_argument('--awp_eps', type=float, default=0,required=False)    \n    parser.add_argument('--awp_start', type=int, default=0,required=False)    \n    parser.add_argument('--awp_prob', type=float, default=0,required=False)   \n    parser.add_argument('--tof_win', type=int, default=16,required=False)    \n    parser.add_argument('--sep', type=int, default=0,required=False)    \n    parser.add_argument('--hflip', type=float, default=0,required=False)   \n    parser.add_argument('--dif_lr', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--add_ave', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--flip_left', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--flip_rot', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--flip_tof', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--sgkf_seed', type=int, default=0,required=False)\n    parser.add_argument('--fill_rot_mean', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--rand_aug', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--data_aug', type=int, default=0,required=False)\n    parser.add_argument('--y_aug_max', type=float, default=7,required=False)   \n    parser.add_argument('--z_aug_max', type=float, default=45,required=False)   \n    parser.add_argument('--new_wce', action=argparse.BooleanOptionalAction)\n    # n_clf classifiers trained independently\n    parser.add_argument('--n_clf', type=int, default=0,required=False)\n    parser.add_argument('--imu_image_weight', type=float, default=0,required=False)   \n    parser.add_argument('--rel_quat', type=int, default=0,required=False)\n    # percentile for robast scaling of rel_quat (e.g. 15 or 25)\n    parser.add_argument('--scale_quant', type=int, default=0,required=False)\n    parser.add_argument('--clip_quat', type=float, default=0,required=False)\n    parser.add_argument('--euler', action=argparse.BooleanOptionalAction)\n    # methods: cosine,angular_sum, gramian\n    parser.add_argument(\"--gaf\", type=str, default='', required=False) \n    parser.add_argument('--fix_outliers', action=argparse.BooleanOptionalAction)\n    parser.add_argument('--tof_drop', type=float, default=0,required=False)\n    return parser.parse_args()\n\n\n# Main script\nif __name__ == \"__main__\":\n    print('start')\n    mlc_t = mlc.time.Timer()\n    # Parse command-line arguments\n    args = parse_args()\n    if args.predict:\n        cfg.train = False\n    if args.full_data:\n        cfg.trn_folds = [-1]\n    cfg.folds = cfg.folds.replace('team_folds',args.folds_file)\n    for key in ['max_len','d_model','nhead','imu_fe','tof_win','sep','n_clf',\n                'imu_image_weight','rel_quat','scale_quant','clip_quat','gaf',\n                'tof_drop']:\n        setattr(cfg, key, getattr(args, key))\n    cfg.model = args.model.lower()\n    image_size_dic = defaultdict(lambda: 224, {'b1':240,'b3':300,'b5':456,'b6':528,\n                                               'v2_s':384,'v2_m':480})\n    cfg.image_size = image_size_dic[cfg.model]\n    cfg.no_decay = args.dif_lr is None\n    for key in ['flip_left','flip_rot','flip_tof','fill_rot_mean','euler']:\n        setattr(cfg, key, getattr(args, key) is not None)\n    if args.data_aug > 0:\n        print(args.data_aug,args.y_aug_max,args.z_aug_max)\n        cfg.z_aug_grid = [d for d in np.linspace(\n            -args.z_aug_max,args.z_aug_max,args.data_aug + 1) if abs(d) > 2.1]\n        cfg.y_aug_grid = [d for d in np.linspace(\n            -args.y_aug_max,args.y_aug_max,args.data_aug + 2) if abs(d) > 2.1]\n        print('y_aug_grid =',cfg.y_aug_grid)\n        print('z_aug_grid =',cfg.z_aug_grid)\n\n    if args.imu_only:\n        cfg.use_groups = ['acc','rot']\n\n    # save after all changes in cfg   \n    save_namespace_json(cfg, cfg.output_dir + 'cfg.json')\n    # Load datasets\n    mlc_t.add('load')\n    # train = pd.read_csv(data_path + 'train.csv')\n    train = pd.read_feather(my_data_path + 'train.feather')\n    print('loaded',mlc_t.fsince('load'))\n    if args.fix_outliers:\n        train = Fix_outlier(train, args.imu_only)\n        print('outliers fixed',mlc_t.fsince('load'))\n    demo = pd.read_csv(\n        data_path + 'train_demographics.csv') if args.flip_left else None\n    ds = CustomDataset(train, cfg, args.mix_alpha, demo)\n    print('ds created',mlc_t.fsince('load'))\n\n    if cfg.train:\n        scores,losses = [],[]\n        oof = np.zeros((train.sequence_id.nunique(),len(cfg.gestures)),dtype=np.float32)\n        imu_oof = oof.copy()\n        fw_preds = np.zeros((len(cfg.trn_folds),) + oof.shape, dtype=np.float32)\n        imu_fw_preds = fw_preds.copy()\n        if args.folds_file in ['ln_folds','minerppdy_folds']:\n            print(f'folds_file {args.folds_file}, sgkf {args.sgkf_seed}')\n        for seed in args.seed: # each seed in args.seed is one digit number\n            print('Training with seed',seed)\n            # mlc_t.add('load')\n            # random.seed(int(seed))\n            # random.shuffle(cfg.y_aug_grid)\n            # random.shuffle(cfg.z_aug_grid)\n            # ds = CustomDataset(train, cfg, args.mix_alpha, demo)\n            # print('ds created',mlc_t.fsince('load'))\n            seed_oof = np.zeros(len(oof),dtype='int8')\n            for fold in cfg.trn_folds:\n                # print(10*int(seed) + fold,type(10*int(seed) + fold))\n                pred,score,val_loss,val_idx,history,fw_pred,imu_fw_pred = train_loop(ds,fold,seed)\n                if args.save_fw_preds:\n                    fw_preds[fold] += fw_pred\n                    if len(imu_fw_pred) > 0:\n                        imu_fw_preds[fold] += imu_fw_pred\n                if cfg.save_history:\n                    if (fold == 0) & (seed == args.seed[0]) & (not KAGGLE):\n                        pd.DataFrame(history,columns=['train_loss','val_loss','score','accuracy']).to_csv(\n                            f'models/{cfg.job_id}/history.csv',index=False)\n                if len(val_idx) > 0:\n                    scores.append(score)\n                    losses.append(val_loss) \n                    oof[val_idx] += pred / len(args.seed)\n                    seed_oof[val_idx] = pred.argmax(-1)\n            if len(val_idx) > 0:\n                print(f'Average score: {np.mean(scores):.4f} loss: {np.mean(losses):.4f} scores std: {np.std(scores):.4f}')\n                if args.folds_file in ['ln_folds','minerppdy_folds']:\n                    out = '\\t'.join(f'={_:.4f}' for _ in scores[-5:])\n                    print(f'Scores and overall: {out}',\n                          f'{score_from_int(ds.target, seed_oof):.4f}')\n\n        if len(val_idx) > 0:\n            tg = ds.target[:len(oof)]\n            oof_score = score_from_int(tg, oof.argmax(-1))\n            logloss = log_loss(tg, softmax(oof,-1),labels=np.arange(oof.shape[-1]))\n            print(f'Total for mean oof score/logloss {oof_score:.4f}/{logloss:.2f}')\n            if args.save_fw_preds:\n                np.save(f'models/{cfg.job_id}/fw_preds.npy',fw_preds / len(args.seed))\n                print('Foldwise predictions with shape',fw_preds.shape,'saved')\n            if args.save_imu_fw_preds:\n                np.save(f'models/{cfg.job_id}/imu_fw_preds.npy',imu_fw_preds / len(args.seed))\n                print('imu only fw_preds with shape',imu_fw_preds.shape,'saved')\n\n            if not KAGGLE:\n                np.save(f'models/{cfg.job_id}/oof_{four_digit(oof_score)}.npy',oof)\n                print(f'Debug output. Saved oof after {seed=} {fold=}')\n    # predict\n    else:\n        pass\n    print(f'Total time ',mlc_t.fsince(0))\n","metadata":{"papermill":{"duration":0.020397,"end_time":"2025-06-14T17:56:59.276951","exception":false,"start_time":"2025-06-14T17:56:59.256554","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T09:40:09.815597Z","iopub.execute_input":"2025-08-25T09:40:09.815919Z","iopub.status.idle":"2025-08-25T09:40:09.851286Z","shell.execute_reply.started":"2025-08-25T09:40:09.815893Z","shell.execute_reply":"2025-08-25T09:40:09.850537Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset\nimport tensorflow as tf\nfrom scipy.special import softmax\nfrom scipy.stats.mstats import gmean\nfrom copy import deepcopy\nfrom pathlib import Path\nfrom main import CustomDataset,CustomNetwork,cfg,test_fn,get_model,load_namespace_json\nfrom cmi_utils import  (quaternion_multiply, quaternion_to_euler, euler_to_quaternion,\n                        custom_objs, prepare_tf79_input)\n\nargs = cfg\ncfg.d_model = 64\ncfg.nhead = 4\ncfg.flip_left = False\ncfg.imu_fe = False\nfull_data = True # keep only full_data models\n\n# Set GPU assignments\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\"  # Both GPUs visible, but we'll control usage\n# TensorFlow setup: Restrict to GPU 0 and enable memory growth\ntf_physical_devices = tf.config.list_physical_devices('GPU')\nif tf_physical_devices:\n    try:\n        tf.config.set_visible_devices([tf_physical_devices[0]], 'GPU')  # Use only GPU 0\n        for gpu in tf.config.get_visible_devices('GPU'):\n            tf.config.experimental.set_memory_growth(gpu, True)\n        print(\"TensorFlow restricted to GPU 0 with memory growth enabled.\")\n    except RuntimeError as e:\n        print(f\"Error restricting TensorFlow to GPU 0: {e}\")\n\n# PyTorch setup: Use GPU 1\ntorch_device = torch.device(\"cuda:1\")  # GPU 1 for PyTorch\ncfg.device = torch_device\n\n\n# Load NN models\nmodels_path = '/kaggle/input/cmi-models-new'\nmodel_names867 = {'imu': ['tf79-imu-fl-aug-8000','tf79-imu-fl-7928',\n                       'tb0-imu-rq2-8095',\n                       'tb3-imu-aug-8065','tb3-imu-rq3-8227',\n                       'tb5-imu-aug-8083','tb5-imu-rq4-8232',\n                       ],\n               'all': ['tf79-imu-fl-aug-8000','tf79-fl-8325',\n                       'tb5-imu-rq4-8232',           \n                       'tb0-8450','tb5-8581','tv2_s-rq2-8573',\n                      ]\n               } # LB 867\nmodel_names = {'imu': ['tf79-imu-fl-aug-8000','nb3-imu-rq5-8220',\n                        'nb5-imu-aug-8095','nb5-imu-rq4-8223',\n                       ],\n               'all': ['tf79-imu-fl-aug-8000','tf79-fl-8325',\n                       'nb5-imu-rq4-8223', \n                       'nb0-8424','nb3-rq2-sc15-aug-8671',\n                       'nb5-rq3-sc15-8694','nv2_m-aug-8704'\n                      ]\n               } # LB 863\nmodel_names['imu'] = ['tf79-imu-fl-aug-8000','tf79-imu-fl-7928',\n                        'nb0-imu-rq2-8036','nb3-imu-rq5-8220',\n                        'nb5-imu-aug-8095','nb5-imu-rq4-8223',\n                       ]\n# model_names = {'imu': ['tf79-imu-fl-aug-8000','tf79-imu-fl-7928',\n#                         'nb0-imu-rq2-8036','nb3-imu-rq5-8220',\n#                         'nb5-imu-aug-8095','nb5-imu-rq4-8223',\n#                        ],\n#                'all': ['tf79-imu-fl-aug-8000','tf79-fl-8325', 'tf79-fl-aug-8392',\n#                        'nb5-imu-rq4-8223', 'nb3-imu-rq5-8220',\n#                        'nb0-8424','nb0-rq3-sc15-aug-8652',\n#                        'nb0-rq5-sc15-8593','nb3-rq2-sc15-aug-8671',\n#                        'nb5-8571','nb5-aug-8754','nv2_m-aug-8704'\n#                       ]\n#                } not submitted with foldwaise selection\n\n\nno_inf_par = ['debug','train','print_freq','num_workers','prefix',\n              'gradient_checkpointing','scheduler','batch_scheduler','num_cycles',\n              'min_lr','eps','betas','max_grad_norm','trn_folds','folds',\n              'no_decay','job_id','save_history','output_dir']\ntof_columns = [f'tof_{i}_v{j}' for i in range(1,6) for j in range(64)]\nthm_columns = ['thm_1', 'thm_2', 'thm_3', 'thm_4', 'thm_5']\n\n\ndef flip_acc_rot(df, demo_df):\n    handedness = (demo_df[['subject','handedness']].set_index('subject')\n            .reindex(df['subject'].values).handedness.to_numpy('int8'))\n    df.loc[handedness==0,'acc_x'] = - df.loc[handedness==0,'acc_x']\n    # clip_rot\n    rot_columns = ['rot_w','rot_x', 'rot_y', 'rot_z']\n    x = df[rot_columns].to_numpy(np.float32)\n    proper_left = (~np.isnan(x).any(axis=1)) & (handedness == 0)\n    d=-110\n    quat_rot_z = np.array([np.cos(np.deg2rad(d)/2),0,0,np.sin(np.deg2rad(d)/2)])\n    if proper_left.any():\n        quat_flip = x[proper_left].copy() \n        quat_flip = quaternion_multiply(quat_rot_z, quat_flip)[:,[1,2,3,0]] \n        euler_x,euler_y,euler_z = quaternion_to_euler(quat_flip)  \n        quat_flip = euler_to_quaternion(euler_x, -euler_y, -euler_z)\n        quat_flip = quat_flip*((quat_flip[:,0:1]>0)*2-1) \n        df.loc[proper_left,rot_columns]= quat_flip.clip(-1,1)\n    return df\n\n\nclass Classifier():\n    def __init__(self):\n        self.cfg = cfg\n        self.nn_ensemble = {'imu':[], 'all':[]}\n        self.load_models()\n        self.threshold = 0.2\n\n    def predict(self, pd_seq: pd.DataFrame, pd_demo: pd.DataFrame,\n                key = None) -> np.ndarray:\n        pd_seq[tof_columns+thm_columns] = pd_seq[tof_columns+thm_columns].astype(np.float32)\n        if key is None:\n            no_tof = pd_seq[tof_columns].isnull().values.mean() > self.threshold\n            key = 'imu' if no_tof else 'all'\n        else:\n            assert key in ['imu','all'], f'Unknown key {key}'\n        nn_preds = []\n        # print(key)\n        for n in self.nn_ensemble[key]:\n            cfg = n['cfg']\n            if n['model'] == 'tf':\n                seq_copy = deepcopy(pd_seq)\n                if 'fl' in n['model']:\n                    seq_copy = flip_acc_rot(seq_copy, pd_demo)\n                batch = prepare_tf79_input(seq_copy)\n                nn_preds.append(\n                    np.mean([m(batch, training=False).numpy()\n                             for m in n['models']],axis=0).flatten()\n                )\n            else:\n                cfg.device = torch_device\n                ds = CustomDataset(deepcopy(pd_seq), cfg,demo=pd_demo)\n                nn_preds.append(\n                    softmax(np.mean([test_fn(m, ds[0], torch_device) \n                                     for m in n['models']],axis=0),axis=-1)\n                    )\n        preds = gmean(nn_preds,axis=0)\n        return preds\n\n    def load_models(self):\n        for key in self.nn_ensemble:\n            print(f'=== {key} ===')\n            for i,model_name in enumerate(model_names[key]):\n                n = {'path': Path(models_path) / model_name}\n                print('model',i,n['path'])\n                saved_cfg =  deepcopy(self.cfg)\n                if 'tf' in model_name:\n                    n['model'] = 'tf'\n                else:\n                    config = load_namespace_json(n['path'] / 'cfg.json')\n                    for sub_k in config.__dict__:\n                        if sub_k in no_inf_par:\n                            continue\n                        new = config.__dict__[sub_k]\n                        if not hasattr(self.cfg, sub_k):\n                            setattr(self.cfg, sub_k, new)\n                            print('set',sub_k,'to',new)\n                        elif new != self.cfg.__dict__[sub_k]:\n                            setattr(self.cfg, sub_k, new)\n                            print('modified',sub_k,'to',new)\n                    n['model'] = self.cfg.model\n                n['cfg'] = self.cfg\n                ext = 'h5' if n['model'] == 'tf' else 'bin'\n                checkpoints = [f for f in os.listdir(n['path']) if ext in f]\n                if full_data: # keep only full_data models\n                    checkpoints = [f for f in checkpoints if 'f' not in f.split('_')[-1]]\n                else: # only models trained on folds\n                    checkpoints = [f for f in checkpoints if 'f' in f.split('_')[-1]]\n                \n                print('Found checkpoints', checkpoints)\n                n['fold'] =[]\n                n['models'] = []\n                for fname in checkpoints:\n                    if n['model'] == 'tf':\n                        model = tf.keras.models.load_model(n['path'] / fname, \n                                        compile=False,custom_objects=custom_objs)\n                    else:\n                        model = get_model(self.cfg).to(torch_device)\n                        model.load_state_dict(torch.load(n['path'] / fname,\n                                                         map_location=torch_device))\n                    n['models'].append(model)\n                    if 'f' in fname:\n                        n['fold'].append(fname[-6:-4])\n                self.nn_ensemble[key].append(n)\n                self.cfg = deepcopy(saved_cfg)\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}