Upload 115 files

PreTrain_MeDSLIP/configs/Pretrain_MeDSLIP.yaml

@@ -0,0 +1,34 @@
+train_file: "setting/rad_graph_metric_train_local.json"
+valid_file: "setting/rad_graph_metric_validate_local.json"
+test_file: "setting/rad_graph_metric_test_local.json"
+label_file: "setting/landmark_observation_adj_mtx.npy"
+pathology_book: "PreTrain_MeDSLIP/data_file/observation explanation.json"
+
+image_res: 224
+patch_size: 16
+num_sentences: 12
+num_tokens: 32
+vision_width: 768
+fea_width: 197
+embed_dim: 256
+batch_size: 64
+test_batch_size: 32
+temp: 0.07
+mlm_probability: 0.15
+queue_size: 8192
+momentum: 0.995
+alpha: 0.4
+d_model: 256
+res_base_model: "resnet50"
+num_queries: 75
+dropout: 0.1
+attribute_set_size: 2
+N: 4
+H: 4
+no_cl: False
+
+exclude_class: False
+text_encoder: "emilyalsentzer/Bio_ClinicalBERT"
+shuffle_ratio: 0.5
+optimizer: {opt: adamW, lr: 1e-4, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 1e-4, epochs: 100, min_lr: 1e-5, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 5, cooldown_epochs: 0}

PreTrain_MeDSLIP/data_file/observation explanation.json

@@ -0,0 +1,77 @@
+{
+    "normal": "It means the absence of diseases and infirmity, indicating the structure is normal.",
+    "clear": "The lungs are clear and normal. No evidence for other diseases on lung.",
+    "sharp": "This means that an anatomical structure's boundary or edge is clear and normal, meaning it is free of diseases.",
+    "sharply": "\u2018Sharply seen\u2019 means that an anatomical structure is clearly visible.",
+    "unremarkable": "This represents some anatomical structures are normal, usually modifying cardiac and mediastinal silhouettes.",
+    "intact": "The bonny structure is complete and normal, meaning no fractures.",
+    "stable": "The modified anatomical structures are normal and stable. No evidence for diseases.",
+    "free": "It usually refers to free air and is associate with pneumothorax,atelectasis,pneumoperitoneum and emphysema.",
+    "effusion": "A pleural effusion is accumulation of excessive fluid in the pleural space, the potential space that surrounds each lung. A pleural effusion infiltrates the space between the visceral pleura and the parietal pleura",
+    "opacity": "It is defined as an area of hazy opacification due to air displacement by fluid, airway collapse, fibrosis, or a neoplastic process. It is causes include infections, interstitial lung disease, and pulmonary edema.",
+    "pneumothorax": "A pneumothorax is an abnormal collection of air in the pleural space between the lung and the chest wall. It may be caused by pneumonia or fibrosis and other diseases. ",
+    "edema": "Pulmonary edema, also known as pulmonary congestion, is excessive liquid accumulation in the tissue and air spaces of the lungs. It will show fluid in the alveolar walls",
+    "atelectasis": "It is the collapse or closure of a lung resulting in reduced or absent gas exchange. Findings can include lung opacification and loss of lung volume.",
+    "tube": "it is a surgical drain that is inserted through the chest wall and into the pleural space or the mediastinum to remove undesired substances such as air (pneumothorax), excess fluid (pleural effusion or hydrothorax), blood (hemothorax), chyle (chylothorax) or pus (empyema) from the intrathoracic space.",
+    "consolidation": "it is a region of normally compressible lung tissue that has filled with liquid instead of air. Consolidation must be present to diagnose pneumonia: the signs of lobar pneumonia are characteristic and clinically referred to as consolidation.",
+    "process": "Acute process' means there is abnormality in the anotomy structure. ",
+    "abnormality": "It means the exist of diseases and infirmity, indicating the structure is abnormal.",
+    "enlarge": "It usually modifies cardiac silhouette and heart. Cardiomegaly is a medical condition in which the heart is enlarged.",
+    "tip": "It refers to the top head of the tube.",
+    "low": "The presence of low lung volumes may be a sign of a restrictive lung condition such as pulmonary fibrosis or sarcoidosis.",
+    "eumonia": "Pneumonia is an inflammatory condition of the lung primarily affecting the small air sacs known as alveoli.  Pneumonia may present with opacities. Complications such as pleural effusion may also be found increasing the diagnostic accuracy of lung consolidation and pleural effusion",
+    "line": "It refers to venous access line ot PICC lines.",
+    "congestion": "Pulmonary congestion is defined as accumulation of fluid in the lungs, resulting in impaired gas exchange and arterial hypoxemia. ",
+    "catheter": "catheter is a tube placed in the body to drain and collect urine from the bladder",
+    "cardiomegaly": "Cardiomegaly (sometimes megacardia or megalocardia) is a medical condition in which the heart is enlarged. ",
+    "fracture": "fracture is a break in a rib bone.",
+    "air": "It refers to the free air or gas in pleural space, indicating pneumothorax. Air displacement by fluid may lead to opacity.",
+    "tortuous": "the Aorta is slightly tortuous. Sometimes it may refer to varicose veins",
+    "lead": "It refers to the leading head of the tube.",
+    "disease": "It means the exist of diseases and abnormalty, indicating the structure is abnormal. ",
+    "calcification": "Pulmonary calcification is a common asymptomatic finding. Pulmonary calcifications are caused mainly by two mechanisms: the dystrophic form and the metastatic form",
+    "prominence": "It means the exist of some observation.",
+    "device": "It refer to some equipments like picc tub, valve catheter, pacemaker hardware, arthroplastmarker icd defib, device support equipment and mediport'",
+    "engorgement": "pulmonary vascular engorgement means obstruction of the normal flux of blood within the blood vessel network of the lung resulting in engorgement of pulmonary vessels",
+    "picc": "A peripherally inserted central catheter (PICC), also called a PICC line, is a long, thin tube that's inserted through a vein in your arm and passed through to the larger veins near your heart. ",
+    "clip": "Surgical clips or vascular clips usually represent the one kind of medical equipments.",
+    "elevation": "If tissues or anatomical structures are elevated, they are raised up higher than the normal location.",
+    "expand": "It means the lungs are normally expanded and clear, indicating the absence of pneumothorax.",
+    "nodule": "A lung nodule or pulmonary nodule is a relatively small focal density in the lung. it may be confused with the projection of a structure of the chest wall or skin, such as a nipple, a healing rib fracture or  lung cancer.",
+    "wire": "sternotomis wires means the center line of the chest.",
+    "fluid": "It refers to the water of liquid in the lung and it may indicate edema and other diseases.",
+    "degenerative": "Degenerative disease is the result of a continuous process based on degenerative cell changes",
+    "pacemaker": "pacemaker device usually represents the one kind of medical equipments.",
+    "thicken": "Pleural thickening is an increase in the bulkiness of one or both of the pulmonary pleurae. It may cause by pulmonary Infection, empyema, tuberculosis or lung cancer.",
+    "marking": "It represents interstitial markings or bronchovascular markings",
+    "scar": "A scar (or scar tissue) is an area of fibrous tissue that replaces normal tissues after an injury.",
+    "hyperinflate": "Hyperinflated lungs are larger-than-normal lungs as a result of trapped air.",
+    "blunt": "Blunting of the costophrenic angles is usually caused by a pleural effusion, as already discussed. Other causes of costophrenic angle blunting include lung disease in the region of the costophrenic angle, and lung hyperexpansion.",
+    "loss": "The etiology of lung volume loss can be listed as follows: airway obstruction or compression, obesity, scoliosis, restrictive diseases such as pulmonary fibrosis and interstitial lung disease, tuberculosis, sarcoidosis, pleural effusions, rib injury (fractures or diaphragm paralysis), and heart failure",
+    "widen": "The mediastinum is not widened or enlarged",
+    "collapse": "collapse lung refers to pneumothorax or atelectasis.",
+    "density": "The density (more precisely, the volumetric mass density; also known as specific mass), of a substance is its mass per unit volume. ",
+    "emphysema": "Emphysema, or pulmonary emphysema, is a lower respiratory tract disease, characterized by air-filled spaces (pneumatosis) in the lungs, that can vary in size and may be very large",
+    "aerate": "Aeration (also called aerification or aeriation) is the process by which air is circulated through, mixed with or dissolved in a liquid or other substances that act as a fluid (such as soil).",
+    "mass": "A lung mass is an abnormal growth or area in the lungs and it can also view as lung cancer.",
+    "crowd": "Crowding of the bronchovascular structures is an important direct sign of volume loss. The atelectatic lung enhances densely after contrast administration because of closeness of the pulmonary arteries and arterioles within the collapsed lobe.",
+    "infiltrate": "A pulmonary infiltrate is a substance denser than air, such as pus, blood, or protein, which lingers within the parenchyma of the lungs. Pulmonary infiltrates are associated with pneumonia, tuberculosis and sarcoidosis.",
+    "obscure": "Some anatomy structures are not clear and is difficult to understand or see",
+    "deformity": "It means some body parts are abnormal or unjuried.",
+    "hernia": "Lung hernia (Sibson hernia) is a protrusion of lung outside of thoracic wall. the hernia is noted after chest trauma, thoracic surgery or certain pulmonary diseases",
+    "drainage": "Tube drainage represents the one kind of medical equipment",
+    "distention": "Distension generally refers to an enlargement, dilation, or ballooning effect. It may refer to: Abdominal distension,",
+    "shift": "The mediastinal shift is the deviation of the mediastinal structures towards one side of the chest cavity, usually seen on chest radiograph. It indicates a severe asymmetry of intrathoracic pressures.",
+    "stent": "tracheal stent represents the one kind of medical equipments",
+    "pressure": "Pulmonary venous pressure is intermediate between mean PAP and LAP over all physiologic pressures",
+    "lesion": "Lung nodules, pulmonary nodules, white spots, lesions\u2014these terms all describe the same phenomenon: an abnormality in the lungs.",
+    "finding": "Some observation on body parts, usually indicating abnormalty.",
+    "borderline": "borderline size of the cardiac silhouette means the cardiac silhouette is not enlarged and normal.",
+    "hardware": "It represents the one kind of medical equipments.",
+    "dilation": "the state of being larger or more open than normal",
+    "chf": "Heart failure \u2014 sometimes known as congestive heart failure \u2014 occurs when the heart muscle doesn't pump blood as well as it should. When this happens, blood often backs up and fluid can build up in the lungs, causing shortness of breath.",
+    "redistribution": "If the pulmonary edema is due to heart failure or fluid overload, you may also see cardiomegaly and distension of the pulmonary veins, particularly in the upper lung fields.",
+    "aspiration": "Aspiration pneumonia occurs when food or liquid is breathed into the airways or lungs, instead of being swallowed. ",
+    "tail_abnorm_obs": "Some very rare diseases.",
+    "excluded_obs": "Some observations that seldom appear in the reports."
+}

PreTrain_MeDSLIP/data_file/preprocessing/adj_matrix.py

