{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"sourceType":"competition"},{"sourceId":12888675,"sourceType":"datasetVersion","datasetId":6093998},{"sourceId":12929851,"sourceType":"datasetVersion","datasetId":8178148},{"sourceId":261252662,"sourceType":"kernelVersion"}],"dockerImageVersionId":31089,"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":"markdown","source":"### Training script","metadata":{}},{"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-09-11T12:54:17.410445Z","iopub.execute_input":"2025-09-11T12:54:17.410709Z","iopub.status.idle":"2025-09-11T12:54:17.439840Z","shell.execute_reply.started":"2025-09-11T12:54:17.410686Z","shell.execute_reply":"2025-09-11T12:54:17.439049Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### The list of @nazarov models used in the final submission","metadata":{}},{"cell_type":"code","source":"['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 'tf79-fl-8325',\n 'nb0-8424','nb3-rq2-sc15-aug-8671',\n 'nb5-rq3-sc15-8694','nv2_m-aug-8704'\n ]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Training instructions\n\n#### Tensorflow models\n\nIn order to get 3 checkpoint of the model `tf79-imu-fl-7928` run [LB_79_work_with_time v96](https://www.kaggle.com/code/nazarov/lb-79-work-with-time?scriptVersionId=255319058)\n\nIn order to get 3 checkpoint of the model `tf79-fl-8325` run [LB_79_work_with_time v97](https://www.kaggle.com/code/nazarov/lb-79-work-with-time?scriptVersionId=255321006)\n\nIn order to get 3 checkpoint of the model `tf79-imu-fl-aug-8000` run [LB_79_work_with_time v113](https://www.kaggle.com/code/nazarov/lb-79-work-with-time?scriptVersionId=255664981)\n\n\n\n#### PyTorch models\nEach model which name starts with \"n\" has subfolder with the same name in the dataset [cmi-models-new](https://www.kaggle.com/datasets/nazarov/cmi-models-new). This subfolder contains three checkpoints and a Slurm job script (.sbatch) that includes the call to the training script. These script calls are collected in the cell below.","metadata":{}},{"cell_type":"code","source":"['nb0-imu-rq2-8036',\n# python main.py --epochs 2 --lr 0.0007 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 32 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b0 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --imu_only --imu_fe --flip_left --sep 2 --data_aug 6 \\\n# --y_aug_max 9 --z_aug_max 30 --rel_quat 2 --full_data\n'nb3-imu-rq5-8220',\n#  python main.py --epochs 2 --lr 0.0011 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 32 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b3 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --imu_fe --rel_quat 5 --imu_only --sep 2 \\\n# --flip_left --flip_rot --fill_rot_mean \\\n# --data_aug 6 --y_aug_max 9 --z_aug_max 30 --full_data\n'nb5-imu-aug-8095',\n#  python main.py --epochs 2 --lr 0.0009 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b5 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --tof_win 4 --imu_fe --imu_only --sep 2 --flip_left --flip_rot \\\n# --data_aug 6 --y_aug_max 9 --z_aug_max 30 --full_data\n'nb5-imu-rq4-8223',\n#  python main.py --epochs 2 --lr 0.0009 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b5 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --imu_fe --rel_quat 4 --imu_only --sep 2 --flip_left --flip_rot \\\n# --data_aug 6 --y_aug_max 9 --z_aug_max 30 --full_data\n'nb0-8424',\n#  python main.py --epochs 11 --lr 0.0014 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 32 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b0 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --tof_win 4 --flip_left --flip_rot --flip_tof --full_data\n'nb3-rq2-sc15-aug-8671',\n#  python main.py --epochs 3 --lr 0.0005 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b3 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --tof_win 4 --imu_fe --sep 2 --flip_left --flip_rot --fill_rot_mean --rel_quat 3 \\\n# --scale_quant 15 --clip_quat 1 --tof_drop 0.5 --data_aug 6 \\\n# --y_aug_max 9 --z_aug_max 30 --full_data\n'nb5-rq3-sc15-8694',\n#  python main.py --epochs 9 --lr 0.001 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.1 --max_len 88 --model b5 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --tof_win 4 --imu_fe --sep 2 --flip_left --flip_rot --fill_rot_mean --rel_quat 3 \\\n# --scale_quant 15 --clip_quat 1 --full_data\n'nv2_m-aug-8704'\n# python main.py --epochs 3 --lr 0.0005 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n# --mix_alpha 0.4 --wd 0.01 --max_len 88 --model v2_m --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n# --ce_for_targets --tof_win 4 --imu_fe --sep 2 --flip_left --flip_rot --rel_quat 3 \\\n# --scale_quant 25 --clip_quat 1 --tof_drop 0.5 --data_aug 6 \\\n# --y_aug_max 9 --z_aug_max 30 --full_data\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-11T12:54:17.441156Z","iopub.execute_input":"2025-09-11T12:54:17.441381Z","iopub.status.idle":"2025-09-11T12:54:17.459736Z","shell.execute_reply.started":"2025-09-11T12:54:17.441363Z","shell.execute_reply":"2025-09-11T12:54:17.459202Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The models have been trained locally with 128 Gb RAM. That's why setting `--data_aug 6` does not enable training on Kaggle. Here is the example of training such model (nb5-imu-rq4-8223) with reduced `data_aug' for 1 epoch. ","metadata":{}},{"cell_type":"code","source":"!python main.py --epochs 1 --lr 0.0009 --low_lr 1e-5 --warmup_ratio 0.1 --batch_size 16 --seed 123 \\\n--mix_alpha 0.4 --wd 0.1 --max_len 88 --model b5 --loss 2 --bce_weight 4 --aux_weight 0.3 \\\n--ce_for_targets --imu_fe --rel_quat 4 --imu_only --sep 2 --flip_left --flip_rot \\\n--data_aug 2 --y_aug_max 9 --z_aug_max 30 --full_data","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-11T12:54:17.460468Z","iopub.execute_input":"2025-09-11T12:54:17.460984Z","execution_failed":"2025-09-11T12:56:32.566Z"}},"outputs":[],"execution_count":null}]}