@@ -0,0 +1,891 @@
+"""
+Code copied from AGXNet:
+https://github.com/batmanlab/AGXNet
+"""
+
+"""Create adjacency matrix for representing the relations between anatomical landmarks and observations."""
+
+import argparse
+import pandas as pd
+import numpy as np
+import pickle
+
+from tqdm import tqdm, trange
+from torch.utils.data import Dataset, DataLoader
+
+
+parser = argparse.ArgumentParser(description="Create Adjacency matrix Matrix.")
+
+parser.add_argument(
+    "--input-path",
+    default="/PROJECT DIR/preprocessing/mimic-cxr-radgraph-sentence-parsed.csv",
+    help="Itemized input data path.",
+)
+
+# List of most common normal observations
+NORM_OBS = [
+    "normal",
+    "clear",
+    "sharp",
+    "sharply",
+    "unremarkable",
+    "intact",
+    "stable",
+    "free",
+]
+
+# exclude
+EXCLUDED_OBS = [
+    "none",
+    "unchanged",
+    "change",
+    "great",
+    "similar",
+    "large",
+    "small",
+    "moderate",
+    "mild",
+    "median",
+    "decrease",
+    "bad",
+    "more",
+    "constant",
+    "worsen",
+    "new",
+    "improve",
+    "status",
+    "position",
+    "sternotomy",
+    "cabg",
+    "replacement",
+    "postoperative",
+    "assessment",
+    "patient",
+]
+
+# top 90% abnormal observations
+ABNORM_OBS = [
+    "effusion",
+    "opacity",
+    "pneumothorax",
+    "edema",
+    "atelectasis",
+    "tube",
+    "consolidation",
+    "process",
+    "abnormality",
+    "enlarge",
+    "tip",
+    "low",
+    "pneumonia",
+    "line",
+    "congestion",
+    "catheter",
+    "cardiomegaly",
+    "fracture",
+    "air",
+    "tortuous",
+    "lead",
+    "disease",
+    "calcification",
+    "prominence",
+    "device",
+    "engorgement",
+    "picc",
+    "clip",
+    "elevation",
+    "expand",
+    "nodule",
+    "wire",
+    "fluid",
+    "degenerative",
+    "pacemaker",
+    "thicken",
+    "marking",
+    "scar",
+    "hyperinflate",
+    "blunt",
+    "loss",
+    "widen",
+    "collapse",
+    "density",
+    "emphysema",
+    "aerate",
+    "mass",
+    "crowd",
+    "infiltrate",
+    "obscure",
+    "deformity",
+    "hernia",
+    "drainage",
+    "distention",
+    "shift",
+    "stent",
+    "pressure",
+    "lesion",
+    "finding",
+    "borderline",
+    "hardware",
+    "dilation",
+    "chf",
+    "redistribution",
+    "aspiration",
+]
+
+# final row and column names in adjacent matrix
+LANDMARK_NAME = [
+    "trachea",
+    "left_hilar",
+    "right_hilar",
+    "hilar_unspec",
+    "left_pleural",
+    "right_pleural",
+    "pleural_unspec",
+    "heart_size",
+    "heart_border",
+    "left_diaphragm",
+    "right_diaphragm",
+    "diaphragm_unspec",
+    "retrocardiac",
+    "lower_left_lobe",
+    "upper_left_lobe",
+    "lower_right_lobe",
+    "middle_right_lobe",
+    "upper_right_lobe",
+    "left_lower_lung",
+    "left_mid_lung",
+    "left_upper_lung",
+    "left_apical_lung",
+    "left_lung_unspec",
+    "right_lower_lung",
+    "right_mid_lung",
+    "right_upper_lung",
+    "right_apical_lung",
+    "right_lung_unspec",
+    "lung_apices",
+    "lung_bases",
+    "left_costophrenic",
+    "right_costophrenic",
+    "costophrenic_unspec",
+    "cardiophrenic_sulcus",
+    "mediastinal",
+    "spine",
+    "clavicle",
+    "rib",
+    "stomach",
+    "right_atrium",
+    "right_ventricle",
+    "aorta",
+    "svc",
+    "interstitium",
+    "parenchymal",
+    "cavoatrial_junction",
+    "cardiopulmonary",
+    "pulmonary",
+    "lung_volumes",
+    "unspecified",
+    "other",
+]
+
+OBSERVATION_CLASS = [
+    "normal",
+    "clear",
+    "sharp",
+    "sharply",
+    "unremarkable",
+    "intact",
+    "stable",
+    "free",
+    "effusion",
+    "opacity",
+    "pneumothorax",
+    "edema",
+    "atelectasis",
+    "tube",
+    "consolidation",
+    "process",
+    "abnormality",
+    "enlarge",
+    "tip",
+    "low",
+    "pneumonia",
+    "line",
+    "congestion",
+    "catheter",
+    "cardiomegaly",
+    "fracture",
+    "air",
+    "tortuous",
+    "lead",
+    "disease",
+    "calcification",
+    "prominence",
+    "device",
+    "engorgement",
+    "picc",
+    "clip",
+    "elevation",
+    "expand",
+    "nodule",
+    "wire",
+    "fluid",
+    "degenerative",
+    "pacemaker",
+    "thicken",
+    "marking",
+    "scar",
+    "hyperinflate",
+    "blunt",
+    "loss",
+    "widen",
+    "collapse",
+    "density",
+    "emphysema",
+    "aerate",
+    "mass",
+    "crowd",
+    "infiltrate",
+    "obscure",
+    "deformity",
+    "hernia",
+    "drainage",
+    "distention",
+    "shift",
+    "stent",
+    "pressure",
+    "lesion",
+    "finding",
+    "borderline",
+    "hardware",
+    "dilation",
+    "chf",
+    "redistribution",
+    "aspiration",
+    "tail_abnorm_obs",
+    "excluded_obs",
+]
+
+DICT_ANATOMICAL_LANDMARKS = {
+    "trachea": {"a": ["trachea", "tracheal"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "left_hilar": {
+        "a": ["hilar", "hilum", "perihilar", "infrahilar"],
+        "m1": ["left"],
+        "m2": ["right"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_hilar": {
+        "a": ["hilar", "hilum", "perihilar", "infrahilar"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "hilar_unspec": {
+        "a": ["hilar", "hilum", "perihilar", "infrahilar"],
+        "m1": ["left", "right"],
+        "m2": [],
+        "sc": ["hila", "perihilar|right|left", "perihilar|left|right"],
+        "t": "m1-",
+    },
+    "left_pleural": {
+        "a": ["pleural"],
+        "m1": ["left"],
+        "m2": ["right"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_pleural": {
+        "a": ["pleural"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "pleural_unspec": {
+        "a": ["pleural"],
+        "m1": ["left", "right"],
+        "m2": [],
+        "sc": [
+            "pleural|left|right",
+            "pleural|right|left",
+            "pleural|bilateral|right|left",
+            "pleural|bilateral|left|right",
+        ],
+        "t": "m1-",
+    },
+    "heart_size": {
+        "a": ["heart", "cardiac"],
+        "m1": ["border", "borders"],
+        "m2": [],
+        "sc": [],
+        "t": "m1-",
+    },
+    "heart_border": {
+        "a": ["heart", "cardiac"],
+        "m1": ["border", "borders"],
+        "m2": [],
+        "sc": [],
+        "t": "m1+",
+    },
+    "left_diaphragm": {
+        "a": ["diaphragm", "hemidiaphragm"],
+        "m1": ["left"],
+        "m2": ["right"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_diaphragm": {
+        "a": ["diaphragm", "hemidiaphragm"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "diaphragm_unspec": {
+        "a": ["diaphragm", "diaphragms", "hemidiaphragms", "hemidiaphragm"],
+        "m1": ["left", "right"],
+        "m2": [],
+        "sc": ["hemidiaphragm|left|right", "hemidiaphragm|right|left"],
+        "t": "m1-",
+    },
+    "retrocardiac": {"a": ["retrocardiac"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "lower_left_lobe": {
+        "a": ["lobe"],
+        "m1": ["left"],
+        "m2": ["lower"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "upper_left_lobe": {
+        "a": ["lobe"],
+        "m1": ["left"],
+        "m2": ["upper"],
+        "sc": ["lingula", "lingular"],
+        "t": "m1+m2+",
+    },
+    "lower_right_lobe": {
+        "a": ["lobe"],
+        "m1": ["right"],
+        "m2": ["lower"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "middle_right_lobe": {
+        "a": ["lobe"],
+        "m1": ["right"],
+        "m2": ["middle"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "upper_right_lobe": {
+        "a": ["lobe"],
+        "m1": ["right"],
+        "m2": ["upper"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "left_lower_lung": {
+        "a": ["lung"],
+        "m1": ["left"],
+        "m2": ["lower", "base", "basilar", "basal", "basis"],
+        "sc": ["base|left", "basilar|left", "basal|left", "lung|left|bases"],
+        "t": "m1+m2+",
+    },
+    "left_mid_lung": {
+        "a": ["lung"],
+        "m1": ["left"],
+        "m2": ["middle", "mid"],
+        "sc": ["midlung|left"],
+        "t": "m1+m2+",
+    },
+    "left_upper_lung": {
+        "a": ["lung"],
+        "m1": ["left"],
+        "m2": ["upper"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "left_apical_lung": {
+        "a": ["apex", "apical", "apical", "apicolateral"],
+        "m1": ["left"],
+        "m2": ["right"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "left_lung_unspec": {
+        "a": ["lung", "hemithorax"],
+        "m1": ["left", "left-sided"],
+        "m2": [
+            "volume",
+            "volumes",
+            "right",
+            "lower",
+            "base",
+            "bases",
+            "basilar",
+            "basilar",
+            "basal",
+            "basis",
+            "middle",
+            "mid",
+            "upper",
+            "apex",
+            "apical",
+            "perihilar",
+        ],
+        "sc": ["left", "left side", "thorax|left|hemi"],
+        "t": "m1+m2-",
+    },
+    "right_lower_lung": {
+        "a": ["lung"],
+        "m1": ["right"],
+        "m2": ["lower", "base", "basilar", "basal", "basis"],
+        "sc": ["base|right", "basilar|right", "basal|right", "lung|right|bases"],
+        "t": "m1+m2+",
+    },
+    "right_mid_lung": {
+        "a": ["lung"],
+        "m1": ["right"],
+        "m2": ["middle", "mid"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "right_upper_lung": {
+        "a": ["lung"],
+        "m1": ["right"],
+        "m2": ["upper"],
+        "sc": [],
+        "t": "m1+m2+",
+    },
+    "right_apical_lung": {
+        "a": ["apex", "apical", "apical", "apicolateral"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_lung_unspec": {
+        "a": ["lung", "hemithorax"],
+        "m1": ["right", "right-sided"],
+        "m2": [
+            "volume",
+            "volumes",
+            "left",
+            "lower",
+            "base",
+            "bases",
+            "basilar",
+            "basilar",
+            "basal",
+            "basis",
+            "middle",
+            "mid",
+            "upper",
+            "apex",
+            "apical",
+            "perihilar",
+        ],
+        "sc": ["right", "right side", "thorax|right|hemi"],
+        "t": "m1+m2-",
+    },
+    "lung_apices": {
+        "a": ["apices", "apical"],
+        "m1": ["left", "right"],
+        "m2": [],
+        "sc": ["biapical", "lungs|upper"],
+        "t": "m1-",
+    },
+    "lung_bases": {
+        "a": ["lung", "lungs"],
+        "m1": ["left", "right"],
+        "m2": ["bibasilar", "basilar", "base", "bases", "bibasal", "basal"],
+        "sc": [
+            "lung|lower",
+            "lungs|lower",
+            "bibasilar",
+            "basilar",
+            "bases",
+            "bibasal",
+            "basal",
+            "basal|bilateral",
+            "lobe|lower",
+            "lobes|lower",
+            "lobe|bilateral|lower",
+            "bases|both",
+            "bibasilar|left|right",
+            "bibasilar|right|left",
+        ],
+        "t": "m1-m2+",
+    },
+    "left_costophrenic": {
+        "a": ["costophrenic"],
+        "m1": ["left"],
+        "m2": ["right"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_costophrenic": {
+        "a": ["costophrenic"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "costophrenic_unspec": {
+        "a": ["costophrenic"],
+        "m1": ["left", "right"],
+        "m2": [],
+        "sc": [],
+        "t": "m1-",
+    },
+    "cardiophrenic_sulcus": {
+        "a": ["cardiophrenic"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "mediastinal": {
+        "a": ["mediastinal", "cardiomediastinal", "mediastinum", "cardiomediastinum"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "spine": {"a": ["spine", "spinal"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "clavicle": {
+        "a": ["clavicle", "clavicles"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "rib": {"a": ["rib", "ribs"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "stomach": {
+        "a": ["stomach", "abdomen", "abdominal"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "right_atrium": {
+        "a": ["atrium", "atrial"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "right_ventricle": {
+        "a": ["ventricle", "ventricular"],
+        "m1": ["right"],
+        "m2": ["left"],
+        "sc": [],
+        "t": "m1+m2-",
+    },
+    "aorta": {"a": ["aorta", "aortic"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "svc": {"a": ["svc"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "interstitium": {
+        "a": ["interstitium", "interstitial"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "parenchymal": {"a": ["parenchymal"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "cavoatrial_junction": {
+        "a": ["cavoatrial junction"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "cardiopulmonary": {
+        "a": ["cardiopulmonary"],
+        "m1": [],
+        "m2": [],
+        "sc": [],
+        "t": "m0",
+    },
+    "pulmonary": {"a": ["pulmonary"], "m1": [], "m2": [], "sc": [], "t": "m0"},
+    "lung_volumes": {
+        "a": ["lungs", "lung", "volume", "volumes"],
+        "m1": [
+            "left",
+            "right",
+            "lower",
+            "base",
+            "bases",
+            "basilar",
+            "basal",
+            "basis",
+            "middle",
+            "mid",
+            "upper",
+            "apex",
+            "apical",
+            "apical",
+        ],
+        "m2": [],
+        "sc": [],
+        "t": "m1-",
+    },
+}
+
+
+class LandmarkObservationAdjacentMatrix(Dataset):
+    def __init__(self, LANDMARK_NAME, OBSERVATION_CLASS, df_anatomy_label):
+        self.LANDMARK_NAME = LANDMARK_NAME
+        self.OBSERVATION_CLASS = OBSERVATION_CLASS
+        self.df_anatomy_label = df_anatomy_label
+
+        # get all study ids
+        self.sids = list(self.df_anatomy_label["study_id"].unique())
+
+    def __getitem__(self, idx):
+        sid = self.sids[idx]
+        df_sid = self.df_anatomy_label[self.df_anatomy_label["study_id"] == sid]
+        landmark_observation_adj_mtx = (
+            np.zeros((len(LANDMARK_NAME), len(OBSERVATION_CLASS))) - 1.0
+        )
+        for index, row in df_sid.iterrows():
+            try:
+                observation_idx = self.OBSERVATION_CLASS.index(
+                    row.obs_lemma_grp
+                )  # if a rare observation, skip this instance
+                landmark_idx = self.LANDMARK_NAME.index(row.landmark_name)
+
+                curr_val = landmark_observation_adj_mtx[landmark_idx, observation_idx]
+
+                # for obs_lemma_grp, such as tail_abnorm_obs
+                # if one observation is DP, then 1.0
+                if row.label == "OBS-DP":
+                    landmark_observation_adj_mtx[landmark_idx, observation_idx] = 1.0
+                elif row.label == "OBS-DA":
+                    landmark_observation_adj_mtx[
+                        landmark_idx, observation_idx
+                    ] = np.maximum(curr_val, 0.0)
+            except:
+                pass
+        return sid, landmark_observation_adj_mtx
+
+    def __len__(self):
+        return len(self.sids)
+
+
+def anatomy_to_landmark(x, a, m1=[], m2=[], sc=[], t="m0"):
+    """
+    Args:
+        x: input anatomy, e.g., "lobe|left|lower"
+        a: base anatomy set, e.g., ["hilar", "hilum", "perihilar"]
+        m1: level 1 modifier, e.g., ["left", "right"]
+        m2: level 2 modifier, e.g., ["upper", "middle", "lower"]
+        s: special cases, e.g., ["chest"]
+        t: type, ["m2+", "m1+m2-"]
+    Return:
+        flag: boolean, matched or not matched
+    """
+    s = set(x.split("|"))
+    if t == "m1+m2+":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) > 0) & (len(s & set(m2)) > 0)
+    elif t == "m1+m2-":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) > 0) & (len(s & set(m2)) == 0)
+    elif t == "m1-m2+":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) == 0) & (len(s & set(m2)) > 0)
+    elif t == "m1-m2-":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) == 0) & (len(s & set(m2)) == 0)
+    elif t == "m1+":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) > 0)
+    elif t == "m2+":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m2)) > 0)
+    elif t == "m1-":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m1)) == 0)
+    elif t == "m2-":
+        flag = (len(s & set(a)) > 0) & (len(s & set(m2)) == 0)
+    elif t == "m0":
+        flag = len(s & set(a)) > 0
+
+    if sc:
+        flag = flag | (x in sc)
+    return flag
+
+
+def create_adj_matrix(args):
+    # load anatomy label table, text table and master table
+    print("Loading parsed RadGraph data...")
+    df_anatomy_label = pd.read_csv(args.input_path, dtype=str)
+
+    # manual lemmatization correction
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["enlargement", "increase"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "enlarge"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["engorge"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "engorgement"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["opacification", "opacity-"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "opacity"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["calcify"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "calcification"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["effusion ;"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "effusion"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(
+        ["atelectatic", "atelectasis ;", "atelectase"]
+    )
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "atelectasis"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["aeration"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "aerate"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["distend", "distension"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "distention"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["wide"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "widen"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["prominent"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "prominence"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["haze"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "haziness"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["masse"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "mass"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["kyphotic"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "kyphosis"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["degenerate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "degenerative"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["obscuration"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "obscure"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["fibrotic"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "fibrosis"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["nodular", "nodularity"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "nodule"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["ventilate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "ventilation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["tortuosity"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "tortuous"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["elongate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "elongation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["elevate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "elevation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["drain"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "drainage"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["deviate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "deviation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["consolidative", "consolidate"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "consolidation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["dilate", "dilatation"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "dilation"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(
+        ["hydropneumothorax", "pneumothoraces", "pneumothorace"]
+    )
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "pneumothorax"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(["improvement", "improved"])
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "improve"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(
+        [
+            "can not be assess",
+            "can not be evaluate",
+            "not well see",
+            "not well assess",
+            "can not be accurately assess",
+            "not well evaluate",
+            "not well visualize",
+            "difficult to evaluate",
+            "poorly see",
+        ]
+    )
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "difficult to assess"
+
+    idx_replace = df_anatomy_label["obs_lemma"] == "pacer"
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "pacemaker"
+
+    idx_replace = df_anatomy_label["obs_lemma"].isin(
+        ["infection", "infectious", "infectious process"]
+    )
+    df_anatomy_label.loc[idx_replace, "obs_lemma"] = "pneumonia"
+
+    df_anatomy_label.loc[df_anatomy_label["label"].isna(), "label"] = "OBS-NA"
+
+    # step 1: map anatomy name to landmark name
+    landmark_name = []
+    for index, row in tqdm(
+        df_anatomy_label.iterrows(), total=df_anatomy_label.shape[0]
+    ):
+        x = row.anatomy
+        flag = False
+        for k, v in DICT_ANATOMICAL_LANDMARKS.items():
+            flag = anatomy_to_landmark(x, v["a"], v["m1"], v["m2"], v["sc"], v["t"])
+            if flag:
+                landmark_name.append(k)
+                break
+        if (not flag) & (row.anatomy == "unspecified"):
+            landmark_name.append("unspecified")
+        elif (not flag) & (row.anatomy != "unspecified"):
+            landmark_name.append("other")
+
+    df_anatomy_label["landmark_name"] = landmark_name
+
+    # create a new obs_lemma column to grouop other abnormal observation class
+    df_anatomy_label["obs_lemma_grp"] = df_anatomy_label["obs_lemma"]
+
+    idx1 = df_anatomy_label["obs_lemma"].isin(NORM_OBS)
+    idx2 = df_anatomy_label["obs_lemma"].isin(ABNORM_OBS)
+    idx3 = df_anatomy_label["obs_lemma"].isin(EXCLUDED_OBS)
+
+    df_anatomy_label.loc[idx3, "obs_lemma_grp"] = "excluded_obs"
+
+    idx = (~idx1) & (~idx2) & (~idx3)  # abnormal observations that are in the tail
+    df_anatomy_label.loc[idx, "obs_lemma_grp"] = "tail_abnorm_obs"
+
+    # step 2: get landmark - observation adjacent matrix
+    dataset = LandmarkObservationAdjacentMatrix(
+        LANDMARK_NAME, OBSERVATION_CLASS, df_anatomy_label
+    )
+    loader = DataLoader(
+        dataset, batch_size=32, shuffle=False, num_workers=8, drop_last=False
+    )
+
+    sid_lst = []
+    adj_mtx_lst = []
+    for index, data in tqdm(enumerate(loader), total=len(loader)):
+        sid, landmark_observation_adj_mtx = data
+        sid_lst.append(sid)
+        adj_mtx_lst.append(landmark_observation_adj_mtx)
+
+    # step 3: convert outputs to a dictionary and then save to a pickel file
+    full_sids = np.concatenate(sid_lst, axis=0)
+    full_adj_mtx = np.concatenate(adj_mtx_lst, axis=0)
+    dict_adj_mtx = {}
+    for i in trange(len(full_sids)):
+        sid = full_sids[i]
+        dict_adj_mtx[sid] = full_adj_mtx[i]
+
+    np.save("landmark_observation_sids.npy", full_sids)
+    print("landmark_observation_sids.npy has been saved!")
+    np.save("landmark_observation_adj_mtx.npy", full_adj_mtx)
+    print("landmark_observation_sids.npy has been saved!")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    create_adj_matrix(args)

PreTrain_MeDSLIP/data_file/preprocessing/radgraph_itemized.py

@@ -0,0 +1,139 @@
+"""
+Code copied from AGXNet:
+https://github.com/batmanlab/AGXNet
+"""
+
+import argparse
+import pandas as pd
+import json
+from tqdm import tqdm
+import nltk
+
+
+parser = argparse.ArgumentParser(description="Itemize RadGraph Dataset.")
+
+parser.add_argument(
+    "--data-path",
+    default="/PATH TO RADGRAPH DATA/RadGraph/physionet.org/files/radgraph/1.0.0/MIMIC-CXR_graphs.json",
+    help="RadGraph data path.",
+)
+parser.add_argument(
+    "--output-path",
+    default="/PROJECT DIR/preprocessing/mimic-cxr-radgraph-itemized.csv",
+    help="Output path for itemized RadGraph data.",
+)
+
+
+def get_ids(key):
+    """Convert keys in the RadGraph file into IDs"""
+    lst = key.split("/")
+    partition = lst[0]  # dataset partition
+    pid = lst[1][1:]  # patient id
+    sid = lst[2].split(".")[0][1:]  # study id, remove .txt
+    return partition, pid, sid
+
+
+def get_sen_from_token_ix(text, ix):
+    """get the sentence to which the input token index belongs."""
+    sen_lst = nltk.sent_tokenize(text)
+    dict_ws = {}
+    ix_w = 0
+    ix_s = 0
+    for s in sen_lst:
+        words = nltk.word_tokenize(s)
+        for w in words:
+            dict_ws[ix_w] = ix_s
+            ix_w += 1
+        ix_s += 1
+    return dict_ws[ix], sen_lst[dict_ws[ix]]
+
+
+def get_entity_relation(value):
+    """itemize each relation"""
+    source_lst = []
+    target_lst = []
+    token_lst = []
+    token_ix_lst = []
+    label_lst = []
+    relation_lst = []
+    sen_lst = []
+    sen_ix_lst = []
+
+    text = value["text"]
+
+    entities = value["entities"]
+    for k, v in entities.items():
+        six, sen = get_sen_from_token_ix(text, v["start_ix"])
+        relations = v["relations"]
+
+        # source node has no out going edge
+        if (len(relations) == 0) or (relations[0] is None):
+            source_lst.append(k)
+            token_ix_lst.append(v["start_ix"])
+            token_lst.append(v["tokens"])
+            label_lst.append(v["label"])
+            relation_lst.append(None)
+            target_lst.append(None)
+            sen_ix_lst.append(six)
+            sen_lst.append(sen)
+        else:
+            for r in relations:
+                source_lst.append(k)
+                token_ix_lst.append(v["start_ix"])
+                token_lst.append(v["tokens"])
+                label_lst.append(v["label"])
+                relation_lst.append(r[0])
+                target_lst.append(r[1])
+                sen_ix_lst.append(six)
+                sen_lst.append(sen)
+
+    # save outputs in a dataframe
+    return pd.DataFrame(
+        {
+            "source": source_lst,
+            "token": token_lst,
+            "token_ix": token_ix_lst,
+            "label": label_lst,
+            "relation": relation_lst,
+            "target": target_lst,
+            "sentence_ix": sen_ix_lst,
+            "sentence": sen_lst,
+        }
+    )
+
+
+def radgraph_itemize(args):
+    """Convert nested RadGraph data to itemized examples."""
+
+    print("Loading RadGraph data...")
+    f = open(args.data_path)
+    data = json.load(f)
+    print("RadGraph data is loaded.")
+
+    # create itemized RadGraph data
+    df_lst = []
+    pid_lst = []
+    sid_lst = []
+    text_lst = []
+    print("Itemizing RadGraph data...")
+    for key, value in tqdm(data.items()):
+        _, pid, sid = get_ids(key)
+        pid_lst.append(pid)
+        sid_lst.append(sid)
+        text_lst.append(data[key]["text"])
+        df = get_entity_relation(value)
+        df["subject_id"] = pid
+        df["study_id"] = sid
+        df_lst.append(df)
+
+    # entity level dataframe
+    df_itemized = pd.concat(df_lst)
+
+    # save dataframes to a .csv file
+    df_itemized.to_csv(args.output_path, index=False)
+    print("Outputs have been saved!")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    radgraph_itemize(args)

PreTrain_MeDSLIP/data_file/preprocessing/radgraph_parsed.py

@@ -0,0 +1,322 @@
+"""
+Code copied from AGXNet:
+https://github.com/batmanlab/AGXNet
+"""
+
+import argparse
+import pandas as pd
+from tqdm import tqdm
+import spacy
+
+sp = spacy.load("en_core_web_sm")
+
+parser = argparse.ArgumentParser(description="Pharse RadGraph Relations.")
+
+parser.add_argument(
+    "--input-path",
+    default="/PROJECT DIR/preprocessing/mimic-cxr-radgraph-itemized.csv",
+    help="Itemized input data path.",
+)
+parser.add_argument(
+    "--output-path",
+    default="/PROJECT DIR/preprocessing/mimic-cxr-radgraph-sentence-parsed.csv",
+    help="Output path for parsed relations.",
+)
+
+
+def obs_lemmatization(x):
+    """
+    Lemmatize observation
+    Args:
+        x: a observation token
+    Return:
+        normalized observation
+    """
+    w_lst = []
+    for word in sp(str(x)):
+        w_lst.append(word.lemma_)
+    return " ".join(w_lst)
+
+
+def radgraph_parse(args):
+    """Pharse RadGraph relations."""
+
+    print("Loading itemized RadGraph data...")
+    df_itemized = pd.read_csv(args.input_path)
+
+    # get all study_id
+    sid_lst = list(df_itemized["study_id"].unique())
+
+    tuple_lst = []
+    print("Preprocessing sentences...")
+    for sid in tqdm(sid_lst):
+        idx_s = df_itemized["study_id"] == sid
+        df_sid = df_itemized[idx_s]
+
+        # unique sentence index
+        sen_ids = list(df_sid["sentence_ix"].unique())
+
+        for si in sen_ids:
+            idx_sen = df_sid["sentence_ix"] == si
+            df_sen = df_sid[idx_sen]
+            sen = df_sen["sentence"].iloc[0]
+
+            # step 1, select all target anatomy entities (e.g., lobe) with label = ANAT-DP and target = NaN
+            idx_a = (df_sen["label"] == "ANAT-DP") & (df_sen["target"].isnull())
+            df_a = df_sen[idx_a]
+
+            if sum(idx_a) > 0:
+                for _, row_a in df_a.iterrows():
+                    anatomy_source_keys = []
+                    sen = row_a.sentence
+                    source_key = row_a.source
+
+                    # step 2, get detailed target anatomy (e.g., lower left lobe)
+                    token_a = [row_a["token"].lower()]
+                    anatomy_source_keys.append(source_key)
+                    idx_t = (df_sen["label"] == "ANAT-DP") & (
+                        df_sen["target"] == source_key
+                    )
+                    if sum(idx_t) > 0:
+                        df_t = df_sen[idx_t]
+                        for _, row in df_t.iterrows():
+                            token_a += [row["token"].lower()]
+                            anatomy_source_keys.append(
+                                row["source"]
+                            )  # save keys of all anatomy token, i.e., lower, left, lobe
+                        anatomy = "|".join(token_a)
+
+                    else:
+                        anatomy = row_a["token"].lower()
+
+                    # step 3: get observations associated with the target anatomy (e.g., normal, effusion)
+                    idx_o = (
+                        (df_sen["label"].isin(["OBS-DA", "OBS-DP", "OBS-U"]))
+                        & (df_sen["target"].isin(anatomy_source_keys))
+                        & (df_sen["relation"] == "located_at")
+                    )
+                    if sum(idx_o) > 0:
+                        df_o = df_sen[idx_o]
+
+                        anatomy_lst = []
+                        obs_lst = []
+                        label_lst = []
+                        obs_modify_lst = []
+                        obs_suggestive_lst = []
+
+                        for _, row_o in df_o.iterrows():
+                            anatomy_lst.append(anatomy)
+                            obs_lst.append(row_o["token"].lower())
+                            label_lst.append(row_o["label"])
+
+                            # step 4: get obs modification
+                            idx_o_m = (df_sen["target"] == row_o.source) & (
+                                df_sen["relation"] == "modify"
+                            )
+                            obs_modify = None
+                            if sum(idx_o_m) > 0:
+                                df_o_m = df_sen[idx_o_m]
+                                temp_lst = []
+                                for _, row_om in df_o_m.iterrows():
+                                    # if the modification is present
+                                    if row_om.label == "OBS-DP":
+                                        temp_lst.append(row_om["token"].lower())
+                                if len(temp_lst) > 0:
+                                    obs_modify = "|".join(temp_lst)
+                            obs_modify_lst.append(obs_modify)
+
+                            # step 5: get suggestive of obs
+                            idx_o_s = (df_sen["target"] == row_o.source) & (
+                                df_sen["relation"] == "suggestive_of"
+                            )
+                            obs_suggestive = None
+                            if sum(idx_o_s) > 0:
+                                df_o_s = df_sen[idx_o_s]
+                                temp_lst = []
+                                for _, row_os in df_o_s.iterrows():
+                                    # if the modification is present
+                                    if row_os.label == "OBS-DP":
+                                        temp_lst.append(row_os["token"].lower())
+                                if len(temp_lst) > 0:
+                                    obs_suggestive = "|".join(temp_lst)
+                            obs_suggestive_lst.append(obs_suggestive)
+
+                    else:
+                        anatomy_lst = [anatomy]
+                        obs_lst = [None]
+                        label_lst = [None]
+                        obs_modify_lst = [None]
+                        obs_suggestive_lst = [None]
+
+                    # step 4: get observations that are not associated with the target anatomy
+                    idx_oo = (
+                        (df_sen["label"].isin(["OBS-DA", "OBS-DP", "OBS-U"]))
+                        & (df_sen["target"].isna())
+                        & (df_sen["relation"].isna())
+                    )
+                    if sum(idx_oo) > 0:
+                        df_oo = df_sen[idx_oo]
+                        for _, row_oo in df_oo.iterrows():
+                            anatomy_lst.append("unspecified")
+                            obs_lst.append(row_oo["token"].lower())
+                            label_lst.append(row_oo["label"])
+                            # obs_modify_lst.append(None)
+                            # obs_suggestive_lst.append(None)
+
+                            # step 5: get obs modification
+                            idx_o_m = (df_sen["target"] == row_oo.source) & (
+                                df_sen["relation"] == "modify"
+                            )
+                            obs_modify = None
+                            if sum(idx_o_m) > 0:
+                                df_o_m = df_sen[idx_o_m]
+                                temp_lst = []
+                                for _, row_om in df_o_m.iterrows():
+                                    # if the modification is present
+                                    if row_om.label == "OBS-DP":
+                                        temp_lst.append(row_om["token"].lower())
+                                if len(temp_lst) > 0:
+                                    obs_modify = "|".join(temp_lst)
+                            obs_modify_lst.append(obs_modify)
+
+                            # step 5: get suggestive of obs
+                            idx_o_s = (df_sen["target"] == row_oo.source) & (
+                                df_sen["relation"] == "suggestive_of"
+                            )
+                            obs_suggestive = None
+                            if sum(idx_o_s) > 0:
+                                df_o_s = df_sen[idx_o_s]
+                                temp_lst = []
+                                for _, row_os in df_o_s.iterrows():
+                                    # if the modification is present
+                                    if row_os.label == "OBS-DP":
+                                        temp_lst.append(row_os["token"].lower())
+                                if len(temp_lst) > 0:
+                                    obs_suggestive = "|".join(temp_lst)
+                            obs_suggestive_lst.append(obs_suggestive)
+
+                    # step 6: create tuple of 7 values (sid, sentence_id, sentence, anatomy, obs, label)
+                    t_lst = []
+                    for i in range(len(obs_lst)):
+                        t_lst.append(
+                            (
+                                sid,
+                                si,
+                                sen,
+                                anatomy_lst[i],
+                                obs_lst[i],
+                                label_lst[i],
+                                obs_modify_lst[i],
+                                obs_suggestive_lst[i],
+                            )
+                        )
+
+                    # remove duplicates caused by 1 obs "located_at" multiple anatomies
+                    tuple_lst.append(list(set(t_lst)))
+
+            # if the sentence does not have any ANATOMY token
+            else:
+                idx_o = (df_sen["label"].isin(["OBS-DA", "OBS-DP", "OBS-U"])) & (
+                    df_sen["target"].isnull()
+                )
+                if sum(idx_o) > 0:
+                    df_o = df_sen[idx_o]
+
+                    obs_lst = []
+                    label_lst = []
+                    obs_modify_lst = []
+                    obs_suggestive_lst = []
+
+                    for _, row_o in df_o.iterrows():
+                        obs_lst.append(row_o["token"].lower())
+                        label_lst.append(row_o["label"])
+
+                        # step 4: get obs modification
+                        idx_o_m = (df_sen["target"] == row_o.source) & (
+                            df_sen["relation"] == "modify"
+                        )
+                        obs_modify = None
+                        if sum(idx_o_m) > 0:
+                            df_o_m = df_sen[idx_o_m]
+                            temp_lst = []
+                            for _, row_om in df_o_m.iterrows():
+                                # if the modification is present
+                                if row_om.label == "OBS-DP":
+                                    temp_lst.append(row_om["token"].lower())
+                            if len(temp_lst) > 0:
+                                obs_modify = "|".join(temp_lst)
+                        obs_modify_lst.append(obs_modify)
+
+                        # step 5: get suggestive of obs
+                        idx_o_s = (df_sen["target"] == row_o.source) & (
+                            df_sen["relation"] == "suggestive_of"
+                        )
+                        obs_suggestive = None
+                        if sum(idx_o_s) > 0:
+                            df_o_s = df_sen[idx_o_s]
+                            temp_lst = []
+                            for _, row_os in df_o_s.iterrows():
+                                # if the modification is present
+                                if row_os.label == "OBS-DP":
+                                    temp_lst.append(row_os["token"].lower())
+                            if len(temp_lst) > 0:
+                                obs_suggestive = "|".join(temp_lst)
+                        obs_suggestive_lst.append(obs_suggestive)
+                else:
+                    obs_lst = [None]
+                    label_lst = [None]
+                    obs_modify_lst = [None]
+                    obs_suggestive_lst = [None]
+
+                # step 6: create tuple of 7 values (sid, sentence_id, sentence, anatomy, obs, label)
+                t_lst = []
+                for i in range(len(obs_lst)):
+                    t_lst.append(
+                        (
+                            sid,
+                            si,
+                            sen,
+                            "unspecified",
+                            obs_lst[i],
+                            label_lst[i],
+                            obs_modify_lst[i],
+                            obs_suggestive_lst[i],
+                        )
+                    )
+
+                # remove duplicates if existing
+                tuple_lst.append(list(set(t_lst)))
+
+    # flatten nested list
+    df_lst = [item for sublist in tuple_lst for item in sublist]
+    df_anatomy_label = pd.DataFrame(
+        df_lst,
+        columns=[
+            "study_id",
+            "sen_id",
+            "sentence",
+            "anatomy",
+            "observation",
+            "label",
+            "obs_modify",
+            "obs_suggestive",
+        ],
+    )
+
+    # lemmatize observation tokens (e.g., normalize opacities to opacity)
+    obs_lemma_lst = []
+    print("Lemmatizing observation tokens...")
+    for t in tqdm(df_lst):
+        obs = t[4]
+        obs_lemma = obs_lemmatization(obs)
+        obs_lemma_lst.append(obs_lemma)
+
+    # save preprocessed sentence level data
+    df_anatomy_label["obs_lemma"] = obs_lemma_lst
+    df_anatomy_label.to_csv(args.output_path, index=False)
+    print("Output file has been saved!")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    radgraph_parse(args)

PreTrain_MeDSLIP/dataset/dataset.py

@@ -0,0 +1,310 @@
+import json
+from torch.utils.data import DataLoader
+import PIL
+from torch.utils.data import Dataset
+import numpy as np
+import pandas as pd
+from torchvision import transforms
+from PIL import Image
+import random
+from dataset.randaugment import RandomAugment
+
+
+class MeDSLIP_Dataset(Dataset):
+    def __init__(self, csv_path, np_path, mode="train", num_neg_samples=7):
+        self.num_neg_samples = num_neg_samples
+        self.ann = json.load(open(csv_path, "r"))
+        self.img_path_list = list(self.ann)
+        self.anaomy_list = [
+            "trachea",
+            "left_hilar",
+            "right_hilar",
+            "hilar_unspec",
+            "left_pleural",
+            "right_pleural",
+            "pleural_unspec",
+            "heart_size",
+            "heart_border",
+            "left_diaphragm",
+            "right_diaphragm",
+            "diaphragm_unspec",
+            "retrocardiac",
+            "lower_left_lobe",
+            "upper_left_lobe",
+            "lower_right_lobe",
+            "middle_right_lobe",
+            "upper_right_lobe",
+            "left_lower_lung",
+            "left_mid_lung",
+            "left_upper_lung",
+            "left_apical_lung",
+            "left_lung_unspec",
+            "right_lower_lung",
+            "right_mid_lung",
+            "right_upper_lung",
+            "right_apical_lung",
+            "right_lung_unspec",
+            "lung_apices",
+            "lung_bases",
+            "left_costophrenic",
+            "right_costophrenic",
+            "costophrenic_unspec",
+            "cardiophrenic_sulcus",
+            "mediastinal",
+            "spine",
+            "clavicle",
+            "rib",
+            "stomach",
+            "right_atrium",
+            "right_ventricle",
+            "aorta",
+            "svc",
+            "interstitium",
+            "parenchymal",
+            "cavoatrial_junction",
+            "cardiopulmonary",
+            "pulmonary",
+            "lung_volumes",
+            "unspecified",
+            "other",
+        ]
+        self.obs_list = [
+            "normal",
+            "clear",
+            "sharp",
+            "sharply",
+            "unremarkable",
+            "intact",
+            "stable",
+            "free",
+            "effusion",
+            "opacity",
+            "pneumothorax",
+            "edema",
+            "atelectasis",
+            "tube",
+            "consolidation",
+            "process",
+            "abnormality",
+            "enlarge",
+            "tip",
+            "low",
+            "pneumonia",
+            "line",
+            "congestion",
+            "catheter",
+            "cardiomegaly",
+            "fracture",
+            "air",
+            "tortuous",
+            "lead",
+            "disease",
+            "calcification",
+            "prominence",
+            "device",
+            "engorgement",
+            "picc",
+            "clip",
+            "elevation",
+            "expand",
+            "nodule",
+            "wire",
+            "fluid",
+            "degenerative",
+            "pacemaker",
+            "thicken",
+            "marking",
+            "scar",
+            "hyperinflate",
+            "blunt",
+            "loss",
+            "widen",
+            "collapse",
+            "density",
+            "emphysema",
+            "aerate",
+            "mass",
+            "crowd",
+            "infiltrate",
+            "obscure",
+            "deformity",
+            "hernia",
+            "drainage",
+            "distention",
+            "shift",
+            "stent",
+            "pressure",
+            "lesion",
+            "finding",
+            "borderline",
+            "hardware",
+            "dilation",
+            "chf",
+            "redistribution",
+            "aspiration",
+            "tail_abnorm_obs",
+            "excluded_obs",
+        ]
+        self.rad_graph_results = np.load(np_path)
+        normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+        if mode == "train":
+            self.transform = transforms.Compose(
+                [
+                    transforms.RandomResizedCrop(
+                        224, scale=(0.2, 1.0), interpolation=Image.BICUBIC
+                    ),
+                    transforms.RandomHorizontalFlip(),
+                    RandomAugment(
+                        2,
+                        7,
+                        isPIL=True,
+                        augs=[
+                            "Identity",
+                            "AutoContrast",
+                            "Equalize",
+                            "Brightness",
+                            "Sharpness",
+                            "ShearX",
+                            "ShearY",
+                            "TranslateX",
+                            "TranslateY",
+                            "Rotate",
+                        ],
+                    ),
+                    transforms.ToTensor(),
+                    normalize,
+                ]
+            )
+        if mode == "test":
+            self.transform = transforms.Compose(
+                [
+                    transforms.Resize([224, 224]),
+                    transforms.ToTensor(),
+                    normalize,
+                ]
+            )
+
+    def __getitem__(self, index):
+        img_path = self.img_path_list[index]
+        class_label = self.rad_graph_results[
+            self.ann[img_path]["labels_id"], :, :
+        ] 
+        labels_pathology = np.zeros(class_label.shape[-1]) - 1
+        labels_anatomy = np.zeros(class_label.shape[0]) - 1
+        labels_pathology, index_list_pathology = self.triplet_extraction_pathology(
+            class_label
+        )
+        labels_anatomy, index_list_anatomy = self.triplet_extraction_anatomy(
+            class_label
+        )
+        index_list_pathology = np.array(index_list_pathology)
+        index_list_anatomy = np.array(index_list_anatomy)
+
+        img = PIL.Image.open(img_path).convert("RGB")
+        image = self.transform(img)
+
+        return {
+            "image": image,
+            "label_pathology": labels_pathology,
+            "index_pathology": index_list_pathology,
+            "label_anatomy": labels_anatomy,
+            "index_anatomy": index_list_anatomy,
+            "matrix": class_label,
+        }
+
+    def triplet_extraction_pathology(self, class_label):
+        """
+        This is for ProtoCL. Therefore, we need to extract anatomies to use in pathology stream.
+        """
+
+        exist_labels = np.zeros(class_label.shape[-1]) - 1
+        anatomy_list = []
+        for i in range(class_label.shape[1]):
+            temp_list = []
+            ### extract the exist label for each pathology and maintain -1 if not mentioned. ###
+            if 0 in class_label[:, i]:
+                exist_labels[i] = 0
+
+            if 1 in class_label[:, i]:
+                exist_labels[i] = 1
+                ### if the pathology exists try to get its anatomy.###
+                ### Note that, the contrastive loss will only be caculated on exist pathology as it is meaningless to predict their anatomy for the non-exist entities###
+                temp_list.append(-1)
+
+                try:
+                    temp_list = temp_list + random.sample(
+                        np.where(class_label[:, i] != 1)[0].tolist(),
+                        self.num_neg_samples,
+                    )
+                except:
+                    print("fatal error")
+            if temp_list == []:
+                temp_list = temp_list + random.sample(
+                    np.where(class_label[:, i] != 1)[0].tolist(),
+                    self.num_neg_samples + 1,
+                )
+            anatomy_list.append(temp_list)
+
+        return exist_labels, anatomy_list
+
+    def triplet_extraction_anatomy(self, class_label):
+        """
+        This is for ProtoCL. Therefore, we need to extract pathological labels to use in anatomy stream.
+        """
+        exist_labels = np.zeros(class_label.shape[0]) - 1
+        pathology_list = []
+        for i in range(class_label.shape[0]):
+            temp_list = []
+            ### extract the exist label for each pathology and maintain -1 if not mentioned. ###
+            if 0 in class_label[i, :]:
+                exist_labels[i] = 0
+
+            if 1 in class_label[i, :]:
+                exist_labels[i] = 1
+                ### if the pathology exists try to get its anatomy.###
+                ### Note that, the contrastive loss will only be caculated on exist pathology as it is meaningless to predict their anatomy for the non-exist entities###
+                temp_list.append(-1)
+
+                try:
+                    temp_list = temp_list + random.sample(
+                        np.where(class_label[i, :] != 1)[0].tolist(),
+                        self.num_neg_samples,
+                    )
+                except:
+                    print("fatal error")
+            if temp_list == []:
+                temp_list = temp_list + random.sample(
+                    np.where(class_label[i, :] != 1)[0].tolist(),
+                    self.num_neg_samples + 1,
+                )
+            pathology_list.append(temp_list)
+
+        return exist_labels, pathology_list
+
+    def __len__(self):
+        return len(self.ann)
+
+
+def create_loader(datasets, samplers, batch_size, num_workers, is_trains, collate_fns):
+    loaders = []
+    for dataset, sampler, bs, n_worker, is_train, collate_fn in zip(
+        datasets, samplers, batch_size, num_workers, is_trains, collate_fns
+    ):
+        if is_train:
+            shuffle = sampler is None
+            drop_last = True
+        else:
+            shuffle = False
+            drop_last = False
+        loader = DataLoader(
+            dataset,
+            batch_size=bs,
+            num_workers=n_worker,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=shuffle,
+            collate_fn=collate_fn,
+            drop_last=drop_last,
+        )
+        loaders.append(loader)
+    return loaders

PreTrain_MeDSLIP/dataset/randaugment.py

@@ -0,0 +1,346 @@
+import cv2
+import numpy as np
+
+
+## aug functions
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    """
+        same output as PIL.ImageOps.autocontrast
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            offset = -low * scale
+            table = np.arange(n_bins) * scale + offset
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    """
+        same output as PIL.ImageOps.equalize
+        PIL's implementation is different from cv2.equalize
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0:
+            return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    """
+    like PIL, rotate by degree, not radians
+    """
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    """
+        same output as PIL.ImageOps.posterize
+    """
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Color
+    """
+    ## implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = np.float32(
+        [[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]]
+    ) * factor + np.float32([[0.114], [0.587], [0.299]])
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = (
+        np.array([(el - mean) * factor + mean for el in range(256)])
+        .clip(0, 255)
+        .astype(np.uint8)
+    )
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    """
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    """
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, factor, 0], [0, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    """
+        same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, -offset], [0, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    """
+        same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [0, 1, -offset]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    """
+        same output as PIL.ImageOps.posterize
+    """
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [factor, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+### level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level,)
+
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level,)
+
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    "Identity": identity_func,
+    "AutoContrast": autocontrast_func,
+    "Equalize": equalize_func,
+    "Rotate": rotate_func,
+    "Solarize": solarize_func,
+    "Color": color_func,
+    "Contrast": contrast_func,
+    "Brightness": brightness_func,
+    "Sharpness": sharpness_func,
+    "ShearX": shear_x_func,
+    "TranslateX": translate_x_func,
+    "TranslateY": translate_y_func,
+    "Posterize": posterize_func,
+    "ShearY": shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    "Identity": none_level_to_args,
+    "AutoContrast": none_level_to_args,
+    "Equalize": none_level_to_args,
+    "Rotate": rotate_level_to_args(MAX_LEVEL, replace_value),
+    "Solarize": solarize_level_to_args(MAX_LEVEL),
+    "Color": enhance_level_to_args(MAX_LEVEL),
+    "Contrast": enhance_level_to_args(MAX_LEVEL),
+    "Brightness": enhance_level_to_args(MAX_LEVEL),
+    "Sharpness": enhance_level_to_args(MAX_LEVEL),
+    "ShearX": shear_level_to_args(MAX_LEVEL, replace_value),
+    "TranslateX": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "TranslateY": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "Posterize": posterize_level_to_args(MAX_LEVEL),
+    "ShearY": shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+    def __init__(self, N=2, M=10, isPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        if augs:
+            self.augs = augs
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args)
+        return img
+
+
+if __name__ == "__main__":
+    a = RandomAugment()
+    img = np.random.randn(32, 32, 3)
+    a(img)

PreTrain_MeDSLIP/models/model_MeDSLIP.py

@@ -0,0 +1,530 @@
+# modified from https://github.com/tensorflow/models/blob/master/research/slim/nets/s3dg.py
+from sklearn.metrics import log_loss
+import torch.nn as nn
+import torch
+import math
+import numpy as np
+from torch.nn.utils.rnn import pad_sequence
+import torch.nn.functional as F
+from .transformer import *
+import torchvision.models as models
+from einops import rearrange
+from transformers import AutoModel
+
+"""
+args.N
+args.d_model
+args.res_base_model
+args.H
+args.num_queries
+args.dropout
+args.attribute_set_size
+"""
+
+
+class MeDSLIP(nn.Module):
+    def __init__(
+        self, config, anatomy_book, pathology_book, mode="train",
+    ):
+        super(MeDSLIP, self).__init__()
+        self.mode = mode
+        self.d_model = config["d_model"]
+        # """ book embedding"""
+        with torch.no_grad():
+            bert_model = self._get_bert_basemodel(
+                config["text_encoder"], freeze_layers=None
+            ).to(anatomy_book["input_ids"].device)
+            self.anatomy_book = bert_model(
+                input_ids=anatomy_book["input_ids"],
+                attention_mask=anatomy_book["attention_mask"],
+            )  # (**encoded_inputs)
+            self.anatomy_book = self.anatomy_book.last_hidden_state[:, 0, :]
+            self.pathology_book = bert_model(
+                input_ids=pathology_book["input_ids"],
+                attention_mask=pathology_book["attention_mask"],
+            )  # (**encoded_inputs)
+            self.pathology_book = self.pathology_book.last_hidden_state[:, 0, :]
+        self.pathology_embedding_layer = nn.Linear(768, 256)
+        self.cl_fc_pathology = nn.Linear(256, 768)
+
+        self.pathology_name = [
+            "normal",
+            "clear",
+            "sharp",
+            "sharply",
+            "unremarkable",
+            "intact",
+            "stable",
+            "free",
+            "effusion",
+            "opacity",
+            "pneumothorax",
+            "edema",
+            "atelectasis",
+            "tube",
+            "consolidation",
+            "process",
+            "abnormality",
+            "enlarge",
+            "tip",
+            "low",
+            "pneumonia",
+            "line",
+            "congestion",
+            "catheter",
+            "cardiomegaly",
+            "fracture",
+            "air",
+            "tortuous",
+            "lead",
+            "pathology",
+            "calcification",
+            "prominence",
+            "device",
+            "engorgement",
+            "picc",
+            "clip",
+            "elevation",
+            "expand",
+            "nodule",
+            "wire",
+            "fluid",
+            "degenerative",
+            "pacemaker",
+            "thicken",
+            "marking",
+            "scar",
+            "hyperinflate",
+            "blunt",
+            "loss",
+            "widen",
+            "coll_eapse",
+            "density",
+            "emphysema",
+            "aerate",
+            "mass",
+            "crowd",
+            "infiltrate",
+            "obscure",
+            "deformity",
+            "hernia",
+            "drainage",
+            "distention",
+            "shift",
+            "stent",
+            "pressure",
+            "lesion",
+            "finding",
+            "borderline",
+            "hardware",
+            "dilation",
+            "chf",
+            "redistribution",
+            "aspiration",
+            "tail_abnorm_obs",
+            "excluded_obs",
+        ]
+
+        self.excluded_pathology = [
+            "pneumonia",
+            "infiltrate",
+            "mass",
+            "nodule",
+            "emphysema",
+            "fibrosis",
+            "thicken",
+            "hernia",
+        ]
+
+        self.keep_class_dim_pathology = [
+            self.pathology_name.index(i)
+            for i in self.pathology_name
+            if i not in self.excluded_pathology
+        ]
+        """ visual backbone"""
+        self.resnet_dict = {
+            "resnet18": models.resnet18(pretrained=False),
+            "resnet50": models.resnet50(pretrained=False),
+        }
+        resnet = self._get_res_basemodel(config["res_base_model"])
+        num_ftrs = int(resnet.fc.in_features / 2)
+        self.res_features = nn.Sequential(*list(resnet.children())[:-3])
+
+        self.res_l1_pathology = nn.Linear(num_ftrs, num_ftrs)
+        self.res_l2_pathology = nn.Linear(num_ftrs, self.d_model)
+
+        self.cl_fc_anatomy = nn.Linear(256, 768)
+        self.res_l1_anatomy = nn.Linear(num_ftrs, num_ftrs)
+        self.res_l2_anatomy = nn.Linear(num_ftrs, self.d_model)
+
+        self.mask_generator = nn.Linear(num_ftrs, num_ftrs)
+
+        ###################################
+        """ Query Decoder"""
+        ###################################
+
+        self.H = config["H"]
+        decoder_layer = TransformerDecoderLayer(
+            self.d_model, config["H"], 1024, 0.1, "relu", normalize_before=True
+        )
+        decoder_norm = nn.LayerNorm(self.d_model)
+        self.decoder_anatomy = TransformerDecoder(
+            decoder_layer, config["N"], decoder_norm, return_intermediate=False
+        )
+        self.decoder_pathology = TransformerDecoder(
+            decoder_layer, config["N"], decoder_norm, return_intermediate=False
+        )
+
+        # Learnable Queries
+        self.dropout_feas_anatomy = nn.Dropout(config["dropout"])
+        self.dropout_feas_pathology = nn.Dropout(config["dropout"])
+
+        # Attribute classifier
+        self.classifier_anatomy = nn.Linear(self.d_model, config["attribute_set_size"])
+        self.classifier_pathology = nn.Linear(
+            self.d_model, config["attribute_set_size"]
+        )
+
+        self.apply(self._init_weights)
+
+    def _get_res_basemodel(self, res_model_name):
+        try:
+            res_model = self.resnet_dict[res_model_name]
+            print("Image feature extractor:", res_model_name)
+            return res_model
+        except:
+            raise (
+                "Invalid model name. Check the config file and pass one of: resnet18 or resnet50"
+            )
+
+    def _get_bert_basemodel(self, bert_model_name, freeze_layers):
+        try:
+            model = AutoModel.from_pretrained(bert_model_name)
+            print("text feature extractor:", bert_model_name)
+        except:
+            raise (
+                "Invalid model name. Check the config file and pass a BERT model from transformers lybrary"
+            )
+
+        if freeze_layers is not None:
+            for layer_idx in freeze_layers:
+                for param in list(model.encoder.layer[layer_idx].parameters()):
+                    param.requires_grad = False
+        return model
+
+    def image_encoder(self, xis):
+        # patch features
+        """
+        16 torch.Size([16, 1024, 14, 14])
+        torch.Size([16, 196, 1024])
+        torch.Size([3136, 1024])
+        torch.Size([16, 196, 256])
+        """
+        batch_size = xis.shape[0]
+        res_fea = self.res_features(xis)  # batch_size,feature_size,patch_num,patch_num
+        res_fea = rearrange(res_fea, "b d n1 n2 -> b (n1 n2) d")
+        x = rearrange(res_fea, "b n d -> (b n) d")
+
+        mask = self.mask_generator(x)
+        x_pathology = mask * x
+        x_anatomy = (1 - mask) * x
+
+        x_pathology = self.res_l1_pathology(x_pathology)
+        x_anatomy = self.res_l1_anatomy(x_anatomy)
+        x_pathology = F.relu(x_pathology)
+        x_anatomy = F.relu(x_anatomy)
+
+        x_pathology = self.res_l2_pathology(x_pathology)
+        x_anatomy = self.res_l2_anatomy(x_anatomy)
+
+        out_emb_pathology = rearrange(x_pathology, "(b n) d -> b n d", b=batch_size)
+        out_emb_anatomy = rearrange(x_anatomy, "(b n) d -> b n d", b=batch_size)
+        return out_emb_pathology, out_emb_anatomy
+
+    def forward(
+        self,
+        images,
+        labels_pathology=None,
+        labels_anatomy=None,
+        matrix=None,
+        sample_index_pathology=None,
+        sample_index_anatomy=None,
+        is_train=True,
+        text_gen=False,
+        no_cl=False,
+        exclude_class=False,
+    ):
+
+        B = images.shape[0]
+        device = images.device
+        """ Visual Backbone """
+        x_pathology, x_anatomy = self.image_encoder(images)  # batch_size,patch_num,dim
+
+        features_pathology = x_pathology.transpose(0, 1)  # patch_num b dim
+        features_anatomy = x_anatomy.transpose(0, 1)  # patch_num b dim
+
+        query_embed_pathology = self.pathology_embedding_layer(self.pathology_book)
+        query_embed_anatomy = self.pathology_embedding_layer(self.anatomy_book)
+        query_embed_pathology = query_embed_pathology.unsqueeze(1).repeat(1, B, 1)
+        query_embed_anatomy = query_embed_anatomy.unsqueeze(1).repeat(1, B, 1)
+
+        features_pathology, ws_pathology = self.decoder_pathology(
+            query_embed_pathology,
+            features_pathology,
+            memory_key_padding_mask=None,
+            pos=None,
+            query_pos=None,
+        )
+        features_anatomy, ws_anatomy = self.decoder_anatomy(
+            query_embed_anatomy,
+            features_anatomy,
+            memory_key_padding_mask=None,
+            pos=None,
+            query_pos=None,
+        )
+
+        ap_pathology = features_pathology
+        ap_anatomy = features_anatomy
+
+        ap_logits = torch.bmm(
+            ap_pathology.transpose(0, 1), ap_anatomy.transpose(0, 1).transpose(1, 2)
+        ).transpose(
+            1, 2
+        ) 
+        if text_gen:
+            output_logits = ap_logits
+        matrix_zero = matrix
+
+        masks = matrix_zero >= 0
+        ap_logits = ap_logits[masks]
+        matrix_zero = matrix_zero[masks]
+
+        loss_ap = F.binary_cross_entropy_with_logits(
+            ap_logits.float(), matrix_zero.float()
+        )
+
+        out_pathology = self.dropout_feas_pathology(features_pathology)
+        out_anatomy = self.dropout_feas_anatomy(features_anatomy)
+
+        if is_train == True and no_cl == False:
+
+            # get anatomytomy query
+            anatomytomy_query = torch.zeros(
+                [
+                    sample_index_pathology.shape[0],
+                    sample_index_pathology.shape[1],
+                    sample_index_pathology.shape[2],
+                    self.anatomy_book.shape[-1],
+                ]
+            ).to(
+                device
+            )
+            entity_query = torch.zeros(
+                [
+                    sample_index_anatomy.shape[0],
+                    sample_index_anatomy.shape[1],
+                    sample_index_anatomy.shape[2],
+                    self.pathology_book.shape[-1],
+                ]
+            ).to(device)
+
+            anatomytomy_query = self.anatomy_book[sample_index_pathology, :] * (
+                sample_index_pathology != -1
+            ).int().unsqueeze(-1).repeat(
+                1, 1, 1, 768
+            )  # batch, Q , position_num ,dim
+            entity_query = self.pathology_book[sample_index_anatomy, :] * (
+                sample_index_anatomy != -1
+            ).int().unsqueeze(-1).repeat(1, 1, 1, 768)
+
+            matrix_zero_pathology = matrix
+            matrix_zero_anatomy = matrix.transpose(1, 2)
+            matrix_zero_pathology[matrix_zero_pathology < 1] = 0
+            matrix_zero_anatomy[matrix_zero_anatomy < 1] = 0
+            matrix_zero_pathology = matrix_zero_pathology.unsqueeze(3).repeat(
+                1, 1, 1, anatomytomy_query.shape[-1]
+            )
+            matrix_zero_anatomy = matrix_zero_anatomy.unsqueeze(3).repeat(
+                1, 1, 1, entity_query.shape[-1]
+            )
+
+            anatomy_temp = self.anatomy_book
+            pathology_temp = self.pathology_book
+            anatomy_temp = anatomy_temp.unsqueeze(0).repeat(
+                anatomytomy_query.shape[0], 1, 1
+            )
+            pathology_temp = pathology_temp.unsqueeze(0).repeat(
+                entity_query.shape[0], 1, 1
+            )
+            anatomy_temp = anatomy_temp.unsqueeze(2).repeat(
+                1, 1, anatomytomy_query.shape[1], 1
+            )
+            pathology_temp = pathology_temp.unsqueeze(2).repeat(
+                1, 1, entity_query.shape[1], 1
+            )
+
+            posi_matrix_pathology = (matrix_zero_pathology * anatomy_temp).transpose(
+                1, 2
+            )
+            posi_matrix_anatomy = (matrix_zero_anatomy * pathology_temp).transpose(1, 2)
+
+            for i in range(anatomytomy_query.shape[0]):
+                for j in range(anatomytomy_query.shape[1]):
+                    if (posi_matrix_pathology[i, j] != 0).sum() > 0:
+                        num_posi = (
+                            torch.nonzero(posi_matrix_pathology[i, j], as_tuple=True)[0]
+                            .unique()
+                            .shape[0]
+                        )
+                        assert anatomytomy_query[i, j, 0, :].sum() == 0
+                        anatomytomy_query[i, j, 0, :] = (
+                            posi_matrix_pathology[i, j, :, :].sum(dim=0) / num_posi
+                        )
+
+            for i in range(entity_query.shape[0]):
+                for j in range(entity_query.shape[1]):
+                    if (posi_matrix_anatomy[i, j] != 0).sum() > 0:
+                        num_posi = (
+                            torch.nonzero(posi_matrix_anatomy[i, j], as_tuple=True)[0]
+                            .unique()
+                            .shape[0]
+                        )
+                        assert entity_query[i, j, 0, :].sum() == 0
+                        entity_query[i, j, 0, :] = (
+                            posi_matrix_anatomy[i, j, :, :].sum(dim=0) / num_posi
+                        )
+            # Got anatomytomy query
+
+            # [Q,B,A]
+            ll_pathology = out_pathology.transpose(0, 1)  # B Q A
+            ll_anatomy = out_anatomy.transpose(0, 1)  # B Q A
+
+            Q_pathology = ll_pathology.shape[1]
+            Q_anatomy = ll_anatomy.shape[1]
+
+            ll_pathology = ll_pathology.reshape(
+                ll_pathology.shape[0] * ll_pathology.shape[1], -1
+            )
+            ll_anatomy = ll_anatomy.reshape(
+                ll_anatomy.shape[0] * ll_anatomy.shape[1], -1
+            )
+
+            ll_pathology = self.cl_fc_pathology(ll_pathology)
+            ll_anatomy = self.cl_fc_anatomy(ll_anatomy)
+
+            ll_pathology = ll_pathology.unsqueeze(dim=-1)
+            ll_anatomy = ll_anatomy.unsqueeze(dim=-1)
+
+            anatomytomy_query = anatomytomy_query.reshape(B * Q_pathology, 8, 768)
+            entity_query = entity_query.reshape(B * Q_anatomy, 8, 768)
+
+            ll_pathology = torch.bmm(
+                anatomytomy_query, ll_pathology
+            ).squeeze()  # B Q position_num
+            ll_anatomy = torch.bmm(
+                entity_query, ll_anatomy
+            ).squeeze()  # B Q position_num
+
+            cl_labels_pathology = torch.zeros((ll_pathology.shape[0])).to(device)
+            cl_labels_anatomy = torch.zeros((ll_anatomy.shape[0])).to(device)
+
+            if exclude_class == True:
+                cl_labels_pathology = cl_labels_pathology.reshape(B, Q_pathology)
+                cl_labels_anatomy = cl_labels_anatomy.reshape(B, Q_anatomy)
+
+                cl_labels_pathology = cl_labels_pathology[
+                    :, self.keep_class_dim_pathology
+                ]
+                cl_labels_anatomy = cl_labels_anatomy[:, self.keep_class_dim_pathology]
+
+                cl_labels_pathology = cl_labels_pathology.reshape(-1)
+                cl_labels_anatomy = cl_labels_anatomy.reshape(-1)
+
+                ll_pathology = ll_pathology.reshape(B, Q_pathology, -1)
+                ll_anatomy = ll_anatomy.reshape(B, Q_anatomy, -1)
+
+                ll_pathology = ll_pathology[:, self.keep_class_dim_pathology, :]
+                ll_pathology = ll_pathology.reshape(
+                    B * (len(self.keep_class_dim_pathology)), -1
+                )
+                ll_anatomy = ll_anatomy.reshape(B * Q_anatomy, -1)
+
+        x_pathology = self.classifier_pathology(out_pathology).transpose(0, 1)
+        x_anatomy = self.classifier_anatomy(out_anatomy).transpose(
+            0, 1
+        )  # B query Atributes
+
+        if exclude_class == True:
+            labels_pathology = labels_pathology[:, self.keep_class_dim_pathology]
+            x_pathology = x_pathology[:, self.keep_class_dim_pathology, :]
+
+        labels_pathology = labels_pathology.reshape(-1, 1)
+        labels_anatomy = labels_anatomy.reshape(-1, 1)
+        logits_pathology = x_pathology.reshape(-1, x_pathology.shape[-1])
+        logits_anatomy = x_anatomy.reshape(-1, x_anatomy.shape[-1])
+        Mask_pathology = ((labels_pathology != -1) & (labels_pathology != 2)).squeeze()
+        Mask_anatomy = ((labels_anatomy != -1) & (labels_anatomy != 2)).squeeze()
+
+        cl_mask_pathology = (labels_pathology == 1).squeeze()
+        cl_mask_anatomy = (labels_anatomy == 1).squeeze()
+        if is_train == True:
+            labels_pathology = labels_pathology[Mask_pathology].long()
+            labels_anatomy = labels_anatomy[Mask_anatomy].long()
+            logits_pathology = logits_pathology[Mask_pathology]
+            logits_anatomy = logits_anatomy[Mask_anatomy]
+            loss_ce_pathology = F.cross_entropy(
+                logits_pathology, labels_pathology[:, 0]
+            )
+            loss_ce_anatomy = F.cross_entropy(logits_anatomy, labels_anatomy[:, 0])
+            if no_cl == False:
+                cl_labels_pathology = cl_labels_pathology[cl_mask_pathology].long()
+                cl_labels_anatomy = cl_labels_anatomy[cl_mask_anatomy].long()
+                ll_pathology = ll_pathology[cl_mask_pathology]
+                ll_anatomy = ll_anatomy[cl_mask_anatomy]
+                loss_cl_pathology = F.cross_entropy(ll_pathology, cl_labels_pathology)
+                loss_cl_anatomy = F.cross_entropy(ll_anatomy, cl_labels_anatomy)
+                loss_ce = loss_ce_pathology + loss_ce_anatomy
+                loss_cl = loss_cl_pathology + loss_cl_anatomy
+                loss = loss_ce + loss_cl + loss_ap
+            else:
+                loss_cl = torch.tensor(0)
+                loss = loss_ce_pathology + loss_ce_anatomy + loss_ap
+        else:
+            loss = 0
+        if is_train == True:
+            if text_gen:
+                return (
+                    loss,
+                    x_pathology,
+                    ws_pathology,
+                    x_anatomy,
+                    ws_anatomy,
+                    output_logits,
+                )
+            else:
+                return (
+                    loss,
+                    loss_ce_pathology,
+                    loss_cl_pathology,
+                    loss_ce_anatomy,
+                    loss_cl_anatomy,
+                    loss_ap,
+                )
+        else:
+            return loss, x_pathology, ws_pathology, x_anatomy, ws_anatomy
+
+    @staticmethod
+    def _init_weights(module):
+        r"""Initialize weights like BERT - N(0.0, 0.02), bias = 0."""
+
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+        elif isinstance(module, nn.MultiheadAttention):
+            module.in_proj_weight.data.normal_(mean=0.0, std=0.02)
+            module.out_proj.weight.data.normal_(mean=0.0, std=0.02)
+
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()

PreTrain_MeDSLIP/models/tokenization_bert.py

@@ -0,0 +1,578 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for Bert."""
+
+
+import collections
+import os
+import unicodedata
+from typing import List, Optional, Tuple
+
+from transformers.tokenization_utils import (
+    PreTrainedTokenizer,
+    _is_control,
+    _is_punctuation,
+    _is_whitespace,
+)
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt",
+        "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt",
+        "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/vocab.txt",
+        "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/vocab.txt",
+        "bert-base-multilingual-uncased": "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/vocab.txt",
+        "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/vocab.txt",
+        "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt",
+        "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/vocab.txt",
+        "bert-large-uncased-whole-word-masking": "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/vocab.txt",
+        "bert-large-cased-whole-word-masking": "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/vocab.txt",
+        "bert-large-uncased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt",
+        "bert-large-cased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt",
+        "bert-base-cased-finetuned-mrpc": "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/vocab.txt",
+        "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/vocab.txt",
+        "bert-base-german-dbmdz-uncased": "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/vocab.txt",
+        "TurkuNLP/bert-base-finnish-cased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/vocab.txt",
+        "TurkuNLP/bert-base-finnish-uncased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/vocab.txt",
+        "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/vocab.txt",
+    }
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "bert-base-uncased": 512,
+    "bert-large-uncased": 512,
+    "bert-base-cased": 512,
+    "bert-large-cased": 512,
+    "bert-base-multilingual-uncased": 512,
+    "bert-base-multilingual-cased": 512,
+    "bert-base-chinese": 512,
+    "bert-base-german-cased": 512,
+    "bert-large-uncased-whole-word-masking": 512,
+    "bert-large-cased-whole-word-masking": 512,
+    "bert-large-uncased-whole-word-masking-finetuned-squad": 512,
+    "bert-large-cased-whole-word-masking-finetuned-squad": 512,
+    "bert-base-cased-finetuned-mrpc": 512,
+    "bert-base-german-dbmdz-cased": 512,
+    "bert-base-german-dbmdz-uncased": 512,
+    "TurkuNLP/bert-base-finnish-cased-v1": 512,
+    "TurkuNLP/bert-base-finnish-uncased-v1": 512,
+    "wietsedv/bert-base-dutch-cased": 512,
+}
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "bert-base-uncased": {"do_lower_case": True},
+    "bert-large-uncased": {"do_lower_case": True},
+    "bert-base-cased": {"do_lower_case": False},
+    "bert-large-cased": {"do_lower_case": False},
+    "bert-base-multilingual-uncased": {"do_lower_case": True},
+    "bert-base-multilingual-cased": {"do_lower_case": False},
+    "bert-base-chinese": {"do_lower_case": False},
+    "bert-base-german-cased": {"do_lower_case": False},
+    "bert-large-uncased-whole-word-masking": {"do_lower_case": True},
+    "bert-large-cased-whole-word-masking": {"do_lower_case": False},
+    "bert-large-uncased-whole-word-masking-finetuned-squad": {"do_lower_case": True},
+    "bert-large-cased-whole-word-masking-finetuned-squad": {"do_lower_case": False},
+    "bert-base-cased-finetuned-mrpc": {"do_lower_case": False},
+    "bert-base-german-dbmdz-cased": {"do_lower_case": False},
+    "bert-base-german-dbmdz-uncased": {"do_lower_case": True},
+    "TurkuNLP/bert-base-finnish-cased-v1": {"do_lower_case": False},
+    "TurkuNLP/bert-base-finnish-uncased-v1": {"do_lower_case": True},
+    "wietsedv/bert-base-dutch-cased": {"do_lower_case": False},
+}
+
+
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+class BertTokenizer(PreTrainedTokenizer):
+    r"""
+    Construct a BERT tokenizer. Based on WordPiece.
+    This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the main methods.
+    Users should refer to this superclass for more information regarding those methods.
+    Args:
+        vocab_file (:obj:`str`):
+            File containing the vocabulary.
+        do_lower_case (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to lowercase the input when tokenizing.
+        do_basic_tokenize (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to do basic tokenization before WordPiece.
+        never_split (:obj:`Iterable`, `optional`):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            :obj:`do_basic_tokenize=True`
+        unk_token (:obj:`str`, `optional`, defaults to :obj:`"[UNK]"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        sep_token (:obj:`str`, `optional`, defaults to :obj:`"[SEP]"`):
+            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+            sequence classification or for a text and a question for question answering. It is also used as the last
+            token of a sequence built with special tokens.
+        pad_token (:obj:`str`, `optional`, defaults to :obj:`"[PAD]"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        cls_token (:obj:`str`, `optional`, defaults to :obj:`"[CLS]"`):
+            The classifier token which is used when doing sequence classification (classification of the whole sequence
+            instead of per-token classification). It is the first token of the sequence when built with special tokens.
+        mask_token (:obj:`str`, `optional`, defaults to :obj:`"[MASK]"`):
+            The token used for masking values. This is the token used when training this model with masked language
+            modeling. This is the token which the model will try to predict.
+        tokenize_chinese_chars (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to tokenize Chinese characters.
+            This should likely be deactivated for Japanese (see this `issue
+            <https://github.com/huggingface/transformers/issues/328>`__).
+        strip_accents: (:obj:`bool`, `optional`):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for :obj:`lowercase` (as in the original BERT).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(
+        self,
+        vocab_file,
+        do_lower_case=True,
+        do_basic_tokenize=True,
+        never_split=None,
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+        tokenize_chinese_chars=True,
+        strip_accents=None,
+        **kwargs
+    ):
+        super().__init__(
+            do_lower_case=do_lower_case,
+            do_basic_tokenize=do_basic_tokenize,
+            never_split=never_split,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            tokenize_chinese_chars=tokenize_chinese_chars,
+            strip_accents=strip_accents,
+            **kwargs,
+        )
+
+        if not os.path.isfile(vocab_file):
+            raise ValueError(
+                "Can't find a vocabulary file at path '{}'. To load the vocabulary from a Google pretrained "
+                "model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    vocab_file
+                )
+            )
+        self.vocab = load_vocab(vocab_file)
+        self.ids_to_tokens = collections.OrderedDict(
+            [(ids, tok) for tok, ids in self.vocab.items()]
+        )
+        self.do_basic_tokenize = do_basic_tokenize
+        if do_basic_tokenize:
+            self.basic_tokenizer = BasicTokenizer(
+                do_lower_case=do_lower_case,
+                never_split=never_split,
+                tokenize_chinese_chars=tokenize_chinese_chars,
+                strip_accents=strip_accents,
+            )
+        self.wordpiece_tokenizer = WordpieceTokenizer(
+            vocab=self.vocab, unk_token=self.unk_token
+        )
+
+    @property
+    def do_lower_case(self):
+        return self.basic_tokenizer.do_lower_case
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return dict(self.vocab, **self.added_tokens_encoder)
+
+    def _tokenize(self, text):
+        split_tokens = []
+        if self.do_basic_tokenize:
+            for token in self.basic_tokenizer.tokenize(
+                text, never_split=self.all_special_tokens
+            ):
+
+                # If the token is part of the never_split set
+                if token in self.basic_tokenizer.never_split:
+                    split_tokens.append(token)
+                else:
+                    split_tokens += self.wordpiece_tokenizer.tokenize(token)
+        else:
+            split_tokens = self.wordpiece_tokenizer.tokenize(text)
+        return split_tokens
+
+    def _convert_token_to_id(self, token):
+        """ Converts a token (str) in an id using the vocab. """
+        return self.vocab.get(token, self.vocab.get(self.unk_token))
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens):
+        """ Converts a sequence of tokens (string) in a single string. """
+        out_string = " ".join(tokens).replace(" ##", "").strip()
+        return out_string
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+        - single sequence: ``[CLS] X ``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        if token_ids_1 is None:
+            return [self.cls_token_id] + token_ids_0
+        cls = [self.cls_token_id]
+        sep = [self.sep_token_id]
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer ``prepare_for_model`` method.
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+        Returns:
+            :obj:`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+
+        if already_has_special_tokens:
+            if token_ids_1 is not None:
+                raise ValueError(
+                    "You should not supply a second sequence if the provided sequence of "
+                    "ids is already formatted with special tokens for the model."
+                )
+            return list(
+                map(
+                    lambda x: 1 if x in [self.sep_token_id, self.cls_token_id] else 0,
+                    token_ids_0,
+                )
+            )
+
+        if token_ids_1 is not None:
+            return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+        return [1] + ([0] * len(token_ids_0)) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
+        pair mask has the following format:
+        ::
+            0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+            | first sequence    | second sequence |
+        If :obj:`token_ids_1` is :obj:`None`, this method only returns the first portion of the mask (0s).
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            :obj:`List[int]`: List of `token type IDs <../glossary.html#token-type-ids>`_ according to the given
+            sequence(s).
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        index = 0
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "")
+                + VOCAB_FILES_NAMES["vocab_file"],
+            )
+        else:
+            vocab_file = (
+                filename_prefix + "-" if filename_prefix else ""
+            ) + save_directory
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        "Saving vocabulary to {}: vocabulary indices are not consecutive."
+                        " Please check that the vocabulary is not corrupted!".format(
+                            vocab_file
+                        )
+                    )
+                    index = token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)
+
+
+class BasicTokenizer(object):
+    """
+    Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
+    Args:
+        do_lower_case (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to lowercase the input when tokenizing.
+        never_split (:obj:`Iterable`, `optional`):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            :obj:`do_basic_tokenize=True`
+        tokenize_chinese_chars (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to tokenize Chinese characters.
+            This should likely be deactivated for Japanese (see this `issue
+            <https://github.com/huggingface/transformers/issues/328>`__).
+        strip_accents: (:obj:`bool`, `optional`):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for :obj:`lowercase` (as in the original BERT).
+    """
+
+    def __init__(
+        self,
+        do_lower_case=True,
+        never_split=None,
+        tokenize_chinese_chars=True,
+        strip_accents=None,
+    ):
+        if never_split is None:
+            never_split = []
+        self.do_lower_case = do_lower_case
+        self.never_split = set(never_split)
+        self.tokenize_chinese_chars = tokenize_chinese_chars
+        self.strip_accents = strip_accents
+
+    def tokenize(self, text, never_split=None):
+        """
+        Basic Tokenization of a piece of text. Split on "white spaces" only, for sub-word tokenization, see
+        WordPieceTokenizer.
+        Args:
+            **never_split**: (`optional`) list of str
+                Kept for backward compatibility purposes. Now implemented directly at the base class level (see
+                :func:`PreTrainedTokenizer.tokenize`) List of token not to split.
+        """
+        # union() returns a new set by concatenating the two sets.
+        never_split = (
+            self.never_split.union(set(never_split))
+            if never_split
+            else self.never_split
+        )
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        if self.tokenize_chinese_chars:
+            text = self._tokenize_chinese_chars(text)
+        orig_tokens = whitespace_tokenize(text)
+        split_tokens = []
+        for token in orig_tokens:
+            if token not in never_split:
+                if self.do_lower_case:
+                    token = token.lower()
+                    if self.strip_accents is not False:
+                        token = self._run_strip_accents(token)
+                elif self.strip_accents:
+                    token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text, never_split=None):
+        """Splits punctuation on a piece of text."""
+        if never_split is not None and text in never_split:
+            return [text]
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if (
+            (cp >= 0x4E00 and cp <= 0x9FFF)
+            or (cp >= 0x3400 and cp <= 0x4DBF)  #
+            or (cp >= 0x20000 and cp <= 0x2A6DF)  #
+            or (cp >= 0x2A700 and cp <= 0x2B73F)  #
+            or (cp >= 0x2B740 and cp <= 0x2B81F)  #
+            or (cp >= 0x2B820 and cp <= 0x2CEAF)  #
+            or (cp >= 0xF900 and cp <= 0xFAFF)
+            or (cp >= 0x2F800 and cp <= 0x2FA1F)  #
+        ):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xFFFD or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenization."""
+
+    def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """
+        Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+        tokenization using the given vocabulary.
+        For example, :obj:`input = "unaffable"` wil return as output :obj:`["un", "##aff", "##able"]`.
+        Args:
+          text: A single token or whitespace separated tokens. This should have
+            already been passed through `BasicTokenizer`.
+        Returns:
+          A list of wordpiece tokens.
+        """
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens

PreTrain_MeDSLIP/models/transformer.py

@@ -0,0 +1,210 @@
+"""
+Code modified from DETR tranformer:
+https://github.com/facebookresearch/detr
+Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+
+import copy
+from typing import Optional, List
+import pickle as cp
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        output = tgt
+        T, B, C = memory.shape
+        intermediate = []
+        atten_layers = []
+        for n, layer in enumerate(self.layers):
+
+            residual = True
+            output, ws = layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                pos=pos,
+                query_pos=query_pos,
+                residual=residual,
+            )
+            atten_layers.append(ws)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+        return output, atten_layers
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation="relu",
+        normalize_before=False,
+    ):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        residual=True,
+    ):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2, ws = self.self_attn(
+            q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )
+        tgt = self.norm1(tgt)
+        tgt2, ws = self.multihead_attn(
+            query=self.with_pos_embed(tgt, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )
+
+        # attn_weights [B,NUM_Q,T]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt, ws
+
+    def forward_pre(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2, ws = self.self_attn(
+            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2, attn_weights = self.multihead_attn(
+            query=self.with_pos_embed(tgt2, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt, attn_weights
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        residual=True,
+    ):
+        if self.normalize_before:
+            return self.forward_pre(
+                tgt,
+                memory,
+                tgt_mask,
+                memory_mask,
+                tgt_key_padding_mask,
+                memory_key_padding_mask,
+                pos,
+                query_pos,
+            )
+        return self.forward_post(
+            tgt,
+            memory,
+            tgt_mask,
+            memory_mask,
+            tgt_key_padding_mask,
+            memory_key_padding_mask,
+            pos,
+            query_pos,
+            residual,
+        )
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")

PreTrain_MeDSLIP/optim/__init__.py

@@ -0,0 +1,13 @@
+from .adamp import AdamP
+from .adamw import AdamW
+from .adafactor import Adafactor
+from .adahessian import Adahessian
+from .lookahead import Lookahead
+from .nadam import Nadam
+from .novograd import NovoGrad
+from .nvnovograd import NvNovoGrad
+from .radam import RAdam
+from .rmsprop_tf import RMSpropTF
+from .sgdp import SGDP
+
+from .optim_factory import create_optimizer

PreTrain_MeDSLIP/optim/adafactor.py

@@ -0,0 +1,206 @@
+""" Adafactor Optimizer
+
+Lifted from https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+Original header/copyright below.
+
+"""
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+import math
+
+
+class Adafactor(torch.optim.Optimizer):
+    """Implements Adafactor algorithm.
+    This implementation is based on: `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`
+    (see https://arxiv.org/abs/1804.04235)
+
+    Note that this optimizer internally adjusts the learning rate depending on the
+    *scale_parameter*, *relative_step* and *warmup_init* options.
+
+    To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining parameter groups
+        lr (float, optional): external learning rate (default: None)
+        eps (tuple[float, float]): regularization constants for square gradient
+            and parameter scale respectively (default: (1e-30, 1e-3))
+        clip_threshold (float): threshold of root mean square of final gradient update (default: 1.0)
+        decay_rate (float): coefficient used to compute running averages of square gradient (default: -0.8)
+        beta1 (float): coefficient used for computing running averages of gradient (default: None)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        scale_parameter (bool): if True, learning rate is scaled by root mean square of parameter (default: True)
+        relative_step (bool): if True, time-dependent learning rate is computed
+            instead of external learning rate (default: True)
+        warmup_init (bool): time-dependent learning rate computation depends on
+            whether warm-up initialization is being used (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=1e-30,
+        eps_scale=1e-3,
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        betas=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        warmup_init=False,
+    ):
+        relative_step = lr is None
+        if warmup_init and not relative_step:
+            raise ValueError("warmup_init requires relative_step=True")
+
+        beta1 = (
+            None if betas is None else betas[0]
+        )  # make it compat with standard betas arg
+        defaults = dict(
+            lr=lr,
+            eps=eps,
+            eps_scale=eps_scale,
+            clip_threshold=clip_threshold,
+            decay_rate=decay_rate,
+            beta1=beta1,
+            weight_decay=weight_decay,
+            scale_parameter=scale_parameter,
+            relative_step=relative_step,
+            warmup_init=warmup_init,
+        )
+        super(Adafactor, self).__init__(params, defaults)
+
+    @staticmethod
+    def _get_lr(param_group, param_state):
+        if param_group["relative_step"]:
+            min_step = (
+                1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2
+            )
+            lr_t = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+            param_scale = 1.0
+            if param_group["scale_parameter"]:
+                param_scale = max(param_group["eps_scale"], param_state["RMS"])
+            param_group["lr"] = lr_t * param_scale
+        return param_group["lr"]
+
+    @staticmethod
+    def _get_options(param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group["beta1"] is not None
+        return factored, use_first_moment
+
+    @staticmethod
+    def _rms(tensor):
+        return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+    def _approx_sq_grad(self, exp_avg_sq_row, exp_avg_sq_col):
+        r_factor = (
+            (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True))
+            .rsqrt_()
+            .unsqueeze(-1)
+        )
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError("Adafactor does not support sparse gradients.")
+
+                state = self.state[p]
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state["exp_avg_sq_col"] = torch.zeros(
+                            grad_shape[:-2] + grad_shape[-1:]
+                        ).to(grad)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                    state["RMS"] = 0
+                else:
+                    if use_first_moment:
+                        state["exp_avg"] = state["exp_avg"].to(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+                    else:
+                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+                p_data_fp32 = p.data
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state["step"] += 1
+                state["RMS"] = self._rms(p_data_fp32)
+                lr_t = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+                update = grad ** 2 + group["eps"]
+                if factored:
+                    exp_avg_sq_row = state["exp_avg_sq_row"]
+                    exp_avg_sq_col = state["exp_avg_sq_col"]
+
+                    exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
+                    exp_avg_sq_col.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-2))
+                    # exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)  # pytorch 1.6+
+                    # exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state["exp_avg_sq"]
+
+                    exp_avg_sq.mul_(beta2t).add_(1.0 - beta2t, update)
+                    # exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)  # pytorch 1.6+
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_(
+                    (self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)
+                )
+                update.mul_(lr_t)
+
+                if use_first_moment:
+                    exp_avg = state["exp_avg"]
+                    exp_avg.mul_(group["beta1"]).add_(1 - group["beta1"], update)
+                    # exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])  # pytorch 1.6+
+                    update = exp_avg
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * lr_t, p_data_fp32)
+                    # p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr_t)  # pytorch 1.6+
+
+                p_data_fp32.add_(-update)
+
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p.data.copy_(p_data_fp32)
+
+        return loss

PreTrain_MeDSLIP/optim/adahessian.py

@@ -0,0 +1,207 @@
+""" AdaHessian Optimizer
+
+Lifted from https://github.com/davda54/ada-hessian/blob/master/ada_hessian.py
+Originally licensed MIT, Copyright 2020, David Samuel
+"""
+import torch
+
+
+class Adahessian(torch.optim.Optimizer):
+    """
+    Implements the AdaHessian algorithm from "ADAHESSIAN: An Adaptive Second OrderOptimizer for Machine Learning"
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining parameter groups
+        lr (float, optional): learning rate (default: 0.1)
+        betas ((float, float), optional): coefficients used for computing running averages of gradient and the
+            squared hessian trace (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.0)
+        hessian_power (float, optional): exponent of the hessian trace (default: 1.0)
+        update_each (int, optional): compute the hessian trace approximation only after *this* number of steps
+            (to save time) (default: 1)
+        n_samples (int, optional): how many times to sample `z` for the approximation of the hessian trace (default: 1)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=0.1,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0.0,
+        hessian_power=1.0,
+        update_each=1,
+        n_samples=1,
+        avg_conv_kernel=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= hessian_power <= 1.0:
+            raise ValueError(f"Invalid Hessian power value: {hessian_power}")
+
+        self.n_samples = n_samples
+        self.update_each = update_each
+        self.avg_conv_kernel = avg_conv_kernel
+
+        # use a separate generator that deterministically generates the same `z`s across all GPUs in case of distributed training
+        self.seed = 2147483647
+        self.generator = torch.Generator().manual_seed(self.seed)
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            hessian_power=hessian_power,
+        )
+        super(Adahessian, self).__init__(params, defaults)
+
+        for p in self.get_params():
+            p.hess = 0.0
+            self.state[p]["hessian step"] = 0
+
+    @property
+    def is_second_order(self):
+        return True
+
+    def get_params(self):
+        """
+        Gets all parameters in all param_groups with gradients
+        """
+
+        return (
+            p for group in self.param_groups for p in group["params"] if p.requires_grad
+        )
+
+    def zero_hessian(self):
+        """
+        Zeros out the accumalated hessian traces.
+        """
+
+        for p in self.get_params():
+            if (
+                not isinstance(p.hess, float)
+                and self.state[p]["hessian step"] % self.update_each == 0
+            ):
+                p.hess.zero_()
+
+    @torch.no_grad()
+    def set_hessian(self):
+        """
+        Computes the Hutchinson approximation of the hessian trace and accumulates it for each trainable parameter.
+        """
+
+        params = []
+        for p in filter(lambda p: p.grad is not None, self.get_params()):
+            if (
+                self.state[p]["hessian step"] % self.update_each == 0
+            ):  # compute the trace only each `update_each` step
+                params.append(p)
+            self.state[p]["hessian step"] += 1
+
+        if len(params) == 0:
+            return
+
+        if (
+            self.generator.device != params[0].device
+        ):  # hackish way of casting the generator to the right device
+            self.generator = torch.Generator(params[0].device).manual_seed(self.seed)
+
+        grads = [p.grad for p in params]
+
+        for i in range(self.n_samples):
+            # Rademacher distribution {-1.0, 1.0}
+            zs = [
+                torch.randint(0, 2, p.size(), generator=self.generator, device=p.device)
+                * 2.0
+                - 1.0
+                for p in params
+            ]
+            h_zs = torch.autograd.grad(
+                grads,
+                params,
+                grad_outputs=zs,
+                only_inputs=True,
+                retain_graph=i < self.n_samples - 1,
+            )
+            for h_z, z, p in zip(h_zs, zs, params):
+                p.hess += (
+                    h_z * z / self.n_samples
+                )  # approximate the expected values of z*(H@z)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """
+        Performs a single optimization step.
+        Arguments:
+            closure (callable, optional) -- a closure that reevaluates the model and returns the loss (default: None)
+        """
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        self.zero_hessian()
+        self.set_hessian()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None or p.hess is None:
+                    continue
+
+                if self.avg_conv_kernel and p.dim() == 4:
+                    p.hess = (
+                        torch.abs(p.hess)
+                        .mean(dim=[2, 3], keepdim=True)
+                        .expand_as(p.hess)
+                        .clone()
+                    )
+
+                # Perform correct stepweight decay as in AdamW
+                p.mul_(1 - group["lr"] * group["weight_decay"])
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 1:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p)
+                    # Exponential moving average of Hessian diagonal square values
+                    state["exp_hessian_diag_sq"] = torch.zeros_like(p)
+
+                exp_avg, exp_hessian_diag_sq = (
+                    state["exp_avg"],
+                    state["exp_hessian_diag_sq"],
+                )
+                beta1, beta2 = group["betas"]
+                state["step"] += 1
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(p.grad, alpha=1 - beta1)
+                exp_hessian_diag_sq.mul_(beta2).addcmul_(
+                    p.hess, p.hess, value=1 - beta2
+                )
+
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                k = group["hessian_power"]
+                denom = (
+                    (exp_hessian_diag_sq / bias_correction2)
+                    .pow_(k / 2)
+                    .add_(group["eps"])
+                )
+
+                # make update
+                step_size = group["lr"] / bias_correction1
+                p.addcdiv_(exp_avg, denom, value=-step_size)
+
+        return loss

PreTrain_MeDSLIP/optim/adamp.py

@@ -0,0 +1,133 @@
+"""
+AdamP Optimizer Implementation copied from https://github.com/clovaai/AdamP/blob/master/adamp/adamp.py
+
+Paper: `Slowing Down the Weight Norm Increase in Momentum-based Optimizers` - https://arxiv.org/abs/2006.08217
+Code: https://github.com/clovaai/AdamP
+
+Copyright (c) 2020-present NAVER Corp.
+MIT license
+"""
+
+import torch
+import torch.nn as nn
+from torch.optim.optimizer import Optimizer, required
+import math
+
+
+class AdamP(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0,
+        delta=0.1,
+        wd_ratio=0.1,
+        nesterov=False,
+    ):
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            delta=delta,
+            wd_ratio=wd_ratio,
+            nesterov=nesterov,
+        )
+        super(AdamP, self).__init__(params, defaults)
+
+    def _channel_view(self, x):
+        return x.view(x.size(0), -1)
+
+    def _layer_view(self, x):
+        return x.view(1, -1)
+
+    def _cosine_similarity(self, x, y, eps, view_func):
+        x = view_func(x)
+        y = view_func(y)
+
+        x_norm = x.norm(dim=1).add_(eps)
+        y_norm = y.norm(dim=1).add_(eps)
+        dot = (x * y).sum(dim=1)
+
+        return dot.abs() / x_norm / y_norm
+
+    def _projection(self, p, grad, perturb, delta, wd_ratio, eps):
+        wd = 1
+        expand_size = [-1] + [1] * (len(p.shape) - 1)
+        for view_func in [self._channel_view, self._layer_view]:
+
+            cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func)
+
+            if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)):
+                p_n = p.data / view_func(p.data).norm(dim=1).view(expand_size).add_(eps)
+                perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view(expand_size)
+                wd = wd_ratio
+
+                return perturb, wd
+
+        return perturb, wd
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                beta1, beta2 = group["betas"]
+                nesterov = group["nesterov"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+
+                # Adam
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                    group["eps"]
+                )
+                step_size = group["lr"] / bias_correction1
+
+                if nesterov:
+                    perturb = (beta1 * exp_avg + (1 - beta1) * grad) / denom
+                else:
+                    perturb = exp_avg / denom
+
+                # Projection
+                wd_ratio = 1
+                if len(p.shape) > 1:
+                    perturb, wd_ratio = self._projection(
+                        p,
+                        grad,
+                        perturb,
+                        group["delta"],
+                        group["wd_ratio"],
+                        group["eps"],
+                    )
+
+                # Weight decay
+                if group["weight_decay"] > 0:
+                    p.data.mul_(1 - group["lr"] * group["weight_decay"] * wd_ratio)
+
+                # Step
+                p.data.add_(-step_size, perturb)
+
+        return loss

PreTrain_MeDSLIP/optim/adamw.py

@@ -0,0 +1,131 @@
+""" AdamW Optimizer
+Impl copied from PyTorch master
+"""
+import math
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class AdamW(Optimizer):
+    r"""Implements AdamW algorithm.
+
+    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay coefficient (default: 1e-2)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=1e-2,
+        amsgrad=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad
+        )
+        super(AdamW, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdamW, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Perform stepweight decay
+                p.data.mul_(1 - group["lr"] * group["weight_decay"])
+
+                # Perform optimization step
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "Adam does not support sparse gradients, please consider SparseAdam instead"
+                    )
+                amsgrad = group["amsgrad"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state["max_exp_avg_sq"] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                        group["eps"]
+                    )
+                else:
+                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                        group["eps"]
+                    )
+
+                step_size = group["lr"] / bias_correction1
+
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+        return loss

PreTrain_MeDSLIP/optim/lookahead.py

@@ -0,0 +1,96 @@
+""" Lookahead Optimizer Wrapper.
+Implementation modified from: https://github.com/alphadl/lookahead.pytorch
+Paper: `Lookahead Optimizer: k steps forward, 1 step back` - https://arxiv.org/abs/1907.08610
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+from torch.optim.optimizer import Optimizer
+from collections import defaultdict
+
+
+class Lookahead(Optimizer):
+    def __init__(self, base_optimizer, alpha=0.5, k=6):
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f"Invalid slow update rate: {alpha}")
+        if not 1 <= k:
+            raise ValueError(f"Invalid lookahead steps: {k}")
+        defaults = dict(lookahead_alpha=alpha, lookahead_k=k, lookahead_step=0)
+        self.base_optimizer = base_optimizer
+        self.param_groups = self.base_optimizer.param_groups
+        self.defaults = base_optimizer.defaults
+        self.defaults.update(defaults)
+        self.state = defaultdict(dict)
+        # manually add our defaults to the param groups
+        for name, default in defaults.items():
+            for group in self.param_groups:
+                group.setdefault(name, default)
+
+    def update_slow(self, group):
+        for fast_p in group["params"]:
+            if fast_p.grad is None:
+                continue
+            param_state = self.state[fast_p]
+            if "slow_buffer" not in param_state:
+                param_state["slow_buffer"] = torch.empty_like(fast_p.data)
+                param_state["slow_buffer"].copy_(fast_p.data)
+            slow = param_state["slow_buffer"]
+            slow.add_(group["lookahead_alpha"], fast_p.data - slow)
+            fast_p.data.copy_(slow)
+
+    def sync_lookahead(self):
+        for group in self.param_groups:
+            self.update_slow(group)
+
+    def step(self, closure=None):
+        # assert id(self.param_groups) == id(self.base_optimizer.param_groups)
+        loss = self.base_optimizer.step(closure)
+        for group in self.param_groups:
+            group["lookahead_step"] += 1
+            if group["lookahead_step"] % group["lookahead_k"] == 0:
+                self.update_slow(group)
+        return loss
+
+    def state_dict(self):
+        fast_state_dict = self.base_optimizer.state_dict()
+        slow_state = {
+            (id(k) if isinstance(k, torch.Tensor) else k): v
+            for k, v in self.state.items()
+        }
+        fast_state = fast_state_dict["state"]
+        param_groups = fast_state_dict["param_groups"]
+        return {
+            "state": fast_state,
+            "slow_state": slow_state,
+            "param_groups": param_groups,
+        }
+
+    def load_state_dict(self, state_dict):
+        fast_state_dict = {
+            "state": state_dict["state"],
+            "param_groups": state_dict["param_groups"],
+        }
+        self.base_optimizer.load_state_dict(fast_state_dict)
+
+        # We want to restore the slow state, but share param_groups reference
+        # with base_optimizer. This is a bit redundant but least code
+        slow_state_new = False
+        if "slow_state" not in state_dict:
+            print("Loading state_dict from optimizer without Lookahead applied.")
+            state_dict["slow_state"] = defaultdict(dict)
+            slow_state_new = True
+        slow_state_dict = {
+            "state": state_dict["slow_state"],
+            "param_groups": state_dict[
+                "param_groups"
+            ],  # this is pointless but saves code
+        }
+        super(Lookahead, self).load_state_dict(slow_state_dict)
+        self.param_groups = (
+            self.base_optimizer.param_groups
+        )  # make both ref same container
+        if slow_state_new:
+            # reapply defaults to catch missing lookahead specific ones
+            for name, default in self.defaults.items():
+                for group in self.param_groups:
+                    group.setdefault(name, default)

PreTrain_MeDSLIP/optim/nadam.py

@@ -0,0 +1,108 @@
+import torch
+from torch.optim import Optimizer
+
+
+class Nadam(Optimizer):
+    """Implements Nadam algorithm (a variant of Adam based on Nesterov momentum).
+
+    It has been proposed in `Incorporating Nesterov Momentum into Adam`__.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 2e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        schedule_decay (float, optional): momentum schedule decay (default: 4e-3)
+
+    __ http://cs229.stanford.edu/proj2015/054_report.pdf
+    __ http://www.cs.toronto.edu/~fritz/absps/momentum.pdf
+
+        Originally taken from: https://github.com/pytorch/pytorch/pull/1408
+        NOTE: Has potential issues but does work well on some problems.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=2e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0,
+        schedule_decay=4e-3,
+    ):
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            schedule_decay=schedule_decay,
+        )
+        super(Nadam, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["m_schedule"] = 1.0
+                    state["exp_avg"] = grad.new().resize_as_(grad).zero_()
+                    state["exp_avg_sq"] = grad.new().resize_as_(grad).zero_()
+
+                # Warming momentum schedule
+                m_schedule = state["m_schedule"]
+                schedule_decay = group["schedule_decay"]
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+                eps = group["eps"]
+                state["step"] += 1
+                t = state["step"]
+
+                if group["weight_decay"] != 0:
+                    grad = grad.add(group["weight_decay"], p.data)
+
+                momentum_cache_t = beta1 * (1.0 - 0.5 * (0.96 ** (t * schedule_decay)))
+                momentum_cache_t_1 = beta1 * (
+                    1.0 - 0.5 * (0.96 ** ((t + 1) * schedule_decay))
+                )
+                m_schedule_new = m_schedule * momentum_cache_t
+                m_schedule_next = m_schedule * momentum_cache_t * momentum_cache_t_1
+                state["m_schedule"] = m_schedule_new
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1.0 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1.0 - beta2, grad, grad)
+                exp_avg_sq_prime = exp_avg_sq / (1.0 - beta2 ** t)
+                denom = exp_avg_sq_prime.sqrt_().add_(eps)
+
+                p.data.addcdiv_(
+                    -group["lr"] * (1.0 - momentum_cache_t) / (1.0 - m_schedule_new),
+                    grad,
+                    denom,
+                )
+                p.data.addcdiv_(
+                    -group["lr"] * momentum_cache_t_1 / (1.0 - m_schedule_next),
+                    exp_avg,
+                    denom,
+                )
+
+        return loss

PreTrain_MeDSLIP/optim/novograd.py

@@ -0,0 +1,90 @@
+"""NovoGrad Optimizer.
+Original impl by Masashi Kimura (Convergence Lab): https://github.com/convergence-lab/novograd
+Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks`
+    - https://arxiv.org/abs/1905.11286
+"""
+
+import torch
+from torch.optim.optimizer import Optimizer
+import math
+
+
+class NovoGrad(Optimizer):
+    def __init__(
+        self,
+        params,
+        grad_averaging=False,
+        lr=0.1,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+    ):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super(NovoGrad, self).__init__(params, defaults)
+        self._lr = lr
+        self._beta1 = betas[0]
+        self._beta2 = betas[1]
+        self._eps = eps
+        self._wd = weight_decay
+        self._grad_averaging = grad_averaging
+
+        self._momentum_initialized = False
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        if not self._momentum_initialized:
+            for group in self.param_groups:
+                for p in group["params"]:
+                    if p.grad is None:
+                        continue
+                    state = self.state[p]
+                    grad = p.grad.data
+                    if grad.is_sparse:
+                        raise RuntimeError("NovoGrad does not support sparse gradients")
+
+                    v = torch.norm(grad) ** 2
+                    m = grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
+                    state["step"] = 0
+                    state["v"] = v
+                    state["m"] = m
+                    state["grad_ema"] = None
+            self._momentum_initialized = True
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                state = self.state[p]
+                state["step"] += 1
+
+                step, v, m = state["step"], state["v"], state["m"]
+                grad_ema = state["grad_ema"]
+
+                grad = p.grad.data
+                g2 = torch.norm(grad) ** 2
+                grad_ema = (
+                    g2
+                    if grad_ema is None
+                    else grad_ema * self._beta2 + g2 * (1.0 - self._beta2)
+                )
+                grad *= 1.0 / (torch.sqrt(grad_ema) + self._eps)
+
+                if self._grad_averaging:
+                    grad *= 1.0 - self._beta1
+
+                g2 = torch.norm(grad) ** 2
+                v = self._beta2 * v + (1.0 - self._beta2) * g2
+                m = self._beta1 * m + (
+                    grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
+                )
+                bias_correction1 = 1 - self._beta1 ** step
+                bias_correction2 = 1 - self._beta2 ** step
+                step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1
+
+                state["v"], state["m"] = v, m
+                state["grad_ema"] = grad_ema
+                p.data.add_(-step_size, m)
+        return loss

PreTrain_MeDSLIP/optim/nvnovograd.py

@@ -0,0 +1,132 @@
+""" Nvidia NovoGrad Optimizer.
+Original impl by Nvidia from Jasper example:
+    - https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecognition/Jasper
+Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks`
+    - https://arxiv.org/abs/1905.11286
+"""
+
+import torch
+from torch.optim.optimizer import Optimizer
+import math
+
+
+class NvNovoGrad(Optimizer):
+    """
+    Implements Novograd algorithm.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.95, 0.98))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        grad_averaging: gradient averaging
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+        grad_averaging=False,
+        amsgrad=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            grad_averaging=grad_averaging,
+            amsgrad=amsgrad,
+        )
+
+        super(NvNovoGrad, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(NvNovoGrad, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+            and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sparse gradients are not supported.")
+                amsgrad = group["amsgrad"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state["max_exp_avg_sq"] = torch.zeros([]).to(
+                            state["exp_avg"].device
+                        )
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                norm = torch.sum(torch.pow(grad, 2))
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(1 - beta2, norm)
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group["eps"])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                grad.div_(denom)
+                if group["weight_decay"] != 0:
+                    grad.add_(group["weight_decay"], p.data)
+                if group["grad_averaging"]:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                p.data.add_(-group["lr"], exp_avg)
+
+        return loss

PreTrain_MeDSLIP/optim/optim_factory.py

@@ -0,0 +1,138 @@
+""" Optimizer Factory w/ Custom Weight Decay
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+from torch import optim as optim
+
+from .adafactor import Adafactor
+from .adahessian import Adahessian
+from .adamp import AdamP
+from .lookahead import Lookahead
+from .nadam import Nadam
+from .novograd import NovoGrad
+from .nvnovograd import NvNovoGrad
+from .radam import RAdam
+from .rmsprop_tf import RMSpropTF
+from .sgdp import SGDP
+
+try:
+    from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+
+def add_weight_decay(model, weight_decay=1e-5, skip_list=()):
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+    return [
+        {"params": no_decay, "weight_decay": 0.0},
+        {"params": decay, "weight_decay": weight_decay},
+    ]
+
+
+def create_optimizer(args, model, filter_bias_and_bn=True):
+    opt_lower = args.opt.lower()
+    weight_decay = args.weight_decay
+    if weight_decay and filter_bias_and_bn:
+        skip = {}
+        if hasattr(model, "no_weight_decay"):
+            skip = model.no_weight_decay()
+        parameters = add_weight_decay(model, weight_decay, skip)
+        weight_decay = 0.0
+    else:
+        parameters = filter(
+            lambda p: p.requires_grad, model.parameters()
+        )  # model.parameters()
+
+    if "fused" in opt_lower:
+        assert (
+            has_apex and torch.cuda.is_available()
+        ), "APEX and CUDA required for fused optimizers"
+
+    opt_args = dict(lr=args.lr, weight_decay=weight_decay)
+    if hasattr(args, "opt_eps") and args.opt_eps is not None:
+        opt_args["eps"] = args.opt_eps
+    if hasattr(args, "opt_betas") and args.opt_betas is not None:
+        opt_args["betas"] = args.opt_betas
+    if hasattr(args, "opt_args") and args.opt_args is not None:
+        opt_args.update(args.opt_args)
+
+    opt_split = opt_lower.split("_")
+    opt_lower = opt_split[-1]
+    if opt_lower == "sgd" or opt_lower == "nesterov":
+        opt_args.pop("eps", None)
+        optimizer = optim.SGD(
+            parameters, momentum=args.momentum, nesterov=True, **opt_args
+        )
+    elif opt_lower == "momentum":
+        opt_args.pop("eps", None)
+        optimizer = optim.SGD(
+            parameters, momentum=args.momentum, nesterov=False, **opt_args
+        )
+    elif opt_lower == "adam":
+        optimizer = optim.Adam(parameters, **opt_args)
+    elif opt_lower == "adamw":
+        optimizer = optim.AdamW(parameters, **opt_args)
+    elif opt_lower == "nadam":
+        optimizer = Nadam(parameters, **opt_args)
+    elif opt_lower == "radam":
+        optimizer = RAdam(parameters, **opt_args)
+    elif opt_lower == "adamp":
+        optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
+    elif opt_lower == "sgdp":
+        optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == "adadelta":
+        optimizer = optim.Adadelta(parameters, **opt_args)
+    elif opt_lower == "adafactor":
+        if not args.lr:
+            opt_args["lr"] = None
+        optimizer = Adafactor(parameters, **opt_args)
+    elif opt_lower == "adahessian":
+        optimizer = Adahessian(parameters, **opt_args)
+    elif opt_lower == "rmsprop":
+        optimizer = optim.RMSprop(
+            parameters, alpha=0.9, momentum=args.momentum, **opt_args
+        )
+    elif opt_lower == "rmsproptf":
+        optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == "novograd":
+        optimizer = NovoGrad(parameters, **opt_args)
+    elif opt_lower == "nvnovograd":
+        optimizer = NvNovoGrad(parameters, **opt_args)
+    elif opt_lower == "fusedsgd":
+        opt_args.pop("eps", None)
+        optimizer = FusedSGD(
+            parameters, momentum=args.momentum, nesterov=True, **opt_args
+        )
+    elif opt_lower == "fusedmomentum":
+        opt_args.pop("eps", None)
+        optimizer = FusedSGD(
+            parameters, momentum=args.momentum, nesterov=False, **opt_args
+        )
+    elif opt_lower == "fusedadam":
+        optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
+    elif opt_lower == "fusedadamw":
+        optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
+    elif opt_lower == "fusedlamb":
+        optimizer = FusedLAMB(parameters, **opt_args)
+    elif opt_lower == "fusednovograd":
+        opt_args.setdefault("betas", (0.95, 0.98))
+        optimizer = FusedNovoGrad(parameters, **opt_args)
+    else:
+        assert False and "Invalid optimizer"
+        raise ValueError
+
+    if len(opt_split) > 1:
+        if opt_split[0] == "lookahead":
+            optimizer = Lookahead(optimizer)
+
+    return optimizer

PreTrain_MeDSLIP/optim/radam.py

@@ -0,0 +1,170 @@
+"""RAdam Optimizer.
+Implementation lifted from: https://github.com/LiyuanLucasLiu/RAdam
+Paper: `On the Variance of the Adaptive Learning Rate and Beyond` - https://arxiv.org/abs/1908.03265
+"""
+import math
+import torch
+from torch.optim.optimizer import Optimizer, required
+
+
+class RAdam(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.buffer = [[None, None, None] for ind in range(10)]
+        super(RAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RAdam, self).__setstate__(state)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError("RAdam does not support sparse gradients")
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
+                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
+                else:
+                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
+                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+
+                state["step"] += 1
+                buffered = self.buffer[int(state["step"] % 10)]
+                if state["step"] == buffered[0]:
+                    N_sma, step_size = buffered[1], buffered[2]
+                else:
+                    buffered[0] = state["step"]
+                    beta2_t = beta2 ** state["step"]
+                    N_sma_max = 2 / (1 - beta2) - 1
+                    N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
+                    buffered[1] = N_sma
+
+                    # more conservative since it's an approximated value
+                    if N_sma >= 5:
+                        step_size = (
+                            group["lr"]
+                            * math.sqrt(
+                                (1 - beta2_t)
+                                * (N_sma - 4)
+                                / (N_sma_max - 4)
+                                * (N_sma - 2)
+                                / N_sma
+                                * N_sma_max
+                                / (N_sma_max - 2)
+                            )
+                            / (1 - beta1 ** state["step"])
+                        )
+                    else:
+                        step_size = group["lr"] / (1 - beta1 ** state["step"])
+                    buffered[2] = step_size
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss
+
+
+class PlainRAdam(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+
+        super(PlainRAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(PlainRAdam, self).__setstate__(state)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError("RAdam does not support sparse gradients")
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
+                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
+                else:
+                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
+                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+
+                state["step"] += 1
+                beta2_t = beta2 ** state["step"]
+                N_sma_max = 2 / (1 - beta2) - 1
+                N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    step_size = (
+                        group["lr"]
+                        * math.sqrt(
+                            (1 - beta2_t)
+                            * (N_sma - 4)
+                            / (N_sma_max - 4)
+                            * (N_sma - 2)
+                            / N_sma
+                            * N_sma_max
+                            / (N_sma_max - 2)
+                        )
+                        / (1 - beta1 ** state["step"])
+                    )
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    step_size = group["lr"] / (1 - beta1 ** state["step"])
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss

PreTrain_MeDSLIP/optim/rmsprop_tf.py

@@ -0,0 +1,160 @@
+""" RMSProp modified to behave like Tensorflow impl
+
+Originally cut & paste from PyTorch RMSProp
+https://github.com/pytorch/pytorch/blob/063946d2b3f3f1e953a2a3b54e0b34f1393de295/torch/optim/rmsprop.py
+Licensed under BSD-Clause 3 (ish), https://github.com/pytorch/pytorch/blob/master/LICENSE
+
+Modifications Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch.optim import Optimizer
+
+
+class RMSpropTF(Optimizer):
+    """Implements RMSprop algorithm (TensorFlow style epsilon)
+
+    NOTE: This is a direct cut-and-paste of PyTorch RMSprop with eps applied before sqrt
+    and a few other modifications to closer match Tensorflow for matching hyper-params.
+
+    Noteworthy changes include:
+    1. Epsilon applied inside square-root
+    2. square_avg initialized to ones
+    3. LR scaling of update accumulated in momentum buffer
+
+    Proposed by G. Hinton in his
+    `course <http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_.
+
+    The centered version first appears in `Generating Sequences
+    With Recurrent Neural Networks <https://arxiv.org/pdf/1308.0850v5.pdf>`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        momentum (float, optional): momentum factor (default: 0)
+        alpha (float, optional): smoothing (decay) constant (default: 0.9)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-10)
+        centered (bool, optional) : if ``True``, compute the centered RMSProp,
+            the gradient is normalized by an estimation of its variance
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        decoupled_decay (bool, optional): decoupled weight decay as per https://arxiv.org/abs/1711.05101
+        lr_in_momentum (bool, optional): learning rate scaling is included in the momentum buffer
+            update as per defaults in Tensorflow
+
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-2,
+        alpha=0.9,
+        eps=1e-10,
+        weight_decay=0,
+        momentum=0.0,
+        centered=False,
+        decoupled_decay=False,
+        lr_in_momentum=True,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= momentum:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= alpha:
+            raise ValueError("Invalid alpha value: {}".format(alpha))
+
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            alpha=alpha,
+            eps=eps,
+            centered=centered,
+            weight_decay=weight_decay,
+            decoupled_decay=decoupled_decay,
+            lr_in_momentum=lr_in_momentum,
+        )
+        super(RMSpropTF, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RMSpropTF, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("momentum", 0)
+            group.setdefault("centered", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("RMSprop does not support sparse gradients")
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["square_avg"] = torch.ones_like(
+                        p.data
+                    )  # PyTorch inits to zero
+                    if group["momentum"] > 0:
+                        state["momentum_buffer"] = torch.zeros_like(p.data)
+                    if group["centered"]:
+                        state["grad_avg"] = torch.zeros_like(p.data)
+
+                square_avg = state["square_avg"]
+                one_minus_alpha = 1.0 - group["alpha"]
+
+                state["step"] += 1
+
+                if group["weight_decay"] != 0:
+                    if "decoupled_decay" in group and group["decoupled_decay"]:
+                        p.data.add_(-group["weight_decay"], p.data)
+                    else:
+                        grad = grad.add(group["weight_decay"], p.data)
+
+                # Tensorflow order of ops for updating squared avg
+                square_avg.add_(one_minus_alpha, grad.pow(2) - square_avg)
+                # square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)  # PyTorch original
+
+                if group["centered"]:
+                    grad_avg = state["grad_avg"]
+                    grad_avg.add_(one_minus_alpha, grad - grad_avg)
+                    # grad_avg.mul_(alpha).add_(1 - alpha, grad)  # PyTorch original
+                    avg = (
+                        square_avg.addcmul(-1, grad_avg, grad_avg)
+                        .add(group["eps"])
+                        .sqrt_()
+                    )  # eps moved in sqrt
+                else:
+                    avg = square_avg.add(group["eps"]).sqrt_()  # eps moved in sqrt
+
+                if group["momentum"] > 0:
+                    buf = state["momentum_buffer"]
+                    # Tensorflow accumulates the LR scaling in the momentum buffer
+                    if "lr_in_momentum" in group and group["lr_in_momentum"]:
+                        buf.mul_(group["momentum"]).addcdiv_(group["lr"], grad, avg)
+                        p.data.add_(-buf)
+                    else:
+                        # PyTorch scales the param update by LR
+                        buf.mul_(group["momentum"]).addcdiv_(grad, avg)
+                        p.data.add_(-group["lr"], buf)
+                else:
+                    p.data.addcdiv_(-group["lr"], grad, avg)
+
+        return loss

PreTrain_MeDSLIP/optim/sgdp.py

@@ -0,0 +1,123 @@
+"""
+SGDP Optimizer Implementation copied from https://github.com/clovaai/AdamP/blob/master/adamp/sgdp.py
+
+Paper: `Slowing Down the Weight Norm Increase in Momentum-based Optimizers` - https://arxiv.org/abs/2006.08217
+Code: https://github.com/clovaai/AdamP
+
+Copyright (c) 2020-present NAVER Corp.
+MIT license
+"""
+
+import torch
+import torch.nn as nn
+from torch.optim.optimizer import Optimizer, required
+import math
+
+
+class SGDP(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=required,
+        momentum=0,
+        dampening=0,
+        weight_decay=0,
+        nesterov=False,
+        eps=1e-8,
+        delta=0.1,
+        wd_ratio=0.1,
+    ):
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            dampening=dampening,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            eps=eps,
+            delta=delta,
+            wd_ratio=wd_ratio,
+        )
+        super(SGDP, self).__init__(params, defaults)
+
+    def _channel_view(self, x):
+        return x.view(x.size(0), -1)
+
+    def _layer_view(self, x):
+        return x.view(1, -1)
+
+    def _cosine_similarity(self, x, y, eps, view_func):
+        x = view_func(x)
+        y = view_func(y)
+
+        x_norm = x.norm(dim=1).add_(eps)
+        y_norm = y.norm(dim=1).add_(eps)
+        dot = (x * y).sum(dim=1)
+
+        return dot.abs() / x_norm / y_norm
+
+    def _projection(self, p, grad, perturb, delta, wd_ratio, eps):
+        wd = 1
+        expand_size = [-1] + [1] * (len(p.shape) - 1)
+        for view_func in [self._channel_view, self._layer_view]:
+
+            cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func)
+
+            if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)):
+                p_n = p.data / view_func(p.data).norm(dim=1).view(expand_size).add_(eps)
+                perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view(expand_size)
+                wd = wd_ratio
+
+                return perturb, wd
+
+        return perturb, wd
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            dampening = group["dampening"]
+            nesterov = group["nesterov"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["momentum"] = torch.zeros_like(p.data)
+
+                # SGD
+                buf = state["momentum"]
+                buf.mul_(momentum).add_(1 - dampening, grad)
+                if nesterov:
+                    d_p = grad + momentum * buf
+                else:
+                    d_p = buf
+
+                # Projection
+                wd_ratio = 1
+                if len(p.shape) > 1:
+                    d_p, wd_ratio = self._projection(
+                        p, grad, d_p, group["delta"], group["wd_ratio"], group["eps"]
+                    )
+
+                # Weight decay
+                if weight_decay != 0:
+                    p.data.mul_(
+                        1
+                        - group["lr"]
+                        * group["weight_decay"]
+                        * wd_ratio
+                        / (1 - momentum)
+                    )
+
+                # Step
+                p.data.add_(-group["lr"], d_p)
+
+        return loss

PreTrain_MeDSLIP/scheduler/__init__.py

@@ -0,0 +1,5 @@
+from .cosine_lr import CosineLRScheduler
+from .plateau_lr import PlateauLRScheduler
+from .step_lr import StepLRScheduler
+from .tanh_lr import TanhLRScheduler
+from .scheduler_factory import create_scheduler

PreTrain_MeDSLIP/scheduler/cosine_lr.py

@@ -0,0 +1,136 @@
+""" Cosine Scheduler
+
+Cosine LR schedule with warmup, cycle/restarts, noise.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import logging
+import math
+import numpy as np
+import torch
+
+from .scheduler import Scheduler
+
+from pdb import set_trace as breakpoint
+
+_logger = logging.getLogger(__name__)
+
+
+class CosineLRScheduler(Scheduler):
+    """
+    Cosine decay with restarts.
+    This is described in the paper https://arxiv.org/abs/1608.03983.
+
+    Inspiration from
+    https://github.com/allenai/allennlp/blob/master/allennlp/training/learning_rate_schedulers/cosine.py
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        t_initial: int,
+        t_mul: float = 1.0,
+        lr_min: float = 0.0,
+        decay_rate: float = 1.0,
+        warmup_t=0,
+        warmup_lr_init=0,
+        warmup_prefix=True,
+        cycle_limit=0,
+        t_in_epochs=True,
+        noise_range_t=None,
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=42,
+        initialize=True,
+    ) -> None:
+        super().__init__(
+            optimizer,
+            param_group_field="lr",
+            noise_range_t=noise_range_t,
+            noise_pct=noise_pct,
+            noise_std=noise_std,
+            noise_seed=noise_seed,
+            initialize=initialize,
+        )
+
+        assert t_initial > 0
+        assert lr_min >= 0
+        if t_initial == 1 and t_mul == 1 and decay_rate == 1:
+            _logger.warning(
+                "Cosine annealing scheduler will have no effect on the learning "
+                "rate since t_initial = t_mul = eta_mul = 1."
+            )
+        self.t_initial = t_initial
+        self.t_mul = t_mul
+        self.lr_min = lr_min
+        self.decay_rate = decay_rate
+        self.cycle_limit = cycle_limit
+        self.warmup_t = warmup_t
+        self.warmup_lr_init = warmup_lr_init
+        self.warmup_prefix = warmup_prefix
+        self.t_in_epochs = t_in_epochs
+        if self.warmup_t:
+            self.warmup_steps = [
+                (v - warmup_lr_init) / self.warmup_t for v in self.base_values
+            ]
+            super().update_groups(self.warmup_lr_init)
+        else:
+            self.warmup_steps = [1 for _ in self.base_values]
+
+    def _get_lr(self, t):
+        if t < self.warmup_t:
+            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
+        else:
+            if self.warmup_prefix:
+                t = t - self.warmup_t
+
+            if self.t_mul != 1:
+                i = math.floor(
+                    math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul)
+                )
+                t_i = self.t_mul ** i * self.t_initial
+                t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial
+            else:
+                i = t // self.t_initial
+                t_i = self.t_initial
+                t_curr = t - (self.t_initial * i)
+
+            gamma = self.decay_rate ** i
+            lr_min = self.lr_min * gamma
+            lr_max_values = [v * gamma for v in self.base_values]
+
+            if self.cycle_limit == 0 or (self.cycle_limit > 0 and i < self.cycle_limit):
+                lrs = [
+                    lr_min
+                    + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t_curr / t_i))
+                    for lr_max in lr_max_values
+                ]
+            else:
+                lrs = [self.lr_min for _ in self.base_values]
+
+        return lrs
+
+    def get_epoch_values(self, epoch: int):
+        if self.t_in_epochs:
+            return self._get_lr(epoch)
+        else:
+            return None
+
+    def get_update_values(self, num_updates: int):
+        if not self.t_in_epochs:
+            return self._get_lr(num_updates)
+        else:
+            return None
+
+    def get_cycle_length(self, cycles=0):
+        if not cycles:
+            cycles = self.cycle_limit
+        cycles = max(1, cycles)
+        if self.t_mul == 1.0:
+            return self.t_initial * cycles
+        else:
+            return int(
+                math.floor(
+                    -self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul)
+                )
+            )

PreTrain_MeDSLIP/scheduler/plateau_lr.py

@@ -0,0 +1,116 @@
+""" Plateau Scheduler
+
+Adapts PyTorch plateau scheduler and allows application of noise, warmup.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+
+from .scheduler import Scheduler
+
+
+class PlateauLRScheduler(Scheduler):
+    """Decay the LR by a factor every time the validation loss plateaus."""
+
+    def __init__(
+        self,
+        optimizer,
+        decay_rate=0.1,
+        patience_t=10,
+        verbose=True,
+        threshold=1e-4,
+        cooldown_t=0,
+        warmup_t=0,
+        warmup_lr_init=0,
+        lr_min=0,
+        mode="max",
+        noise_range_t=None,
+        noise_type="normal",
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=None,
+        initialize=True,
+    ):
+        super().__init__(optimizer, "lr", initialize=initialize)
+
+        self.lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            self.optimizer,
+            patience=patience_t,
+            factor=decay_rate,
+            verbose=verbose,
+            threshold=threshold,
+            cooldown=cooldown_t,
+            mode=mode,
+            min_lr=lr_min,
+        )
+
+        self.noise_range = noise_range_t
+        self.noise_pct = noise_pct
+        self.noise_type = noise_type
+        self.noise_std = noise_std
+        self.noise_seed = noise_seed if noise_seed is not None else 42
+        self.warmup_t = warmup_t
+        self.warmup_lr_init = warmup_lr_init
+        if self.warmup_t:
+            self.warmup_steps = [
+                (v - warmup_lr_init) / self.warmup_t for v in self.base_values
+            ]
+            super().update_groups(self.warmup_lr_init)
+        else:
+            self.warmup_steps = [1 for _ in self.base_values]
+        self.restore_lr = None
+
+    def state_dict(self):
+        return {
+            "best": self.lr_scheduler.best,
+            "last_epoch": self.lr_scheduler.last_epoch,
+        }
+
+    def load_state_dict(self, state_dict):
+        self.lr_scheduler.best = state_dict["best"]
+        if "last_epoch" in state_dict:
+            self.lr_scheduler.last_epoch = state_dict["last_epoch"]
+
+    # override the base class step fn completely
+    def step(self, epoch, metric=None):
+        if epoch <= self.warmup_t:
+            lrs = [self.warmup_lr_init + epoch * s for s in self.warmup_steps]
+            super().update_groups(lrs)
+        else:
+            if self.restore_lr is not None:
+                # restore actual LR from before our last noise perturbation before stepping base
+                for i, param_group in enumerate(self.optimizer.param_groups):
+                    param_group["lr"] = self.restore_lr[i]
+                self.restore_lr = None
+
+            self.lr_scheduler.step(metric, epoch)  # step the base scheduler
+
+            if self.noise_range is not None:
+                if isinstance(self.noise_range, (list, tuple)):
+                    apply_noise = self.noise_range[0] <= epoch < self.noise_range[1]
+                else:
+                    apply_noise = epoch >= self.noise_range
+                if apply_noise:
+                    self._apply_noise(epoch)
+
+    def _apply_noise(self, epoch):
+        g = torch.Generator()
+        g.manual_seed(self.noise_seed + epoch)
+        if self.noise_type == "normal":
+            while True:
+                # resample if noise out of percent limit, brute force but shouldn't spin much
+                noise = torch.randn(1, generator=g).item()
+                if abs(noise) < self.noise_pct:
+                    break
+        else:
+            noise = 2 * (torch.rand(1, generator=g).item() - 0.5) * self.noise_pct
+
+        # apply the noise on top of previous LR, cache the old value so we can restore for normal
+        # stepping of base scheduler
+        restore_lr = []
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            old_lr = float(param_group["lr"])
+            restore_lr.append(old_lr)
+            new_lr = old_lr + old_lr * noise
+            param_group["lr"] = new_lr
+        self.restore_lr = restore_lr

PreTrain_MeDSLIP/scheduler/scheduler.py

@@ -0,0 +1,120 @@
+from typing import Dict, Any
+
+import torch
+
+
+class Scheduler:
+    """ Parameter Scheduler Base Class
+    A scheduler base class that can be used to schedule any optimizer parameter groups.
+
+    Unlike the builtin PyTorch schedulers, this is intended to be consistently called
+    * At the END of each epoch, before incrementing the epoch count, to calculate next epoch's value
+    * At the END of each optimizer update, after incrementing the update count, to calculate next update's value
+
+    The schedulers built on this should try to remain as stateless as possible (for simplicity).
+
+    This family of schedulers is attempting to avoid the confusion of the meaning of 'last_epoch'
+    and -1 values for special behaviour. All epoch and update counts must be tracked in the training
+    code and explicitly passed in to the schedulers on the corresponding step or step_update call.
+
+    Based on ideas from:
+     * https://github.com/pytorch/fairseq/tree/master/fairseq/optim/lr_scheduler
+     * https://github.com/allenai/allennlp/tree/master/allennlp/training/learning_rate_schedulers
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        param_group_field: str,
+        noise_range_t=None,
+        noise_type="normal",
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=None,
+        initialize: bool = True,
+    ) -> None:
+        self.optimizer = optimizer
+        self.param_group_field = param_group_field
+        self._initial_param_group_field = f"initial_{param_group_field}"
+        if initialize:
+            for i, group in enumerate(self.optimizer.param_groups):
+                if param_group_field not in group:
+                    raise KeyError(
+                        f"{param_group_field} missing from param_groups[{i}]"
+                    )
+                group.setdefault(
+                    self._initial_param_group_field, group[param_group_field]
+                )
+        else:
+            for i, group in enumerate(self.optimizer.param_groups):
+                if self._initial_param_group_field not in group:
+                    raise KeyError(
+                        f"{self._initial_param_group_field} missing from param_groups[{i}]"
+                    )
+        self.base_values = [
+            group[self._initial_param_group_field]
+            for group in self.optimizer.param_groups
+        ]
+        self.metric = None  # any point to having this for all?
+        self.noise_range_t = noise_range_t
+        self.noise_pct = noise_pct
+        self.noise_type = noise_type
+        self.noise_std = noise_std
+        self.noise_seed = noise_seed if noise_seed is not None else 42
+        self.update_groups(self.base_values)
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {
+            key: value for key, value in self.__dict__.items() if key != "optimizer"
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+
+    def get_epoch_values(self, epoch: int):
+        return None
+
+    def get_update_values(self, num_updates: int):
+        return None
+
+    def step(self, epoch: int, metric: float = None) -> None:
+        self.metric = metric
+        values = self.get_epoch_values(epoch)
+        if values is not None:
+            values = self._add_noise(values, epoch)
+            self.update_groups(values)
+
+    def step_update(self, num_updates: int, metric: float = None):
+        self.metric = metric
+        values = self.get_update_values(num_updates)
+        if values is not None:
+            values = self._add_noise(values, num_updates)
+            self.update_groups(values)
+
+    def update_groups(self, values):
+        if not isinstance(values, (list, tuple)):
+            values = [values] * len(self.optimizer.param_groups)
+        for param_group, value in zip(self.optimizer.param_groups, values):
+            param_group[self.param_group_field] = value
+
+    def _add_noise(self, lrs, t):
+        if self.noise_range_t is not None:
+            if isinstance(self.noise_range_t, (list, tuple)):
+                apply_noise = self.noise_range_t[0] <= t < self.noise_range_t[1]
+            else:
+                apply_noise = t >= self.noise_range_t
+            if apply_noise:
+                g = torch.Generator()
+                g.manual_seed(self.noise_seed + t)
+                if self.noise_type == "normal":
+                    while True:
+                        # resample if noise out of percent limit, brute force but shouldn't spin much
+                        noise = torch.randn(1, generator=g).item()
+                        if abs(noise) < self.noise_pct:
+                            break
+                else:
+                    noise = (
+                        2 * (torch.rand(1, generator=g).item() - 0.5) * self.noise_pct
+                    )
+                lrs = [v + v * noise for v in lrs]
+        return lrs

PreTrain_MeDSLIP/scheduler/scheduler_factory.py

@@ -0,0 +1,87 @@
+""" Scheduler Factory
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from .cosine_lr import CosineLRScheduler
+from .tanh_lr import TanhLRScheduler
+from .step_lr import StepLRScheduler
+from .plateau_lr import PlateauLRScheduler
+
+
+def create_scheduler(args, optimizer):
+    num_epochs = args.epochs
+
+    if getattr(args, "lr_noise", None) is not None:
+        lr_noise = getattr(args, "lr_noise")
+        if isinstance(lr_noise, (list, tuple)):
+            noise_range = [n * num_epochs for n in lr_noise]
+            if len(noise_range) == 1:
+                noise_range = noise_range[0]
+        else:
+            noise_range = lr_noise * num_epochs
+    else:
+        noise_range = None
+
+    lr_scheduler = None
+    if args.sched == "cosine":
+        lr_scheduler = CosineLRScheduler(
+            optimizer,
+            t_initial=num_epochs,
+            t_mul=getattr(args, "lr_cycle_mul", 1.0),
+            lr_min=args.min_lr,
+            decay_rate=args.decay_rate,
+            warmup_lr_init=args.warmup_lr,
+            warmup_t=args.warmup_epochs,
+            cycle_limit=getattr(args, "lr_cycle_limit", 1),
+            t_in_epochs=True,
+            noise_range_t=noise_range,
+            noise_pct=getattr(args, "lr_noise_pct", 0.67),
+            noise_std=getattr(args, "lr_noise_std", 1.0),
+            noise_seed=getattr(args, "seed", 42),
+        )
+        num_epochs = lr_scheduler.get_cycle_length() + args.cooldown_epochs
+    elif args.sched == "tanh":
+        lr_scheduler = TanhLRScheduler(
+            optimizer,
+            t_initial=num_epochs,
+            t_mul=getattr(args, "lr_cycle_mul", 1.0),
+            lr_min=args.min_lr,
+            warmup_lr_init=args.warmup_lr,
+            warmup_t=args.warmup_epochs,
+            cycle_limit=getattr(args, "lr_cycle_limit", 1),
+            t_in_epochs=True,
+            noise_range_t=noise_range,
+            noise_pct=getattr(args, "lr_noise_pct", 0.67),
+            noise_std=getattr(args, "lr_noise_std", 1.0),
+            noise_seed=getattr(args, "seed", 42),
+        )
+        num_epochs = lr_scheduler.get_cycle_length() + args.cooldown_epochs
+    elif args.sched == "step":
+        lr_scheduler = StepLRScheduler(
+            optimizer,
+            decay_t=args.decay_epochs,
+            decay_rate=args.decay_rate,
+            warmup_lr_init=args.warmup_lr,
+            warmup_t=args.warmup_epochs,
+            noise_range_t=noise_range,
+            noise_pct=getattr(args, "lr_noise_pct", 0.67),
+            noise_std=getattr(args, "lr_noise_std", 1.0),
+            noise_seed=getattr(args, "seed", 42),
+        )
+    elif args.sched == "plateau":
+        mode = "min" if "loss" in getattr(args, "eval_metric", "") else "max"
+        lr_scheduler = PlateauLRScheduler(
+            optimizer,
+            decay_rate=args.decay_rate,
+            patience_t=args.patience_epochs,
+            lr_min=args.min_lr,
+            mode=mode,
+            warmup_lr_init=args.warmup_lr,
+            warmup_t=args.warmup_epochs,
+            cooldown_t=0,
+            noise_range_t=noise_range,
+            noise_pct=getattr(args, "lr_noise_pct", 0.67),
+            noise_std=getattr(args, "lr_noise_std", 1.0),
+            noise_seed=getattr(args, "seed", 42),
+        )
+
+    return lr_scheduler, num_epochs

PreTrain_MeDSLIP/scheduler/step_lr.py

@@ -0,0 +1,73 @@
+""" Step Scheduler
+
+Basic step LR schedule with warmup, noise.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import math
+import torch
+
+from .scheduler import Scheduler
+
+
+class StepLRScheduler(Scheduler):
+    """
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        decay_t: float,
+        decay_rate: float = 1.0,
+        warmup_t=0,
+        warmup_lr_init=0,
+        t_in_epochs=True,
+        noise_range_t=None,
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=42,
+        initialize=True,
+    ) -> None:
+        super().__init__(
+            optimizer,
+            param_group_field="lr",
+            noise_range_t=noise_range_t,
+            noise_pct=noise_pct,
+            noise_std=noise_std,
+            noise_seed=noise_seed,
+            initialize=initialize,
+        )
+
+        self.decay_t = decay_t
+        self.decay_rate = decay_rate
+        self.warmup_t = warmup_t
+        self.warmup_lr_init = warmup_lr_init
+        self.t_in_epochs = t_in_epochs
+        if self.warmup_t:
+            self.warmup_steps = [
+                (v - warmup_lr_init) / self.warmup_t for v in self.base_values
+            ]
+            super().update_groups(self.warmup_lr_init)
+        else:
+            self.warmup_steps = [1 for _ in self.base_values]
+
+    def _get_lr(self, t):
+        if t < self.warmup_t:
+            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
+        else:
+            lrs = [
+                v * (self.decay_rate ** (t // self.decay_t)) for v in self.base_values
+            ]
+        return lrs
+
+    def get_epoch_values(self, epoch: int):
+        if self.t_in_epochs:
+            return self._get_lr(epoch)
+        else:
+            return None
+
+    def get_update_values(self, num_updates: int):
+        if not self.t_in_epochs:
+            return self._get_lr(num_updates)
+        else:
+            return None

PreTrain_MeDSLIP/scheduler/tanh_lr.py

@@ -0,0 +1,141 @@
+""" TanH Scheduler
+
+TanH schedule with warmup, cycle/restarts, noise.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import logging
+import math
+import numpy as np
+import torch
+
+from .scheduler import Scheduler
+
+
+_logger = logging.getLogger(__name__)
+
+
+class TanhLRScheduler(Scheduler):
+    """
+    Hyberbolic-Tangent decay with restarts.
+    This is described in the paper https://arxiv.org/abs/1806.01593
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        t_initial: int,
+        lb: float = -6.0,
+        ub: float = 4.0,
+        t_mul: float = 1.0,
+        lr_min: float = 0.0,
+        decay_rate: float = 1.0,
+        warmup_t=0,
+        warmup_lr_init=0,
+        warmup_prefix=False,
+        cycle_limit=0,
+        t_in_epochs=True,
+        noise_range_t=None,
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=42,
+        initialize=True,
+    ) -> None:
+        super().__init__(
+            optimizer,
+            param_group_field="lr",
+            noise_range_t=noise_range_t,
+            noise_pct=noise_pct,
+            noise_std=noise_std,
+            noise_seed=noise_seed,
+            initialize=initialize,
+        )
+
+        assert t_initial > 0
+        assert lr_min >= 0
+        assert lb < ub
+        assert cycle_limit >= 0
+        assert warmup_t >= 0
+        assert warmup_lr_init >= 0
+        self.lb = lb
+        self.ub = ub
+        self.t_initial = t_initial
+        self.t_mul = t_mul
+        self.lr_min = lr_min
+        self.decay_rate = decay_rate
+        self.cycle_limit = cycle_limit
+        self.warmup_t = warmup_t
+        self.warmup_lr_init = warmup_lr_init
+        self.warmup_prefix = warmup_prefix
+        self.t_in_epochs = t_in_epochs
+        if self.warmup_t:
+            t_v = (
+                self.base_values if self.warmup_prefix else self._get_lr(self.warmup_t)
+            )
+            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in t_v]
+            super().update_groups(self.warmup_lr_init)
+        else:
+            self.warmup_steps = [1 for _ in self.base_values]
+
+    def _get_lr(self, t):
+        if t < self.warmup_t:
+            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
+        else:
+            if self.warmup_prefix:
+                t = t - self.warmup_t
+
+            if self.t_mul != 1:
+                i = math.floor(
+                    math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul)
+                )
+                t_i = self.t_mul ** i * self.t_initial
+                t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial
+            else:
+                i = t // self.t_initial
+                t_i = self.t_initial
+                t_curr = t - (self.t_initial * i)
+
+            if self.cycle_limit == 0 or (self.cycle_limit > 0 and i < self.cycle_limit):
+                gamma = self.decay_rate ** i
+                lr_min = self.lr_min * gamma
+                lr_max_values = [v * gamma for v in self.base_values]
+
+                tr = t_curr / t_i
+                lrs = [
+                    lr_min
+                    + 0.5
+                    * (lr_max - lr_min)
+                    * (1 - math.tanh(self.lb * (1.0 - tr) + self.ub * tr))
+                    for lr_max in lr_max_values
+                ]
+            else:
+                lrs = [
+                    self.lr_min * (self.decay_rate ** self.cycle_limit)
+                    for _ in self.base_values
+                ]
+        return lrs
+
+    def get_epoch_values(self, epoch: int):
+        if self.t_in_epochs:
+            return self._get_lr(epoch)
+        else:
+            return None
+
+    def get_update_values(self, num_updates: int):
+        if not self.t_in_epochs:
+            return self._get_lr(num_updates)
+        else:
+            return None
+
+    def get_cycle_length(self, cycles=0):
+        if not cycles:
+            cycles = self.cycle_limit
+        cycles = max(1, cycles)
+        if self.t_mul == 1.0:
+            return self.t_initial * cycles
+        else:
+            return int(
+                math.floor(
+                    -self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul)
+                )
+            )

PreTrain_MeDSLIP/train_MeDSLIP.py

@@ -0,0 +1,446 @@
+import argparse
+import os
+import ruamel_yaml as yaml
+import numpy as np
+import random
+import time
+import datetime
+import json
+from pathlib import Path
+import warnings
+
+warnings.filterwarnings("ignore")
+
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+
+from tensorboardX import SummaryWriter
+
+import utils
+from scheduler import create_scheduler
+from optim import create_optimizer
+from dataset.dataset import MeDSLIP_Dataset
+from models.model_MeDSLIP import MeDSLIP
+from models.tokenization_bert import BertTokenizer
+
+
+def get_tokenizer(tokenizer, target_text):
+
+    target_tokenizer = tokenizer(
+        list(target_text),
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+        return_tensors="pt",
+    )
+
+    return target_tokenizer
+
+
+def train(
+    model,
+    data_loader,
+    optimizer,
+    epoch,
+    warmup_steps,
+    device,
+    scheduler,
+    args,
+    config,
+    writer,
+):
+    model.train()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter(
+        "lr", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss_ce_p", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss_cl_p", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss_ce_a", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss_cl_a", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.add_meter(
+        "loss_ap", utils.SmoothedValue(window_size=50, fmt="{value:.6f}")
+    )
+    metric_logger.update(loss=1.0)
+    metric_logger.update(loss_ce_p=1.0)
+    metric_logger.update(loss_cl_p=1.0)
+    metric_logger.update(loss_ce_a=1.0)
+    metric_logger.update(loss_cl_a=1.0)
+    metric_logger.update(loss_ap=1.0)
+    metric_logger.update(lr=scheduler._get_lr(epoch)[0])
+
+    header = "Train Epoch: [{}]".format(epoch)
+    print_freq = 1
+    step_size = 100
+    warmup_iterations = warmup_steps * step_size
+    scalar_step = epoch * len(data_loader)
+
+    for i, sample in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+
+        images = sample["image"].to(device)
+        labels_pathology = sample["label_pathology"].to(device)
+        labels_anatomy = sample["label_anatomy"].to(device)
+        index_pathology = sample["index_pathology"].to(device)
+        index_anatomy = sample["index_anatomy"].to(device)
+        matrix = sample["matrix"].to(device)
+
+        optimizer.zero_grad()
+
+        (
+            loss,
+            loss_ce_pathology,
+            loss_cl_pathology,
+            loss_ce_anatomy,
+            loss_cl_anatomy,
+            loss_ap,
+        ) = model(
+            images,
+            labels_pathology=labels_pathology,
+            labels_anatomy=labels_anatomy,
+            matrix=matrix,
+            sample_index_pathology=index_pathology,
+            sample_index_anatomy=index_anatomy,
+            is_train=True,
+            no_cl=config["no_cl"],
+            exclude_class=config["exclude_class"],
+        )
+        loss.backward()
+        optimizer.step()
+        writer.add_scalar("loss/loss", loss, scalar_step)
+        writer.add_scalar("loss/loss_ce_pathology", loss_ce_pathology, scalar_step)
+        writer.add_scalar("loss/loss_cl_pathology", loss_cl_pathology, scalar_step)
+        writer.add_scalar("loss/loss_ce_anatomy", loss_ce_anatomy, scalar_step)
+        writer.add_scalar("loss/loss_cl_anatomy", loss_cl_anatomy, scalar_step)
+        writer.add_scalar("loss/loss_ap", loss_ap, scalar_step)
+        scalar_step += 1
+        metric_logger.update(loss_ce_p=loss_ce_pathology.item())
+        metric_logger.update(loss_cl_p=loss_cl_pathology.item())
+        metric_logger.update(loss_ce_a=loss_ce_anatomy.item())
+        metric_logger.update(loss_cl_a=loss_cl_anatomy.item())
+        metric_logger.update(loss_ap=loss_ap.item())
+        metric_logger.update(loss=loss.item())
+        # metric_logger.update(loss_cl=loss_cl.item())
+        if epoch == 0 and i % step_size == 0 and i <= warmup_iterations:
+            scheduler.step(i // step_size)
+        metric_logger.update(lr=scheduler._get_lr(epoch)[0])
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger.global_avg())
+    return {
+        k: "{:.3f}".format(meter.global_avg)
+        for k, meter in metric_logger.meters.items()
+    }
+
+
+def valid(model, data_loader, epoch, device, config, writer):
+    model.eval()
+    val_scalar_step = epoch * len(data_loader)
+    val_loss = []
+    for i, sample in enumerate(data_loader):
+
+        images = sample["image"].to(device)
+        labels_pathology = sample["label_pathology"].to(device)
+        labels_anatomy = sample["label_anatomy"].to(device)
+        index_pathology = sample["index_pathology"].to(device)
+        index_anatomy = sample["index_anatomy"].to(device)
+        matrix = sample["matrix"].to(device)
+
+        with torch.no_grad():
+            (
+                loss,
+                loss_ce_pathology,
+                loss_cl_pathology,
+                loss_ce_anatomy,
+                loss_cl_anatomy,
+                loss_ap,
+            ) = model(
+                images,
+                labels_pathology=labels_pathology,
+                labels_anatomy=labels_anatomy,
+                matrix=matrix,
+                sample_index_pathology=index_pathology,
+                sample_index_anatomy=index_anatomy,
+                is_train=True,
+                no_cl=config["no_cl"],
+                exclude_class=config["exclude_class"],
+            )
+            val_loss.append(loss.item())
+            writer.add_scalar("val_loss/loss", loss, val_scalar_step)
+            writer.add_scalar(
+                "val_loss/loss_ce_pathology", loss_ce_pathology, val_scalar_step
+            )
+            writer.add_scalar(
+                "val_loss/loss_cl_pathology", loss_cl_pathology, val_scalar_step
+            )
+            writer.add_scalar(
+                "val_loss/loss_ce_anatomy", loss_ce_anatomy, val_scalar_step
+            )
+            writer.add_scalar(
+                "val_loss/loss_cl_anatomy", loss_cl_anatomy, val_scalar_step
+            )
+            writer.add_scalar("val_loss/loss_ap", loss_ap, val_scalar_step)
+            val_scalar_step += 1
+    avg_val_loss = np.array(val_loss).mean()
+    return avg_val_loss
+
+
+def main(args, config):
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if args.computing == "parallel":
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        device = torch.device("cuda", rank)
+        print("World size: ", world_size, "; Rank: ", rank)
+
+    print("Total CUDA devices: ", torch.cuda.device_count())
+    torch.set_default_tensor_type("torch.FloatTensor")
+    cudnn.benchmark = True
+
+    start_epoch = 0
+    max_epoch = config["schedular"]["epochs"]
+    warmup_steps = config["schedular"]["warmup_epochs"]
+
+    #### Dataset ####
+    print("Creating dataset")
+    train_datasets = MeDSLIP_Dataset(
+        config["train_file"], config["label_file"], mode="train"
+    )
+    val_datasets = MeDSLIP_Dataset(
+        config["valid_file"], config["label_file"], mode="train"
+    )
+    if args.computing == "parallel":
+        # shuffl
+        train_sampler = torch.utils.data.distributed.DistributedSampler(
+            train_datasets, num_replicas=world_size, rank=rank, shuffle=True
+        )
+        val_sampler = torch.utils.data.distributed.DistributedSampler(
+            val_datasets, num_replicas=world_size, rank=rank, shuffle=True
+        )
+    else:
+        train_sampler = torch.utils.data.RandomSampler(train_datasets)
+        val_sampler = torch.utils.data.RandomSampler(val_datasets)
+    train_dataloader = DataLoader(
+        train_datasets,
+        batch_size=config["batch_size"],
+        num_workers=30,
+        pin_memory=True,
+        sampler=train_sampler,
+        collate_fn=None,
+        drop_last=True,
+    )
+
+    val_dataloader = DataLoader(
+        val_datasets,
+        batch_size=config["batch_size"],
+        num_workers=30,
+        pin_memory=True,
+        sampler=val_sampler,
+        collate_fn=None,
+        drop_last=True,
+    )
+
+    print("Creating book")
+    json_book = json.load(open(config["pathology_book"], "r"))
+    pathology_book = [json_book[i] for i in json_book]
+    anatomy_list = [
+        "trachea",
+        "left_hilar",
+        "right_hilar",
+        "hilar_unspec",
+        "left_pleural",
+        "right_pleural",
+        "pleural_unspec",
+        "heart_size",
+        "heart_border",
+        "left_diaphragm",
+        "right_diaphragm",
+        "diaphragm_unspec",
+        "retrocardiac",
+        "lower_left_lobe",
+        "upper_left_lobe",
+        "lower_right_lobe",
+        "middle_right_lobe",
+        "upper_right_lobe",
+        "left_lower_lung",
+        "left_mid_lung",
+        "left_upper_lung",
+        "left_apical_lung",
+        "left_lung_unspec",
+        "right_lower_lung",
+        "right_mid_lung",
+        "right_upper_lung",
+        "right_apical_lung",
+        "right_lung_unspec",
+        "lung_apices",
+        "lung_bases",
+        "left_costophrenic",
+        "right_costophrenic",
+        "costophrenic_unspec",
+        "cardiophrenic_sulcus",
+        "mediastinal",
+        "spine",
+        "clavicle",
+        "rib",
+        "stomach",
+        "right_atrium",
+        "right_ventricle",
+        "aorta",
+        "svc",
+        "interstitium",
+        "parenchymal",
+        "cavoatrial_junction",
+        "cardiopulmonary",
+        "pulmonary",
+        "lung_volumes",
+        "unspecified",
+        "other",
+    ]
+    anatomy_book = []
+    for i in anatomy_list:
+        anatomy_book.append("It is located at " + i + ". ")
+
+    tokenizer = BertTokenizer.from_pretrained(config["text_encoder"])
+    anatomy_book_tokenizer = get_tokenizer(tokenizer, anatomy_book).to(device)
+    pathology_book_tokenizer = get_tokenizer(tokenizer, pathology_book).to(device)
+    print("Creating model")
+    model = MeDSLIP(
+        config, anatomy_book_tokenizer, pathology_book_tokenizer, mode="train"
+    )
+    model = model.to(device)
+    if args.computing == "parallel":
+        model = nn.parallel.DistributedDataParallel(
+            model, device_ids=[rank], find_unused_parameters=True
+        )
+
+    arg_opt = utils.AttrDict(config["optimizer"])
+    optimizer = create_optimizer(arg_opt, model)
+    arg_sche = utils.AttrDict(config["schedular"])
+    lr_scheduler, _ = create_scheduler(arg_sche, optimizer)
+
+    if args.checkpoint:
+        checkpoint = torch.load(args.checkpoint, map_location="cpu")
+        state_dict = checkpoint["model"]
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        start_epoch = checkpoint["epoch"] + 1
+        model.load_state_dict(state_dict)
+        print("load checkpoint from %s" % args.checkpoint)
+
+    print("Start training")
+    start_time = time.time()
+
+    writer = SummaryWriter(os.path.join(args.output_dir, "log"))
+    for epoch in range(start_epoch, max_epoch):
+        if epoch > 0:
+            lr_scheduler.step(epoch + warmup_steps)
+        train_stats = train(
+            model,
+            train_dataloader,
+            optimizer,
+            epoch,
+            warmup_steps,
+            device,
+            lr_scheduler,
+            args,
+            config,
+            writer,
+        )
+
+        for k, v in train_stats.items():
+            train_loss_epoch = v
+
+        writer.add_scalar("loss/train_loss_epoch", float(train_loss_epoch), epoch)
+        writer.add_scalar("loss/leaning_rate", lr_scheduler._get_lr(epoch)[0], epoch)
+
+        val_loss = valid(model, val_dataloader, epoch, device, config, writer)
+        writer.add_scalar("loss/val_loss_epoch", val_loss, epoch)
+
+        if utils.is_main_process():
+            log_stats = {
+                **{f"train_{k}": v for k, v in train_stats.items()},
+                "epoch": epoch,
+                "val_loss": val_loss.item(),
+            }
+            save_obj = {
+                "model": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "config": config,
+                "epoch": epoch,
+            }
+            torch.save(save_obj, os.path.join(args.output_dir, "checkpoint_state.pth"))
+
+            with open(os.path.join(args.output_dir, "log.txt"), "a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+        if epoch % 1 == 0 and epoch > 15:
+            save_obj = {
+                "model": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "config": config,
+                "epoch": epoch,
+            }
+            torch.save(
+                save_obj,
+                os.path.join(args.output_dir, "checkpoint_" + str(epoch) + ".pth"),
+            )
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print("Training time {}".format(total_time_str))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config", default="PreTrain_MeDSLIP/configs/Pretrain_MeDSLIP.yaml"
+    )
+    parser.add_argument("--checkpoint", default="")
+    parser.add_argument("--output_dir", default="runs/")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--local_rank", default=0, type=int)
+    parser.add_argument("--world_size", default=1, type=int)
+    parser.add_argument(
+        "--computing", type=str, default="single", help="number of gpus"
+    )
+    args = parser.parse_args()
+    import datetime
+
+    args.output_dir = os.path.join(
+        args.output_dir, datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
+    )
+
+    gpus = torch.cuda.device_count()
+    if gpus > 1:
+        args.computing = "parallel"
+
+    config = yaml.load(open(args.config, "r"), Loader=yaml.Loader)
+
+    if not Path(args.output_dir).exists():
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    yaml.dump(config, open(os.path.join(args.output_dir, "config.yaml"), "w"))
+
+    if args.computing == "parallel":
+        torch.distributed.init_process_group(backend="nccl", init_method="env://")
+
+    main(args, config)

PreTrain_MeDSLIP/utils.py

@@ -0,0 +1,277 @@
+import numpy as np
+import io
+import os
+import time
+from collections import defaultdict, deque
+import datetime
+
+import torch
+import torch.distributed as dist
+from tqdm import tqdm
+
+import warnings
+
+warnings.filterwarnings("ignore")
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        if self.count == 0:
+            return self.total
+        else:
+            return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def global_avg(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {:.4f}".format(name, meter.global_avg))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        log_msg = ["{meters}"]
+        if torch.cuda.is_available():
+            log_msg.append("max mem: {memory:.0f}")
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+
+        loop = tqdm(iterable)
+        loop.set_description(header)
+
+        for obj in loop:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(loop) - 1:
+                eta_seconds = iter_time.global_avg * (len(loop) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    loop.set_postfix_str(
+                        log_msg.format(
+                            i,
+                            len(loop),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    loop.set_postfix_str(
+                        log_msg.format(
+                            i,
+                            len(loop),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def compute_acc(logits, label, reduction="mean"):
+    ret = (torch.argmax(logits, dim=1) == label).float()
+    if reduction == "none":
+        return ret.detach()
+    elif reduction == "mean":
+        return ret.mean().item()
+
+
+def compute_n_params(model, return_str=True):
+    tot = 0
+    for p in model.parameters():
+        w = 1
+        for x in p.shape:
+            w *= x
+        tot += w
+    if return_str:
+        if tot >= 1e6:
+            return "{:.1f}M".format(tot / 1e6)
+        else:
+            return "{:.1f}K".format(tot / 1e3)
+    else:
+        return tot
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.local_rank = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.local_rank = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.local_rank)
+    args.dist_backend = "nccl"
+    print(
+        "| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True
+    )
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        init_method=args.dist_url,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)

README.md

@@ -1,3 +1,49 @@
----
-license: cc-by-nc-4.0
----
+# MeDSLIP: Medical Knowledge Enhanced Language-Image Pre-Training in Radiology
+
+## Introduction:
+
+The official implementation  code for "MeDSLIP: Medical Knowledge Enhanced Language-Image Pre-Training in Radiology".
+
+[**Arxiv Version**](https://arxiv.org/abs/2403.10635)
+
+## Quick Start:
+Check checkpoints directory to download our pre-trained model from [Hugging Face: MeDSLIP](https://huggingface.co/pykale/MeDSLIP). It can be used for all zero-shot and finetuning tasks.
+
+* **Zero-Shot Classification:**
+
+    We give an example on CXR14 in ```Sample_Zero-Shot_Classification_CXR14```. Change the data paths, and test our model by ```python test.py```.
+    We give an example on RSNA in ```Sample_Zero-Shot_Classification_RSNA```. Change the data paths, and test our model by ```python test.py```.
+
+* **Zero-Shot Grounding:**
+
+    We give an example on RSNA_Pneumonia in ```Sample_Zero-Shot_Grounding_RSNA```. Change the data paths, and test our model by ```python test.py```.
+
+* **Finetuning:**
+
+    We give segmentation and classification finetune code on SIIM_ACR dataset in ```Sample_Finetuning_SIIMACR```. Change the data paths, and finetune our model by ```python I1_classification/train_res_ft.py``` or ```python I2_segementation/train_res_ft.py```.
+
+## Pre-train:
+### Data Preparation
+All files for data preparation files can be downloaded from [Hugging Face: MeDSLIP](https://huggingface.co/pykale/MeDSLIP).
+- Extracted triplets: `landmark_observation_adj_mtx.npy`
+- Training list: `train.json`
+- Validation list: `valid.json`
+- Test list: `test.json`
+
+### Pre-training
+Our pre-train code is given in ```PreTrain_MeDSLIP```.
+* Check the ```PreTrain_MeDSLIP/data_file``` dir and download the files for data preparation.
+* Change the data and preparation files paths as you disire in ```PreTrain_MeDSLIP/configs/Pretrain_MeDSLIP.yaml```, and ```python PreTrain_MeDSLIP/train_MeDSLIP.py``` to pre-train.
+
+## Reference
+```
+@article{fan2024medslip,
+  title={MeDSLIP: Medical Dual-Stream Language-Image Pre-training for Fine-grained Alignment},
+  author={Fan, Wenrui and Suvon, Mohammod Naimul Islam and Zhou, Shuo and Liu, Xianyuan and Alabed, Samer and Osmani, Venet and Swift, Andrew and Chen, Chen and Lu, Haiping},
+  journal={arXiv preprint arXiv:2403.10635},
+  year={2024}
+}
+```
+
+## Contact
+If you have any question, please feel free to contact [email protected].

Sample_Finetuning_SIIMACR/I1_classification/configs/Res_train.yaml

@@ -0,0 +1,17 @@
+train_file: "SIIM-CLS/siim-acr-pneumothorax/stage_1_train_images.csv"
+valid_file: "SIIM-CLS/siim-acr-pneumothorax/stage_1_test_images.csv"
+test_file: "SIIM-CLS/siim-acr-pneumothorax/stage_1_test_images.csv"
+
+image_res: 224
+batch_size: 64
+test_batch_size: 64
+num_classes: 1
+temp: 0.07
+mlm_probability: 0.15
+queue_size: 8192
+momentum: 0.995
+alpha: 0.4
+percentage: 1.0
+
+optimizer: {opt: adamW, lr: 1e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 1e-5, epochs: 200, min_lr: 1e-5, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 20, cooldown_epochs: 0}

Sample_Finetuning_SIIMACR/I1_classification/dataset/dataset_siim_acr.py

@@ -0,0 +1,124 @@
+from cmath import nan
+import csv
+import json
+import logging
+import os
+import sys
+import pydicom
+
+from abc import abstractmethod
+from itertools import islice
+from typing import List, Tuple, Dict, Any
+from torch.utils.data import DataLoader
+import PIL
+from torch.utils.data import Dataset
+import numpy as np
+import pandas as pd
+from torchvision import transforms
+from PIL import Image
+from skimage import exposure
+import torch
+from torchvision.transforms import InterpolationMode
+from dataset.randaugment import RandomAugment
+
+
+class SIIM_ACR_Dataset(Dataset):
+    def __init__(self, csv_path, is_train=True, percentage=0.01):
+        data_info = pd.read_csv(csv_path)
+        if is_train == True:
+            total_len = int(percentage * len(data_info))
+            choice_list = np.random.choice(
+                range(len(data_info)), size=total_len, replace=False
+            )
+            self.img_path_list = data_info["image_path"][choice_list].tolist()
+        else:
+            self.img_path_list = data_info["image_path"].tolist()
+
+        self.img_root = "SIIM-CLS/siim-acr-pneumothorax/png_images/"
+        self.seg_root = "SIIM-CLS/siim-acr-pneumothorax/png_masks/"  # We have pre-processed the original SIIM_ACR data, you may change this to fix your data
+
+        normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+
+        if is_train:
+            self.transform = transforms.Compose(
+                [
+                    transforms.RandomResizedCrop(
+                        224, scale=(0.2, 1.0), interpolation=Image.BICUBIC
+                    ),
+                    transforms.RandomHorizontalFlip(),
+                    RandomAugment(
+                        2,
+                        7,
+                        isPIL=True,
+                        augs=[
+                            "Identity",
+                            "AutoContrast",
+                            "Equalize",
+                            "Brightness",
+                            "Sharpness",
+                            "ShearX",
+                            "ShearY",
+                            "TranslateX",
+                            "TranslateY",
+                            "Rotate",
+                        ],
+                    ),
+                    transforms.ToTensor(),
+                    normalize,
+                ]
+            )
+        else:
+            self.transform = transforms.Compose(
+                [transforms.Resize([224, 224]), transforms.ToTensor(), normalize,]
+            )
+
+        self.seg_transfrom = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Resize([224, 224], interpolation=InterpolationMode.NEAREST),
+            ]
+        )
+
+    def __getitem__(self, index):
+        img_path = self.img_root + self.img_path_list[index].split("/")[-1]  # + ".png"
+        seg_path = (
+            self.seg_root + self.img_path_list[index].split("/")[-1]  # + ".png"
+        )  # We have pre-processed the original SIIM_ACR data, you may change this to fix your data
+        img = PIL.Image.open(img_path).convert("RGB")
+        image = self.transform(img)
+
+        seg_map = PIL.Image.open(seg_path)
+        seg_map = self.seg_transfrom(seg_map)
+        seg_map = (seg_map > 0).type(torch.int)
+        class_label = np.array([int(torch.sum(seg_map) > 0)])
+        return {"image": image, "label": class_label}
+
+    def __len__(self):
+        return len(self.img_path_list)
+
+
+def create_loader_RSNA(
+    datasets, samplers, batch_size, num_workers, is_trains, collate_fns
+):
+    loaders = []
+    for dataset, sampler, bs, n_worker, is_train, collate_fn in zip(
+        datasets, samplers, batch_size, num_workers, is_trains, collate_fns
+    ):
+        if is_train:
+            shuffle = sampler is None
+            drop_last = True
+        else:
+            shuffle = False
+            drop_last = False
+        loader = DataLoader(
+            dataset,
+            batch_size=bs,
+            num_workers=n_worker,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=shuffle,
+            collate_fn=collate_fn,
+            drop_last=drop_last,
+        )
+        loaders.append(loader)
+    return loaders

Sample_Finetuning_SIIMACR/I1_classification/dataset/randaugment.py

@@ -0,0 +1,346 @@
+import cv2
+import numpy as np
+
+
+## aug functions
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    """
+        same output as PIL.ImageOps.autocontrast
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            offset = -low * scale
+            table = np.arange(n_bins) * scale + offset
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    """
+        same output as PIL.ImageOps.equalize
+        PIL's implementation is different from cv2.equalize
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0:
+            return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    """
+    like PIL, rotate by degree, not radians
+    """
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    """
+        same output as PIL.ImageOps.posterize
+    """
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Color
+    """
+    ## implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = np.float32(
+        [[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]]
+    ) * factor + np.float32([[0.114], [0.587], [0.299]])
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = (
+        np.array([(el - mean) * factor + mean for el in range(256)])
+        .clip(0, 255)
+        .astype(np.uint8)
+    )
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    """
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    """
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, factor, 0], [0, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    """
+        same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, -offset], [0, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    """
+        same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [0, 1, -offset]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    """
+        same output as PIL.ImageOps.posterize
+    """
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [factor, 1, 0]])
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+### level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level,)
+
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level,)
+
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    "Identity": identity_func,
+    "AutoContrast": autocontrast_func,
+    "Equalize": equalize_func,
+    "Rotate": rotate_func,
+    "Solarize": solarize_func,
+    "Color": color_func,
+    "Contrast": contrast_func,
+    "Brightness": brightness_func,
+    "Sharpness": sharpness_func,
+    "ShearX": shear_x_func,
+    "TranslateX": translate_x_func,
+    "TranslateY": translate_y_func,
+    "Posterize": posterize_func,
+    "ShearY": shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    "Identity": none_level_to_args,
+    "AutoContrast": none_level_to_args,
+    "Equalize": none_level_to_args,
+    "Rotate": rotate_level_to_args(MAX_LEVEL, replace_value),
+    "Solarize": solarize_level_to_args(MAX_LEVEL),
+    "Color": enhance_level_to_args(MAX_LEVEL),
+    "Contrast": enhance_level_to_args(MAX_LEVEL),
+    "Brightness": enhance_level_to_args(MAX_LEVEL),
+    "Sharpness": enhance_level_to_args(MAX_LEVEL),
+    "ShearX": shear_level_to_args(MAX_LEVEL, replace_value),
+    "TranslateX": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "TranslateY": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "Posterize": posterize_level_to_args(MAX_LEVEL),
+    "ShearY": shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+    def __init__(self, N=2, M=10, isPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        if augs:
+            self.augs = augs
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args)
+        return img
+
+
+if __name__ == "__main__":
+    a = RandomAugment()
+    img = np.random.randn(32, 32, 3)
+    a(img)

Sample_Finetuning_SIIMACR/I1_classification/models/resnet.py

@@ -0,0 +1,88 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import torch
+from einops import rearrange
+
+
+class ModelRes_ft(nn.Module):
+    def __init__(
+        self,
+        res_base_model,
+        out_size,
+        imagenet_pretrain=False,
+        linear_probe=False,
+        use_base=True,
+    ):
+        super(ModelRes_ft, self).__init__()
+        self.resnet_dict = {
+            "resnet18": models.resnet18(pretrained=imagenet_pretrain),
+            "resnet50": models.resnet50(pretrained=imagenet_pretrain),
+        }
+        resnet = self._get_res_basemodel(res_base_model)
+        self.use_base = use_base
+
+        if not self.use_base:
+            num_ftrs = int(resnet.fc.in_features / 2)
+            self.res_features = nn.Sequential(*list(resnet.children())[:-3])
+            self.res_l1_anatomy = nn.Linear(num_ftrs, num_ftrs)
+            self.res_l2_anatomy = nn.Linear(num_ftrs, 256)
+            self.res_l1_pathology = nn.Linear(num_ftrs, num_ftrs)
+            self.res_l2_pathology = nn.Linear(num_ftrs, 256)
+
+            self.mask_generator = nn.Linear(num_ftrs, num_ftrs)
+            self.back = nn.Linear(256, num_ftrs)
+            self.last_res = nn.Sequential(*list(resnet.children())[-3:-1])
+        else:
+            self.res_features = nn.Sequential(*list(resnet.children())[:-1])
+        self.res_out = nn.Linear(int(resnet.fc.in_features), out_size)
+
+    def _get_res_basemodel(self, res_model_name):
+        try:
+            res_model = self.resnet_dict[res_model_name]
+            print("Image feature extractor:", res_model_name)
+            return res_model
+        except:
+            raise (
+                "Invalid model name. Check the config file and pass one of: resnet18 or resnet50"
+            )
+
+    def image_encoder(self, xis):
+        # patch features
+        """
+        16 torch.Size([16, 1024, 14, 14])
+        torch.Size([16, 196, 1024])
+        torch.Size([3136, 1024])
+        torch.Size([16, 196, 256])
+        """
+        batch_size = xis.shape[0]
+        res_fea = self.res_features(xis)  # batch_size,feature_size,patch_num,patch_num
+        res_fea = rearrange(res_fea, "b d n1 n2 -> b (n1 n2) d")
+        x = rearrange(res_fea, "b n d -> (b n) d")
+        mask = self.mask_generator(x)
+        x_pathology = mask * x
+        x_pathology = self.res_l1_pathology(x_pathology)
+        x_pathology = F.relu(x_pathology)
+
+        x_pathology = self.res_l2_pathology(x_pathology)
+
+        out_emb_pathology = rearrange(x_pathology, "(b n) d -> b n d", b=batch_size)
+        out_emb_pathology = self.back(out_emb_pathology)
+        out_emb_pathology = rearrange(out_emb_pathology, "b (n1 n2) d -> b d n1 n2", n1=14, n2=14)
+        out_emb_pathology = self.last_res(out_emb_pathology)
+        out_emb_pathology = out_emb_pathology.squeeze()
+
+        return out_emb_pathology
+
+    def forward(self, img, linear_probe=False):
+        if self.use_base:
+            x = self.res_features(img)
+        else:
+            x = self.image_encoder(img)
+
+        x = x.squeeze()
+        if linear_probe:
+            return x
+        else:
+            x = self.res_out(x)
+            return x

Sample_Finetuning_SIIMACR/I1_classification/optim/__init__.py

@@ -0,0 +1,13 @@
+from .adamp import AdamP
+from .adamw import AdamW
+from .adafactor import Adafactor
+from .adahessian import Adahessian
+from .lookahead import Lookahead
+from .nadam import Nadam
+from .novograd import NovoGrad
+from .nvnovograd import NvNovoGrad
+from .radam import RAdam
+from .rmsprop_tf import RMSpropTF
+from .sgdp import SGDP
+
+from .optim_factory import create_optimizer

Sample_Finetuning_SIIMACR/I1_classification/optim/adafactor.py

@@ -0,0 +1,206 @@
+""" Adafactor Optimizer
+
+Lifted from https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+Original header/copyright below.
+
+"""
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+import math
+
+
+class Adafactor(torch.optim.Optimizer):
+    """Implements Adafactor algorithm.
+    This implementation is based on: `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`
+    (see https://arxiv.org/abs/1804.04235)
+
+    Note that this optimizer internally adjusts the learning rate depending on the
+    *scale_parameter*, *relative_step* and *warmup_init* options.
+
+    To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining parameter groups
+        lr (float, optional): external learning rate (default: None)
+        eps (tuple[float, float]): regularization constants for square gradient
+            and parameter scale respectively (default: (1e-30, 1e-3))
+        clip_threshold (float): threshold of root mean square of final gradient update (default: 1.0)
+        decay_rate (float): coefficient used to compute running averages of square gradient (default: -0.8)
+        beta1 (float): coefficient used for computing running averages of gradient (default: None)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        scale_parameter (bool): if True, learning rate is scaled by root mean square of parameter (default: True)
+        relative_step (bool): if True, time-dependent learning rate is computed
+            instead of external learning rate (default: True)
+        warmup_init (bool): time-dependent learning rate computation depends on
+            whether warm-up initialization is being used (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=1e-30,
+        eps_scale=1e-3,
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        betas=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        warmup_init=False,
+    ):
+        relative_step = lr is None
+        if warmup_init and not relative_step:
+            raise ValueError("warmup_init requires relative_step=True")
+
+        beta1 = (
+            None if betas is None else betas[0]
+        )  # make it compat with standard betas arg
+        defaults = dict(
+            lr=lr,
+            eps=eps,
+            eps_scale=eps_scale,
+            clip_threshold=clip_threshold,
+            decay_rate=decay_rate,
+            beta1=beta1,
+            weight_decay=weight_decay,
+            scale_parameter=scale_parameter,
+            relative_step=relative_step,
+            warmup_init=warmup_init,
+        )
+        super(Adafactor, self).__init__(params, defaults)
+
+    @staticmethod
+    def _get_lr(param_group, param_state):
+        if param_group["relative_step"]:
+            min_step = (
+                1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2
+            )
+            lr_t = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+            param_scale = 1.0
+            if param_group["scale_parameter"]:
+                param_scale = max(param_group["eps_scale"], param_state["RMS"])
+            param_group["lr"] = lr_t * param_scale
+        return param_group["lr"]
+
+    @staticmethod
+    def _get_options(param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group["beta1"] is not None
+        return factored, use_first_moment
+
+    @staticmethod
+    def _rms(tensor):
+        return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+    def _approx_sq_grad(self, exp_avg_sq_row, exp_avg_sq_col):
+        r_factor = (
+            (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True))
+            .rsqrt_()
+            .unsqueeze(-1)
+        )
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError("Adafactor does not support sparse gradients.")
+
+                state = self.state[p]
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state["exp_avg_sq_col"] = torch.zeros(
+                            grad_shape[:-2] + grad_shape[-1:]
+                        ).to(grad)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                    state["RMS"] = 0
+                else:
+                    if use_first_moment:
+                        state["exp_avg"] = state["exp_avg"].to(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+                    else:
+                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+                p_data_fp32 = p.data
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state["step"] += 1
+                state["RMS"] = self._rms(p_data_fp32)
+                lr_t = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+                update = grad ** 2 + group["eps"]
+                if factored:
+                    exp_avg_sq_row = state["exp_avg_sq_row"]
+                    exp_avg_sq_col = state["exp_avg_sq_col"]
+
+                    exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
+                    exp_avg_sq_col.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-2))
+                    # exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)  # pytorch 1.6+
+                    # exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state["exp_avg_sq"]
+
+                    exp_avg_sq.mul_(beta2t).add_(1.0 - beta2t, update)
+                    # exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)  # pytorch 1.6+
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_(
+                    (self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)
+                )
+                update.mul_(lr_t)
+
+                if use_first_moment:
+                    exp_avg = state["exp_avg"]
+                    exp_avg.mul_(group["beta1"]).add_(1 - group["beta1"], update)
+                    # exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])  # pytorch 1.6+
+                    update = exp_avg
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * lr_t, p_data_fp32)
+                    # p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr_t)  # pytorch 1.6+
+
+                p_data_fp32.add_(-update)
+
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p.data.copy_(p_data_fp32)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/adahessian.py

@@ -0,0 +1,207 @@
+""" AdaHessian Optimizer
+
+Lifted from https://github.com/davda54/ada-hessian/blob/master/ada_hessian.py
+Originally licensed MIT, Copyright 2020, David Samuel
+"""
+import torch
+
+
+class Adahessian(torch.optim.Optimizer):
+    """
+    Implements the AdaHessian algorithm from "ADAHESSIAN: An Adaptive Second OrderOptimizer for Machine Learning"
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining parameter groups
+        lr (float, optional): learning rate (default: 0.1)
+        betas ((float, float), optional): coefficients used for computing running averages of gradient and the
+            squared hessian trace (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.0)
+        hessian_power (float, optional): exponent of the hessian trace (default: 1.0)
+        update_each (int, optional): compute the hessian trace approximation only after *this* number of steps
+            (to save time) (default: 1)
+        n_samples (int, optional): how many times to sample `z` for the approximation of the hessian trace (default: 1)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=0.1,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0.0,
+        hessian_power=1.0,
+        update_each=1,
+        n_samples=1,
+        avg_conv_kernel=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= hessian_power <= 1.0:
+            raise ValueError(f"Invalid Hessian power value: {hessian_power}")
+
+        self.n_samples = n_samples
+        self.update_each = update_each
+        self.avg_conv_kernel = avg_conv_kernel
+
+        # use a separate generator that deterministically generates the same `z`s across all GPUs in case of distributed training
+        self.seed = 2147483647
+        self.generator = torch.Generator().manual_seed(self.seed)
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            hessian_power=hessian_power,
+        )
+        super(Adahessian, self).__init__(params, defaults)
+
+        for p in self.get_params():
+            p.hess = 0.0
+            self.state[p]["hessian step"] = 0
+
+    @property
+    def is_second_order(self):
+        return True
+
+    def get_params(self):
+        """
+        Gets all parameters in all param_groups with gradients
+        """
+
+        return (
+            p for group in self.param_groups for p in group["params"] if p.requires_grad
+        )
+
+    def zero_hessian(self):
+        """
+        Zeros out the accumalated hessian traces.
+        """
+
+        for p in self.get_params():
+            if (
+                not isinstance(p.hess, float)
+                and self.state[p]["hessian step"] % self.update_each == 0
+            ):
+                p.hess.zero_()
+
+    @torch.no_grad()
+    def set_hessian(self):
+        """
+        Computes the Hutchinson approximation of the hessian trace and accumulates it for each trainable parameter.
+        """
+
+        params = []
+        for p in filter(lambda p: p.grad is not None, self.get_params()):
+            if (
+                self.state[p]["hessian step"] % self.update_each == 0
+            ):  # compute the trace only each `update_each` step
+                params.append(p)
+            self.state[p]["hessian step"] += 1
+
+        if len(params) == 0:
+            return
+
+        if (
+            self.generator.device != params[0].device
+        ):  # hackish way of casting the generator to the right device
+            self.generator = torch.Generator(params[0].device).manual_seed(self.seed)
+
+        grads = [p.grad for p in params]
+
+        for i in range(self.n_samples):
+            # Rademacher distribution {-1.0, 1.0}
+            zs = [
+                torch.randint(0, 2, p.size(), generator=self.generator, device=p.device)
+                * 2.0
+                - 1.0
+                for p in params
+            ]
+            h_zs = torch.autograd.grad(
+                grads,
+                params,
+                grad_outputs=zs,
+                only_inputs=True,
+                retain_graph=i < self.n_samples - 1,
+            )
+            for h_z, z, p in zip(h_zs, zs, params):
+                p.hess += (
+                    h_z * z / self.n_samples
+                )  # approximate the expected values of z*(H@z)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """
+        Performs a single optimization step.
+        Arguments:
+            closure (callable, optional) -- a closure that reevaluates the model and returns the loss (default: None)
+        """
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        self.zero_hessian()
+        self.set_hessian()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None or p.hess is None:
+                    continue
+
+                if self.avg_conv_kernel and p.dim() == 4:
+                    p.hess = (
+                        torch.abs(p.hess)
+                        .mean(dim=[2, 3], keepdim=True)
+                        .expand_as(p.hess)
+                        .clone()
+                    )
+
+                # Perform correct stepweight decay as in AdamW
+                p.mul_(1 - group["lr"] * group["weight_decay"])
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 1:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p)
+                    # Exponential moving average of Hessian diagonal square values
+                    state["exp_hessian_diag_sq"] = torch.zeros_like(p)
+
+                exp_avg, exp_hessian_diag_sq = (
+                    state["exp_avg"],
+                    state["exp_hessian_diag_sq"],
+                )
+                beta1, beta2 = group["betas"]
+                state["step"] += 1
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(p.grad, alpha=1 - beta1)
+                exp_hessian_diag_sq.mul_(beta2).addcmul_(
+                    p.hess, p.hess, value=1 - beta2
+                )
+
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                k = group["hessian_power"]
+                denom = (
+                    (exp_hessian_diag_sq / bias_correction2)
+                    .pow_(k / 2)
+                    .add_(group["eps"])
+                )
+
+                # make update
+                step_size = group["lr"] / bias_correction1
+                p.addcdiv_(exp_avg, denom, value=-step_size)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/adamp.py

@@ -0,0 +1,133 @@
+"""
+AdamP Optimizer Implementation copied from https://github.com/clovaai/AdamP/blob/master/adamp/adamp.py
+
+Paper: `Slowing Down the Weight Norm Increase in Momentum-based Optimizers` - https://arxiv.org/abs/2006.08217
+Code: https://github.com/clovaai/AdamP
+
+Copyright (c) 2020-present NAVER Corp.
+MIT license
+"""
+
+import torch
+import torch.nn as nn
+from torch.optim.optimizer import Optimizer, required
+import math
+
+
+class AdamP(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0,
+        delta=0.1,
+        wd_ratio=0.1,
+        nesterov=False,
+    ):
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            delta=delta,
+            wd_ratio=wd_ratio,
+            nesterov=nesterov,
+        )
+        super(AdamP, self).__init__(params, defaults)
+
+    def _channel_view(self, x):
+        return x.view(x.size(0), -1)
+
+    def _layer_view(self, x):
+        return x.view(1, -1)
+
+    def _cosine_similarity(self, x, y, eps, view_func):
+        x = view_func(x)
+        y = view_func(y)
+
+        x_norm = x.norm(dim=1).add_(eps)
+        y_norm = y.norm(dim=1).add_(eps)
+        dot = (x * y).sum(dim=1)
+
+        return dot.abs() / x_norm / y_norm
+
+    def _projection(self, p, grad, perturb, delta, wd_ratio, eps):
+        wd = 1
+        expand_size = [-1] + [1] * (len(p.shape) - 1)
+        for view_func in [self._channel_view, self._layer_view]:
+
+            cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func)
+
+            if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)):
+                p_n = p.data / view_func(p.data).norm(dim=1).view(expand_size).add_(eps)
+                perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view(expand_size)
+                wd = wd_ratio
+
+                return perturb, wd
+
+        return perturb, wd
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                beta1, beta2 = group["betas"]
+                nesterov = group["nesterov"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+
+                # Adam
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                    group["eps"]
+                )
+                step_size = group["lr"] / bias_correction1
+
+                if nesterov:
+                    perturb = (beta1 * exp_avg + (1 - beta1) * grad) / denom
+                else:
+                    perturb = exp_avg / denom
+
+                # Projection
+                wd_ratio = 1
+                if len(p.shape) > 1:
+                    perturb, wd_ratio = self._projection(
+                        p,
+                        grad,
+                        perturb,
+                        group["delta"],
+                        group["wd_ratio"],
+                        group["eps"],
+                    )
+
+                # Weight decay
+                if group["weight_decay"] > 0:
+                    p.data.mul_(1 - group["lr"] * group["weight_decay"] * wd_ratio)
+
+                # Step
+                p.data.add_(-step_size, perturb)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/adamw.py

@@ -0,0 +1,131 @@
+""" AdamW Optimizer
+Impl copied from PyTorch master
+"""
+import math
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class AdamW(Optimizer):
+    r"""Implements AdamW algorithm.
+
+    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay coefficient (default: 1e-2)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=1e-2,
+        amsgrad=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad
+        )
+        super(AdamW, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdamW, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Perform stepweight decay
+                p.data.mul_(1 - group["lr"] * group["weight_decay"])
+
+                # Perform optimization step
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "Adam does not support sparse gradients, please consider SparseAdam instead"
+                    )
+                amsgrad = group["amsgrad"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state["max_exp_avg_sq"] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                        group["eps"]
+                    )
+                else:
+                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                        group["eps"]
+                    )
+
+                step_size = group["lr"] / bias_correction1
+
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/lookahead.py

@@ -0,0 +1,96 @@
+""" Lookahead Optimizer Wrapper.
+Implementation modified from: https://github.com/alphadl/lookahead.pytorch
+Paper: `Lookahead Optimizer: k steps forward, 1 step back` - https://arxiv.org/abs/1907.08610
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+from torch.optim.optimizer import Optimizer
+from collections import defaultdict
+
+
+class Lookahead(Optimizer):
+    def __init__(self, base_optimizer, alpha=0.5, k=6):
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f"Invalid slow update rate: {alpha}")
+        if not 1 <= k:
+            raise ValueError(f"Invalid lookahead steps: {k}")
+        defaults = dict(lookahead_alpha=alpha, lookahead_k=k, lookahead_step=0)
+        self.base_optimizer = base_optimizer
+        self.param_groups = self.base_optimizer.param_groups
+        self.defaults = base_optimizer.defaults
+        self.defaults.update(defaults)
+        self.state = defaultdict(dict)
+        # manually add our defaults to the param groups
+        for name, default in defaults.items():
+            for group in self.param_groups:
+                group.setdefault(name, default)
+
+    def update_slow(self, group):
+        for fast_p in group["params"]:
+            if fast_p.grad is None:
+                continue
+            param_state = self.state[fast_p]
+            if "slow_buffer" not in param_state:
+                param_state["slow_buffer"] = torch.empty_like(fast_p.data)
+                param_state["slow_buffer"].copy_(fast_p.data)
+            slow = param_state["slow_buffer"]
+            slow.add_(group["lookahead_alpha"], fast_p.data - slow)
+            fast_p.data.copy_(slow)
+
+    def sync_lookahead(self):
+        for group in self.param_groups:
+            self.update_slow(group)
+
+    def step(self, closure=None):
+        # assert id(self.param_groups) == id(self.base_optimizer.param_groups)
+        loss = self.base_optimizer.step(closure)
+        for group in self.param_groups:
+            group["lookahead_step"] += 1
+            if group["lookahead_step"] % group["lookahead_k"] == 0:
+                self.update_slow(group)
+        return loss
+
+    def state_dict(self):
+        fast_state_dict = self.base_optimizer.state_dict()
+        slow_state = {
+            (id(k) if isinstance(k, torch.Tensor) else k): v
+            for k, v in self.state.items()
+        }
+        fast_state = fast_state_dict["state"]
+        param_groups = fast_state_dict["param_groups"]
+        return {
+            "state": fast_state,
+            "slow_state": slow_state,
+            "param_groups": param_groups,
+        }
+
+    def load_state_dict(self, state_dict):
+        fast_state_dict = {
+            "state": state_dict["state"],
+            "param_groups": state_dict["param_groups"],
+        }
+        self.base_optimizer.load_state_dict(fast_state_dict)
+
+        # We want to restore the slow state, but share param_groups reference
+        # with base_optimizer. This is a bit redundant but least code
+        slow_state_new = False
+        if "slow_state" not in state_dict:
+            print("Loading state_dict from optimizer without Lookahead applied.")
+            state_dict["slow_state"] = defaultdict(dict)
+            slow_state_new = True
+        slow_state_dict = {
+            "state": state_dict["slow_state"],
+            "param_groups": state_dict[
+                "param_groups"
+            ],  # this is pointless but saves code
+        }
+        super(Lookahead, self).load_state_dict(slow_state_dict)
+        self.param_groups = (
+            self.base_optimizer.param_groups
+        )  # make both ref same container
+        if slow_state_new:
+            # reapply defaults to catch missing lookahead specific ones
+            for name, default in self.defaults.items():
+                for group in self.param_groups:
+                    group.setdefault(name, default)

Sample_Finetuning_SIIMACR/I1_classification/optim/nadam.py

@@ -0,0 +1,108 @@
+import torch
+from torch.optim import Optimizer
+
+
+class Nadam(Optimizer):
+    """Implements Nadam algorithm (a variant of Adam based on Nesterov momentum).
+
+    It has been proposed in `Incorporating Nesterov Momentum into Adam`__.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 2e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        schedule_decay (float, optional): momentum schedule decay (default: 4e-3)
+
+    __ http://cs229.stanford.edu/proj2015/054_report.pdf
+    __ http://www.cs.toronto.edu/~fritz/absps/momentum.pdf
+
+        Originally taken from: https://github.com/pytorch/pytorch/pull/1408
+        NOTE: Has potential issues but does work well on some problems.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=2e-3,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        weight_decay=0,
+        schedule_decay=4e-3,
+    ):
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            schedule_decay=schedule_decay,
+        )
+        super(Nadam, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["m_schedule"] = 1.0
+                    state["exp_avg"] = grad.new().resize_as_(grad).zero_()
+                    state["exp_avg_sq"] = grad.new().resize_as_(grad).zero_()
+
+                # Warming momentum schedule
+                m_schedule = state["m_schedule"]
+                schedule_decay = group["schedule_decay"]
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+                eps = group["eps"]
+                state["step"] += 1
+                t = state["step"]
+
+                if group["weight_decay"] != 0:
+                    grad = grad.add(group["weight_decay"], p.data)
+
+                momentum_cache_t = beta1 * (1.0 - 0.5 * (0.96 ** (t * schedule_decay)))
+                momentum_cache_t_1 = beta1 * (
+                    1.0 - 0.5 * (0.96 ** ((t + 1) * schedule_decay))
+                )
+                m_schedule_new = m_schedule * momentum_cache_t
+                m_schedule_next = m_schedule * momentum_cache_t * momentum_cache_t_1
+                state["m_schedule"] = m_schedule_new
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1.0 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1.0 - beta2, grad, grad)
+                exp_avg_sq_prime = exp_avg_sq / (1.0 - beta2 ** t)
+                denom = exp_avg_sq_prime.sqrt_().add_(eps)
+
+                p.data.addcdiv_(
+                    -group["lr"] * (1.0 - momentum_cache_t) / (1.0 - m_schedule_new),
+                    grad,
+                    denom,
+                )
+                p.data.addcdiv_(
+                    -group["lr"] * momentum_cache_t_1 / (1.0 - m_schedule_next),
+                    exp_avg,
+                    denom,
+                )
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/novograd.py

@@ -0,0 +1,90 @@
+"""NovoGrad Optimizer.
+Original impl by Masashi Kimura (Convergence Lab): https://github.com/convergence-lab/novograd
+Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks`
+    - https://arxiv.org/abs/1905.11286
+"""
+
+import torch
+from torch.optim.optimizer import Optimizer
+import math
+
+
+class NovoGrad(Optimizer):
+    def __init__(
+        self,
+        params,
+        grad_averaging=False,
+        lr=0.1,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+    ):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super(NovoGrad, self).__init__(params, defaults)
+        self._lr = lr
+        self._beta1 = betas[0]
+        self._beta2 = betas[1]
+        self._eps = eps
+        self._wd = weight_decay
+        self._grad_averaging = grad_averaging
+
+        self._momentum_initialized = False
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        if not self._momentum_initialized:
+            for group in self.param_groups:
+                for p in group["params"]:
+                    if p.grad is None:
+                        continue
+                    state = self.state[p]
+                    grad = p.grad.data
+                    if grad.is_sparse:
+                        raise RuntimeError("NovoGrad does not support sparse gradients")
+
+                    v = torch.norm(grad) ** 2
+                    m = grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
+                    state["step"] = 0
+                    state["v"] = v
+                    state["m"] = m
+                    state["grad_ema"] = None
+            self._momentum_initialized = True
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                state = self.state[p]
+                state["step"] += 1
+
+                step, v, m = state["step"], state["v"], state["m"]
+                grad_ema = state["grad_ema"]
+
+                grad = p.grad.data
+                g2 = torch.norm(grad) ** 2
+                grad_ema = (
+                    g2
+                    if grad_ema is None
+                    else grad_ema * self._beta2 + g2 * (1.0 - self._beta2)
+                )
+                grad *= 1.0 / (torch.sqrt(grad_ema) + self._eps)
+
+                if self._grad_averaging:
+                    grad *= 1.0 - self._beta1
+
+                g2 = torch.norm(grad) ** 2
+                v = self._beta2 * v + (1.0 - self._beta2) * g2
+                m = self._beta1 * m + (
+                    grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
+                )
+                bias_correction1 = 1 - self._beta1 ** step
+                bias_correction2 = 1 - self._beta2 ** step
+                step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1
+
+                state["v"], state["m"] = v, m
+                state["grad_ema"] = grad_ema
+                p.data.add_(-step_size, m)
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/nvnovograd.py

@@ -0,0 +1,132 @@
+""" Nvidia NovoGrad Optimizer.
+Original impl by Nvidia from Jasper example:
+    - https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecognition/Jasper
+Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks`
+    - https://arxiv.org/abs/1905.11286
+"""
+
+import torch
+from torch.optim.optimizer import Optimizer
+import math
+
+
+class NvNovoGrad(Optimizer):
+    """
+    Implements Novograd algorithm.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.95, 0.98))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        grad_averaging: gradient averaging
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+        grad_averaging=False,
+        amsgrad=False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            grad_averaging=grad_averaging,
+            amsgrad=amsgrad,
+        )
+
+        super(NvNovoGrad, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(NvNovoGrad, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+            and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sparse gradients are not supported.")
+                amsgrad = group["amsgrad"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state["max_exp_avg_sq"] = torch.zeros([]).to(
+                            state["exp_avg"].device
+                        )
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                norm = torch.sum(torch.pow(grad, 2))
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(1 - beta2, norm)
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group["eps"])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                grad.div_(denom)
+                if group["weight_decay"] != 0:
+                    grad.add_(group["weight_decay"], p.data)
+                if group["grad_averaging"]:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                p.data.add_(-group["lr"], exp_avg)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/optim_factory.py

@@ -0,0 +1,138 @@
+""" Optimizer Factory w/ Custom Weight Decay
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+from torch import optim as optim
+
+from .adafactor import Adafactor
+from .adahessian import Adahessian
+from .adamp import AdamP
+from .lookahead import Lookahead
+from .nadam import Nadam
+from .novograd import NovoGrad
+from .nvnovograd import NvNovoGrad
+from .radam import RAdam
+from .rmsprop_tf import RMSpropTF
+from .sgdp import SGDP
+
+try:
+    from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+
+def add_weight_decay(model, weight_decay=1e-5, skip_list=()):
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+    return [
+        {"params": no_decay, "weight_decay": 0.0},
+        {"params": decay, "weight_decay": weight_decay},
+    ]
+
+
+def create_optimizer(args, model, filter_bias_and_bn=True):
+    opt_lower = args.opt.lower()
+    weight_decay = args.weight_decay
+    if weight_decay and filter_bias_and_bn:
+        skip = {}
+        if hasattr(model, "no_weight_decay"):
+            skip = model.no_weight_decay()
+        parameters = add_weight_decay(model, weight_decay, skip)
+        weight_decay = 0.0
+    else:
+        parameters = filter(
+            lambda p: p.requires_grad, model.parameters()
+        )  # model.parameters()
+
+    if "fused" in opt_lower:
+        assert (
+            has_apex and torch.cuda.is_available()
+        ), "APEX and CUDA required for fused optimizers"
+
+    opt_args = dict(lr=args.lr, weight_decay=weight_decay)
+    if hasattr(args, "opt_eps") and args.opt_eps is not None:
+        opt_args["eps"] = args.opt_eps
+    if hasattr(args, "opt_betas") and args.opt_betas is not None:
+        opt_args["betas"] = args.opt_betas
+    if hasattr(args, "opt_args") and args.opt_args is not None:
+        opt_args.update(args.opt_args)
+
+    opt_split = opt_lower.split("_")
+    opt_lower = opt_split[-1]
+    if opt_lower == "sgd" or opt_lower == "nesterov":
+        opt_args.pop("eps", None)
+        optimizer = optim.SGD(
+            parameters, momentum=args.momentum, nesterov=True, **opt_args
+        )
+    elif opt_lower == "momentum":
+        opt_args.pop("eps", None)
+        optimizer = optim.SGD(
+            parameters, momentum=args.momentum, nesterov=False, **opt_args
+        )
+    elif opt_lower == "adam":
+        optimizer = optim.Adam(parameters, **opt_args)
+    elif opt_lower == "adamw":
+        optimizer = optim.AdamW(parameters, **opt_args)
+    elif opt_lower == "nadam":
+        optimizer = Nadam(parameters, **opt_args)
+    elif opt_lower == "radam":
+        optimizer = RAdam(parameters, **opt_args)
+    elif opt_lower == "adamp":
+        optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
+    elif opt_lower == "sgdp":
+        optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == "adadelta":
+        optimizer = optim.Adadelta(parameters, **opt_args)
+    elif opt_lower == "adafactor":
+        if not args.lr:
+            opt_args["lr"] = None
+        optimizer = Adafactor(parameters, **opt_args)
+    elif opt_lower == "adahessian":
+        optimizer = Adahessian(parameters, **opt_args)
+    elif opt_lower == "rmsprop":
+        optimizer = optim.RMSprop(
+            parameters, alpha=0.9, momentum=args.momentum, **opt_args
+        )
+    elif opt_lower == "rmsproptf":
+        optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == "novograd":
+        optimizer = NovoGrad(parameters, **opt_args)
+    elif opt_lower == "nvnovograd":
+        optimizer = NvNovoGrad(parameters, **opt_args)
+    elif opt_lower == "fusedsgd":
+        opt_args.pop("eps", None)
+        optimizer = FusedSGD(
+            parameters, momentum=args.momentum, nesterov=True, **opt_args
+        )
+    elif opt_lower == "fusedmomentum":
+        opt_args.pop("eps", None)
+        optimizer = FusedSGD(
+            parameters, momentum=args.momentum, nesterov=False, **opt_args
+        )
+    elif opt_lower == "fusedadam":
+        optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
+    elif opt_lower == "fusedadamw":
+        optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
+    elif opt_lower == "fusedlamb":
+        optimizer = FusedLAMB(parameters, **opt_args)
+    elif opt_lower == "fusednovograd":
+        opt_args.setdefault("betas", (0.95, 0.98))
+        optimizer = FusedNovoGrad(parameters, **opt_args)
+    else:
+        assert False and "Invalid optimizer"
+        raise ValueError
+
+    if len(opt_split) > 1:
+        if opt_split[0] == "lookahead":
+            optimizer = Lookahead(optimizer)
+
+    return optimizer

Sample_Finetuning_SIIMACR/I1_classification/optim/radam.py

@@ -0,0 +1,170 @@
+"""RAdam Optimizer.
+Implementation lifted from: https://github.com/LiyuanLucasLiu/RAdam
+Paper: `On the Variance of the Adaptive Learning Rate and Beyond` - https://arxiv.org/abs/1908.03265
+"""
+import math
+import torch
+from torch.optim.optimizer import Optimizer, required
+
+
+class RAdam(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.buffer = [[None, None, None] for ind in range(10)]
+        super(RAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RAdam, self).__setstate__(state)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError("RAdam does not support sparse gradients")
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
+                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
+                else:
+                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
+                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+
+                state["step"] += 1
+                buffered = self.buffer[int(state["step"] % 10)]
+                if state["step"] == buffered[0]:
+                    N_sma, step_size = buffered[1], buffered[2]
+                else:
+                    buffered[0] = state["step"]
+                    beta2_t = beta2 ** state["step"]
+                    N_sma_max = 2 / (1 - beta2) - 1
+                    N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
+                    buffered[1] = N_sma
+
+                    # more conservative since it's an approximated value
+                    if N_sma >= 5:
+                        step_size = (
+                            group["lr"]
+                            * math.sqrt(
+                                (1 - beta2_t)
+                                * (N_sma - 4)
+                                / (N_sma_max - 4)
+                                * (N_sma - 2)
+                                / N_sma
+                                * N_sma_max
+                                / (N_sma_max - 2)
+                            )
+                            / (1 - beta1 ** state["step"])
+                        )
+                    else:
+                        step_size = group["lr"] / (1 - beta1 ** state["step"])
+                    buffered[2] = step_size
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss
+
+
+class PlainRAdam(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+
+        super(PlainRAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(PlainRAdam, self).__setstate__(state)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError("RAdam does not support sparse gradients")
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
+                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
+                else:
+                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
+                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+
+                state["step"] += 1
+                beta2_t = beta2 ** state["step"]
+                N_sma_max = 2 / (1 - beta2) - 1
+                N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(-group["weight_decay"] * group["lr"], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    step_size = (
+                        group["lr"]
+                        * math.sqrt(
+                            (1 - beta2_t)
+                            * (N_sma - 4)
+                            / (N_sma_max - 4)
+                            * (N_sma - 2)
+                            / N_sma
+                            * N_sma_max
+                            / (N_sma_max - 2)
+                        )
+                        / (1 - beta1 ** state["step"])
+                    )
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    step_size = group["lr"] / (1 - beta1 ** state["step"])
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss

Sample_Finetuning_SIIMACR/I1_classification/optim/rmsprop_tf.py

@@ -0,0 +1,160 @@
+""" RMSProp modified to behave like Tensorflow impl
+
+Originally cut & paste from PyTorch RMSProp
+https://github.com/pytorch/pytorch/blob/063946d2b3f3f1e953a2a3b54e0b34f1393de295/torch/optim/rmsprop.py
+Licensed under BSD-Clause 3 (ish), https://github.com/pytorch/pytorch/blob/master/LICENSE
+
+Modifications Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch.optim import Optimizer
+
+
+class RMSpropTF(Optimizer):
+    """Implements RMSprop algorithm (TensorFlow style epsilon)
+
+    NOTE: This is a direct cut-and-paste of PyTorch RMSprop with eps applied before sqrt
+    and a few other modifications to closer match Tensorflow for matching hyper-params.
+
+    Noteworthy changes include:
+    1. Epsilon applied inside square-root
+    2. square_avg initialized to ones
+    3. LR scaling of update accumulated in momentum buffer
+
+    Proposed by G. Hinton in his
+    `course <http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_.
+
+    The centered version first appears in `Generating Sequences
+    With Recurrent Neural Networks <https://arxiv.org/pdf/1308.0850v5.pdf>`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        momentum (float, optional): momentum factor (default: 0)
+        alpha (float, optional): smoothing (decay) constant (default: 0.9)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-10)
+        centered (bool, optional) : if ``True``, compute the centered RMSProp,
+            the gradient is normalized by an estimation of its variance
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        decoupled_decay (bool, optional): decoupled weight decay as per https://arxiv.org/abs/1711.05101
+        lr_in_momentum (bool, optional): learning rate scaling is included in the momentum buffer
+            update as per defaults in Tensorflow
+
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-2,
+        alpha=0.9,
+        eps=1e-10,
+        weight_decay=0,
+        momentum=0.0,
+        centered=False,
+        decoupled_decay=False,
+        lr_in_momentum=True,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= momentum:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= alpha:
+            raise ValueError("Invalid alpha value: {}".format(alpha))
+
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            alpha=alpha,
+            eps=eps,
+            centered=centered,
+            weight_decay=weight_decay,
+            decoupled_decay=decoupled_decay,
+            lr_in_momentum=lr_in_momentum,
+        )
+        super(RMSpropTF, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RMSpropTF, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("momentum", 0)
+            group.setdefault("centered", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("RMSprop does not support sparse gradients")
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["square_avg"] = torch.ones_like(
+                        p.data
+                    )  # PyTorch inits to zero
+                    if group["momentum"] > 0:
+                        state["momentum_buffer"] = torch.zeros_like(p.data)
+                    if group["centered"]:
+                        state["grad_avg"] = torch.zeros_like(p.data)
+
+                square_avg = state["square_avg"]
+                one_minus_alpha = 1.0 - group["alpha"]
+
+                state["step"] += 1
+
+                if group["weight_decay"] != 0:
+                    if "decoupled_decay" in group and group["decoupled_decay"]:
+                        p.data.add_(-group["weight_decay"], p.data)
+                    else:
+                        grad = grad.add(group["weight_decay"], p.data)
+
+                # Tensorflow order of ops for updating squared avg
+                square_avg.add_(one_minus_alpha, grad.pow(2) - square_avg)
+                # square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)  # PyTorch original
+
+                if group["centered"]:
+                    grad_avg = state["grad_avg"]
+                    grad_avg.add_(one_minus_alpha, grad - grad_avg)
+                    # grad_avg.mul_(alpha).add_(1 - alpha, grad)  # PyTorch original
+                    avg = (
+                        square_avg.addcmul(-1, grad_avg, grad_avg)
+                        .add(group["eps"])
+                        .sqrt_()
+                    )  # eps moved in sqrt
+                else:
+                    avg = square_avg.add(group["eps"]).sqrt_()  # eps moved in sqrt
+
+                if group["momentum"] > 0:
+                    buf = state["momentum_buffer"]
+                    # Tensorflow accumulates the LR scaling in the momentum buffer
+                    if "lr_in_momentum" in group and group["lr_in_momentum"]:
+                        buf.mul_(group["momentum"]).addcdiv_(group["lr"], grad, avg)
+                        p.data.add_(-buf)
+                    else:
+                        # PyTorch scales the param update by LR
+                        buf.mul_(group["momentum"]).addcdiv_(grad, avg)
+                        p.data.add_(-group["lr"], buf)
+                else:
+                    p.data.addcdiv_(-group["lr"], grad, avg)
+
+        return loss