posture2melody.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating landmark and mel-spectrogram data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa\n",
    "import librosa.display\n",
    "import numpy as np\n",
    "import soundfile as sf\n",
    "import moviepy.editor as movp\n",
    "import os\n",
    "import warnings\n",
    "import mediapipe as mp\n",
    "import cv2\n",
    "import pandas as pd\n",
    "# mediapipe\n",
    "mp_drawing = mp.solutions.drawing_utils #visualizing poses using visual indicators\n",
    "mp_pose = mp.solutions.pose #pose estimation models (solutions)\n",
    "# extracting audio from video (ffmpeg)\n",
    "os.environ[\"IMAGEIO_FFMPEG_EXE\"] = \"/opt/homebrew/opt/ffmpeg/bin/ffmpeg\"\n",
    "\n",
    "def Euc_dist(pts, cm, img_h, img_w):\n",
    "        distances_sq = (pts - cm) ** 2\n",
    "        euc_distances = distances_sq.sum(axis = 1)\n",
    "        diag_length = np.sqrt(img_h**2 + img_w**2)\n",
    "        return euc_distances * diag_length\n",
    "\n",
    "def landmark_data_gen(vid_dir, landmark_dir):\n",
    "    # suppress warning messages\n",
    "    warnings.filterwarnings('ignore')\n",
    "\n",
    "    landmarks = [f'landmark_{i}' for i in range(0, 33)]\n",
    "    features = ['x', 'y', 'distance_to_cm']\n",
    "    comprehensive_landmarks_df_columns = pd.MultiIndex.from_product([landmarks, features], names=['landmark', 'feature'])\n",
    "    comprehensive_landmarks_df = pd.DataFrame(columns = comprehensive_landmarks_df_columns)\n",
    "\n",
    "    cap = cv2.VideoCapture(vid_dir)\n",
    "    with mp_pose.Pose(min_detection_confidence=0.5, # using the pose estimation model\n",
    "                    min_tracking_confidence=0.5) as pose:\n",
    "        while cap.isOpened():\n",
    "            ret, frame = cap.read() # frame is the actual image of each frame\n",
    "\n",
    "            if not ret: # ending the loop when the video is done playing\n",
    "                break\n",
    "            \n",
    "            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "            image.flags.writeable = False #not modifiable\n",
    "            results = pose.process(image) #making the actual detection\n",
    "            image.flags.writeable = True #now modifiable\n",
    "            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) #opencv wants image file in BGR format, we need to rerender the images\n",
    "            image_h, image_w, _ = image.shape\n",
    "            \n",
    "            if results.pose_landmarks: #sometimes cannot capture any landmarks (mediapipe requires full body view for processing)\n",
    "                landmarks = results.pose_landmarks.landmark #list of pose landmarks\n",
    "\n",
    "                # convert landmark object into dataframe\n",
    "                landmarks_df = pd.DataFrame(columns=['x', 'y'])\n",
    "                for idx, landmark in enumerate(landmarks):\n",
    "                    new_entry = np.array([landmark.x, landmark.y])\n",
    "                    landmarks_df.loc[idx,:] = new_entry\n",
    "\n",
    "                ######## For All landmarks #########\n",
    "                # for all landmark points, keep track of their relative position w.r.t to center of mass (important for posture reconstruction)\n",
    "                all_cm_x, all_cm_y = landmarks_df['x'].mean(), landmarks_df['y'].mean()\n",
    "                landmarks_df['distance_to_cm'] = Euc_dist(np.array(landmarks_df.loc[:, ['x', 'y']]), \n",
    "                                                        np.array([all_cm_x, all_cm_y]),\n",
    "                                                        image_h, image_w)\n",
    "\n",
    "                ######## Storing all important features #########\n",
    "                # x, y coordinates, distance of each landmark from the center of mass\n",
    "                new_data = pd.DataFrame(columns=comprehensive_landmarks_df_columns)\n",
    "                for landmark_idx in range(33):\n",
    "                    landmark_name = f\"landmark_{landmark_idx}\"\n",
    "                    for feature_idx in range(len(features)):\n",
    "                        feature_name = features[feature_idx]\n",
    "                        new_data[(landmark_name, feature_name)] = [landmarks_df.loc[:, feature_name].iloc[landmark_idx]]\n",
    "                comprehensive_landmarks_df = pd.concat([comprehensive_landmarks_df, new_data], ignore_index=True)\n",
    "\n",
    "                # finding and visualizing center of mass\n",
    "                cm_x, cm_y = landmarks_df['x'].mean(), landmarks_df['y'].mean()\n",
    "                cm_x_coord, cm_y_coord = int(cm_x * image_w), int(cm_y * image_h)\n",
    "                cv2.circle(img = image, \n",
    "                        center = (cm_x_coord, cm_y_coord), \n",
    "                        radius=5, \n",
    "                        color = (0, 0, 255), \n",
    "                        thickness = -1)\n",
    "\n",
    "            else:\n",
    "                pass #when landmarks are not detected\n",
    "\n",
    "            mp_drawing.draw_landmarks(image = image, \n",
    "                                    landmark_list = results.pose_landmarks, #coordinate of each landmark detected by the pose estimation model\n",
    "                                    connections = mp_pose.POSE_CONNECTIONS, #the connections of each pose landmark\n",
    "                                    landmark_drawing_spec = mp_drawing.DrawingSpec(color = (245, 117, 66), thickness=2, circle_radius=2),\n",
    "                                    connection_drawing_spec = mp_drawing.DrawingSpec(color = (245, 117, 66), thickness=2, circle_radius=2))\n",
    "\n",
    "            cv2.imshow('Webcam Feed', image)\n",
    "            \n",
    "            if cv2.waitKey(10) & 0xFF == ord('q'): #0xFF is checking which key is being pressed\n",
    "                break\n",
    "\n",
    "        np.save(landmark_dir, comprehensive_landmarks_df)\n",
    "        cap.release() #release the use of webcame\n",
    "        cv2.destroyAllWindows() #close all cv2 associated windows\n",
    "        cv2.waitKey(1) #need this line!! else the mediapipe window maynot close properly for running on local vscode (basically giving the kernel enough time to close the windows)\n",
    "\n",
    "    # Re-enable warnings\n",
    "    warnings.filterwarnings('default')\n",
    "\n",
    "def audio_data_gen(vid_dir, audio_dir, mel_dir):\n",
    "\n",
    "    video = movp.VideoFileClip(vid_dir)\n",
    "\n",
    "    # Extract the audio from the video\n",
    "    audio_path = audio_dir\n",
    "    video.audio.write_audiofile(audio_path)\n",
    "\n",
    "    # create spectrogram based on audio\n",
    "    y, sr = librosa.load(audio_path)\n",
    "    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=2048, hop_length=512).T # [mel_bins, n_segments] --> [n_segments, mel_bins]\n",
    "    np.save(mel_dir, mel_spectrogram) # saving the numpy array\n",
    "\n",
    "def data_generation(main_dir):\n",
    "    vid_folder = os.path.join(main_dir, 'video')\n",
    "\n",
    "    for vid_file in os.listdir(vid_folder):\n",
    "        if vid_file != \".DS_Store\":\n",
    "            vid_dir = os.path.join(vid_folder, vid_file)\n",
    "            vid_name = vid_file.split('.')[0]\n",
    "\n",
    "            # create landmark data\n",
    "            landmark_dir = os.path.join(main_dir, 'landmark', f'{vid_name}_landmark.npy')\n",
    "            landmark_data_gen(vid_dir, landmark_dir)\n",
    "\n",
    "            # create mel_spectrogram data\n",
    "            audio_dir = os.path.join(main_dir, 'audio', f'{vid_name}_audio.wav')\n",
    "            mel_dir = os.path.join(main_dir, 'mel_spectrogram', f'{vid_name}_mel_spectrogram.npy')\n",
    "            audio_data_gen(vid_dir, audio_dir, mel_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_generation('/Users/liuchenshu/Documents/Research/Human Posture/mediadata')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First Option: Unify the length for each batch (not across the entire dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/A good time on vacation _Alexander A ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Love you _Alexander A !_landmark.npy', \"/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Let's have a party！ #dance _Alexander A !_landmark.npy\", '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Shining in the darkness of the night_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Let me tell you a secret…_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Break through yourself ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/The Walker Out of the Night！ Alexander A !_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/scorching summer day _Alexander A ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/call me_Alexander A_landmark.npy', \"/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Dance music that's getting to the top_landmark.npy\", '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/midnight _Alexander A！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Classic reappearance_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Sparkling Snap Dance_Alexander A！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/become progressively uncontrollable_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Get in the groove and start moving  _ Alexander A ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/bad boy Alexander A_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Today you are the master of the stage!_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Bomb the street_Entertainment only_Alexander A !!_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Simple Snap Dance, for those who go above and beyond!_Alexander A ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Please come to center stage^ ^_Alexander A！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/return of the king！！_Alexander A_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Breaking the Rhythm_ Alexander A !_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/When you come to center stage_Alexander A !_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Shining in the darkness～～_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Fly with me!_Alexander  A_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Something different_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/The Black Elf  _Alexander A ！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/A memorable season_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/The macho man is coming_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/kicking back and relaxing^ ^_Alexander  A！_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Dancing lightly with everyone_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/Freedom comes from being together!_landmark.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/landmark/I Need U_landmark.npy']\n",
      "['/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/bad boy Alexander A_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Please come to center stage^ ^_Alexander A！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/When you come to center stage_Alexander A !_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/midnight _Alexander A！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Break through yourself ！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/I Need U_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Freedom comes from being together!_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/The Walker Out of the Night！ Alexander A !_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Shining in the darkness of the night_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/A memorable season_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Dancing lightly with everyone_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/A good time on vacation _Alexander A ！_mel_spectrogram.npy', \"/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Dance music that's getting to the top_mel_spectrogram.npy\", '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Love you _Alexander A !_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/The Black Elf  _Alexander A ！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/The macho man is coming_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Classic reappearance_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Today you are the master of the stage!_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Fly with me!_Alexander  A_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/return of the king！！_Alexander A_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/scorching summer day _Alexander A ！_mel_spectrogram.npy', \"/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Let's have a party！ #dance _Alexander A !_mel_spectrogram.npy\", '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Get in the groove and start moving  _ Alexander A ！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Something different_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Let me tell you a secret…_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Sparkling Snap Dance_Alexander A！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/become progressively uncontrollable_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Bomb the street_Entertainment only_Alexander A !!_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Shining in the darkness～～_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/kicking back and relaxing^ ^_Alexander  A！_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Breaking the Rhythm_ Alexander A !_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/call me_Alexander A_mel_spectrogram.npy', '/Users/liuchenshu/Documents/Research/Human Posture/mediadata/mel_spectrogram/Simple Snap Dance, for those who go above and beyond!_Alexander A ！_mel_spectrogram.npy']\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import os\n",
    "\n",
    "class posture2melDataset(Dataset):\n",
    "    def __init__(self, landmark_dir, mel_dir):\n",
    "        self.landmark_dir = landmark_dir\n",
    "        self.mel_dir = mel_dir\n",
    "    \n",
    "    def __len__(self): # return number of samples\n",
    "        return len(self.landmark_dir)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        landmark_data = np.load(self.landmark_dir[idx])\n",
    "        mel_data = np.load(self.mel_dir[idx])\n",
    "        landmark_tensor = torch.tensor(landmark_data, dtype = torch.float32)\n",
    "        mel_tensor = torch.tensor(mel_data, dtype = torch.float32)\n",
    "        return landmark_tensor, mel_tensor\n",
    "\n",
    "# for unifying sequence length in batch\n",
    "def collate_fn(batch):\n",
    "    landmark_data = [item[0] for item in batch]\n",
    "    mel_data = [item[1] for item in batch]\n",
    "    # make sure the sequences are padded according to the longest sequence in the batch\n",
    "    landmark_data_padded = pad_sequence(landmark_data, batch_first=True) # [batch_size, max_seq_length, n_features]\n",
    "    mel_data_padded = pad_sequence(mel_data, batch_first=True)\n",
    "\n",
    "    # for ignoring padding when feeding into transformer (i.d.ing where the paddings are)\n",
    "    landmark_mask = (landmark_data_padded == 0)\n",
    "    mel_mask = (mel_data_padded == 0)\n",
    "\n",
    "    return landmark_data_padded, mel_data_padded, landmark_mask, mel_mask\n",
    "\n",
    "main_dir = '/Users/liuchenshu/Documents/Research/Human Posture/mediadata'  \n",
    "landmark_paths = [os.path.join(main_dir, 'landmark', landmark_name) for landmark_name in os.listdir(os.path.join(main_dir, 'landmark')) if landmark_name != '.DS_Store']\n",
    "print(landmark_paths)\n",
    "mel_paths = [os.path.join(main_dir, 'mel_spectrogram', landmark_name) for landmark_name in os.listdir(os.path.join(main_dir, 'mel_spectrogram')) if landmark_name != '.DS_Store']\n",
    "print(mel_paths)\n",
    "\n",
    "dataset = posture2melDataset(landmark_paths, mel_paths)\n",
    "data_loader = DataLoader(dataset, batch_size = 4, shuffle = True, collate_fn=collate_fn) #collate_fn make sure the window length is same in each batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for landmark, mel_spectrogram, landmark_mask, mel_mask in data_loader:\n",
    "    print(landmark.shape)\n",
    "    print(mel_spectrogram.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Second Option: Unify length according to longest sequence across the entire dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch 0\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 1\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 2\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 3\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 4\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 5\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 6\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 7\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 8\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 9\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 10\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 11\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 12\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 13\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 14\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 15\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 16\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 17\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 18\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 19\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 20\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 21\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 22\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 23\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 24\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 25\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 26\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 27\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 28\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 29\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 30\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 31\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n",
      "batch 32\n",
      "torch.Size([1, 1867, 99]) torch.Size([1, 1352, 128]) torch.Size([1, 1867]) torch.Size([1, 1352])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "class PaddedDataset(Dataset): # unifying across the entire dataset\n",
    "    def __init__(self, input_file_dir, output_file_dir):\n",
    "        # Load the input and output datasets\n",
    "        self.input_file_dir = input_file_dir\n",
    "        self.output_file_dir = output_file_dir\n",
    "        self.inputs = []\n",
    "        for input_dir in self.input_file_dir:\n",
    "            self.inputs.append(np.load(input_dir))\n",
    "        self.outputs = []\n",
    "        for output_dir in self.output_file_dir:\n",
    "            self.outputs.append(np.load(output_dir))\n",
    "\n",
    "        # Calculate the maximum length across the whole dataset for inputs and outputs (NUCLEAR option!!!)\n",
    "        self.max_input_len = max(inp.shape[0] for inp in self.inputs)\n",
    "        self.max_output_len = max(out.shape[0] for out in self.outputs)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_file_dir)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        input_seq = torch.tensor(self.inputs[idx], dtype=torch.float32)\n",
    "        output_seq = torch.tensor(self.outputs[idx], dtype=torch.float32)\n",
    "\n",
    "        # Pad the input sequence to max_input_len\n",
    "        padded_input = torch.cat([input_seq, torch.zeros(self.max_input_len - input_seq.size(0), input_seq.size(1))], dim=0)\n",
    "\n",
    "        # Create input mask\n",
    "        input_mask = (padded_input.sum(dim=1) != 0)  # Shape: (max_input_len, )\n",
    "\n",
    "        # Pad the output sequence to max_output_len\n",
    "        padded_output = torch.cat([output_seq, torch.zeros(self.max_output_len - output_seq.size(0), output_seq.size(1))], dim=0)\n",
    "\n",
    "        # Create output mask\n",
    "        output_mask = (padded_output.sum(dim=1) != 0)  # Shape: (max_output_len, )\n",
    "\n",
    "        return {\n",
    "            'input': padded_input,\n",
    "            'target': padded_output,\n",
    "            'input_mask': input_mask,\n",
    "            'target_mask': output_mask\n",
    "        }\n",
    "\n",
    "main_dir = '/Users/liuchenshu/Documents/Research/Posture2Melody/mediadata'  \n",
    "landmark_paths = [os.path.join(main_dir, 'landmark', landmark_name) for landmark_name in os.listdir(os.path.join(main_dir, 'landmark')) if landmark_name != '.DS_Store']\n",
    "mel_paths = [os.path.join(main_dir, 'mel_spectrogram', landmark_name) for landmark_name in os.listdir(os.path.join(main_dir, 'mel_spectrogram')) if landmark_name != '.DS_Store']\n",
    "dataset = PaddedDataset(landmark_paths, mel_paths)\n",
    "\n",
    "# Create the DataLoader\n",
    "dataloader = DataLoader(dataset, batch_size=1)\n",
    "\n",
    "# Example of iterating over the DataLoader\n",
    "for idx, batch in enumerate(dataloader):\n",
    "    print(f\"batch {idx}\")\n",
    "    inputs = batch['input']  # Shape: (batch_size, max_input_len, input_dim)\n",
    "    targets = batch['target']  # Shape: (batch_size, max_output_len, output_dim)\n",
    "    input_masks = batch['input_mask']  # Shape: (batch_size, max_input_len, 1)\n",
    "    target_masks = batch['target_mask']  # Shape: (batch_size, max_output_len, 1)\n",
    "    print(inputs.shape, targets.shape, input_masks.shape, target_masks.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer + GAN Mel-spectrogram reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, hidden_dim, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        # Create constant 'pe' matrix with values dependent on position and i\n",
    "        pe = torch.zeros(max_len, hidden_dim)  # Initialize a positional encoding matrix of shape (max_len, hidden_dim)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # A column vector [0, 1, ..., max_len-1]\n",
    "        div_term = torch.exp(torch.arange(0, hidden_dim, 2).float() * (-math.log(10000.0) / hidden_dim))\n",
    "        \n",
    "        # Apply sine to even indices in the encoding (2i)\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        \n",
    "        # Apply cosine to odd indices in the encoding (2i+1)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        \n",
    "        # Add a batch dimension at the beginning of the positional encoding matrix\n",
    "        pe = pe.unsqueeze(0)\n",
    "        \n",
    "        # Store the positional encodings in the buffer (they are constants, not learnable parameters)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Add positional encoding to the input embeddings (x)\n",
    "        # x is expected to be of shape (batch_size, sequence_length, hidden_dim)\n",
    "        x = x + self.pe[:, :x.size(1), :]  # Add positional encodings to the input embeddings\n",
    "        return x\n",
    "    \n",
    "def generate_square_subsequent_mask(size):\n",
    "    \"\"\"\n",
    "    Generate a mask for self-attention in the decoder.\n",
    "    Creates a triangular matrix with the segment on top of diagnal is true, below is false\n",
    "    \"\"\"\n",
    "    mask = torch.triu(torch.ones(size, size), diagonal=1).bool()  # Upper triangular mask\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, hidden_dim, nhead, num_layers, n_features, max_len = 2000):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.embedding = nn.Linear(n_features, hidden_dim)  # Embedding layer\n",
    "        self.pos_encoder = PositionalEncoding(hidden_dim, max_len)\n",
    "        self.transformer_encoder = nn.TransformerEncoder(\n",
    "            nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead), \n",
    "            num_layers=num_layers\n",
    "        )\n",
    "\n",
    "    def forward(self, src, src_key_padding_mask=None):\n",
    "        src_emb = self.embedding(src)  # Shape: (batch_size, n_frame1, hidden_dim)\n",
    "        src_emb = self.pos_encoder(src_emb) # adding the position encoding to the semantic encoding\n",
    "        #print(f\"src_emb is {src_emb}\")\n",
    "        src_emb = src_emb.transpose(0, 1)  # Transpose to (n_frame1, batch_size, hidden_dim)\n",
    "        #print(f\"src_emb is {src_emb}\")\n",
    "        memory = self.transformer_encoder(src_emb, src_key_padding_mask = src_key_padding_mask)\n",
    "        #print(f\"memory within encoder is {memory}\")\n",
    "        return memory  # Transpose back to (n_frame1, batch_size, hidden_dim)\n",
    "\n",
    "class TransformerDecoder(nn.Module):\n",
    "    def __init__(self, hidden_dim, nhead, num_layers, mel_dim, max_len = 2000):\n",
    "        super(TransformerDecoder, self).__init__()\n",
    "        self.embedding = nn.Linear(mel_dim, hidden_dim)  # Embedding layer\n",
    "        self.pos_encoder = PositionalEncoding(hidden_dim, max_len)\n",
    "        self.transformer_decoder = nn.TransformerDecoder(\n",
    "            nn.TransformerDecoderLayer(d_model = hidden_dim, nhead = nhead),\n",
    "            num_layers = num_layers\n",
    "        )\n",
    "        self.fc = nn.Linear(hidden_dim, mel_dim)\n",
    "\n",
    "    def forward(self, tgt, memory, tgt_key_padding_mask = None, memory_key_padding_mask = None):\n",
    "        tgt_emb = self.embedding(tgt)  # Shape: (batch_size, n_frame2, hidden_dim)\n",
    "        tgt_emb = self.pos_encoder(tgt_emb)\n",
    "        tgt_emb = tgt_emb.transpose(0, 1)  # Transpose to (n_frame2, batch_size, hidden_dim)\n",
    "        print(f\"tgt_emb in decoder is {tgt_emb}\")\n",
    "        print(f\"Memory shape from within decoder function: {memory.shape}\")  # Should be (n_frame1, batch_size, hidden_dim)\n",
    "        print(f\"Tgt Embedding shape from within decoder function: {tgt_emb.shape}\")  # Should be (n_frame2, batch_size, hidden_dim)\n",
    "        tgt_len = tgt_emb.size(0)\n",
    "        subsequent_mask = generate_square_subsequent_mask(tgt_len)\n",
    "        print(torch.isnan(tgt_emb).any(), torch.isnan(memory).any())\n",
    "        output = self.transformer_decoder(tgt_emb, memory, # both need to have structure [n_frames1/2, batch_size, hidden_dim]\n",
    "                                          tgt_key_padding_mask = tgt_key_padding_mask, \n",
    "                                          memory_key_padding_mask=memory_key_padding_mask,\n",
    "                                          tgt_mask = subsequent_mask)\n",
    "        # print(f\"Output dimension from decoder is {output.shape}\")\n",
    "        print(f\"first output in decoder {output}\")\n",
    "        output = self.fc(output)\n",
    "        print(f\"output after fc is {output}\")\n",
    "        #print(output.shape)\n",
    "        return output.transpose(0, 1)  # Transpose back to (batch_size, n_frame2, hidden_dim)\n",
    "\n",
    "class GANDiscriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "device = torch.device(\"mps\") if torch.backends.mps.is_available() else torch.device(\"cpu\")\n",
    "device = \"cpu\"\n",
    "# Define model parameters\n",
    "hidden_dim = 256  # Example dimension, adjust as needed\n",
    "nhead = 8  # Number of attention heads\n",
    "num_layers = 4  # Number of transformer layers\n",
    "mel_dim = 128\n",
    "n_features = 99\n",
    "\n",
    "# Instantiate model\n",
    "encoder = TransformerEncoder(hidden_dim, nhead, num_layers, n_features).to(device)\n",
    "decoder = TransformerDecoder(hidden_dim, nhead, num_layers, mel_dim).to(device)\n",
    "transformer_model = nn.Sequential(encoder, decoder).to(device)  # Or define a custom model class that combines them\n",
    "\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.MSELoss()  # Assuming you're predicting continuous values (modify as needed)\n",
    "optimizer = optim.Adam(transformer_model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.mps.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input mask shapetorch.Size([1, 1867])\n",
      "memory is tensor([[[-0.7063, -2.1459,  0.5156,  ...,  0.0467,  1.5292,  1.7966]],\n",
      "\n",
      "        [[-0.2906, -2.7098,  0.0094,  ..., -0.0948,  1.2681,  1.2571]],\n",
      "\n",
      "        [[-1.1288, -3.1621,  0.5975,  ..., -0.1013,  1.5313,  1.4220]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[-1.5757, -2.0522,  1.1208,  ...,  1.0341,  0.5154,  1.8668]],\n",
      "\n",
      "        [[-1.1248, -0.9563,  2.3334,  ...,  1.3376,  0.3556,  1.7938]],\n",
      "\n",
      "        [[-0.4516, -0.7848,  2.1915,  ...,  0.9469,  0.2729,  1.4137]]],\n",
      "       grad_fn=<NativeLayerNormBackward0>)\n",
      "tgt_emb in decoder is tensor([[[-1.9967e-02,  1.0822e+00,  7.1077e-02,  ...,  9.2854e-01,\n",
      "          -7.6038e-03,  1.0490e+00]],\n",
      "\n",
      "        [[ 4.6548e-01,  8.1928e-01,  3.5497e-01,  ..., -2.4307e-01,\n",
      "           1.5783e+00,  1.3891e+00]],\n",
      "\n",
      "        [[-2.3020e+00, -1.2746e-01, -2.8987e+00,  ..., -3.6861e+00,\n",
      "           1.0426e+01,  1.4453e+00]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[-9.7114e-01, -2.2646e-01, -8.9270e-01,  ...,  9.1877e-01,\n",
      "           1.3496e-01,  1.0380e+00]],\n",
      "\n",
      "        [[-7.9386e-01,  7.1586e-01, -2.8719e-01,  ...,  9.1875e-01,\n",
      "           1.3506e-01,  1.0380e+00]],\n",
      "\n",
      "        [[ 9.4858e-02,  1.0758e+00,  6.0561e-01,  ...,  9.1873e-01,\n",
      "           1.3517e-01,  1.0379e+00]]], grad_fn=<TransposeBackward0>)\n",
      "Memory shape from within decoder function: torch.Size([1867, 1, 256])\n",
      "Tgt Embedding shape from within decoder function: torch.Size([1352, 1, 256])\n",
      "tensor(False) tensor(False)\n",
      "first output in decoder tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]],\n",
      "       grad_fn=<NativeLayerNormBackward0>)\n",
      "output after fc is tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]], grad_fn=<ViewBackward0>)\n",
      "output is tensor([[[nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         ...,\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan]]], grad_fn=<TransposeBackward0>)\n",
      "shape of output torch.Size([1, 1352, 128])\n",
      "shape of target_seq torch.Size([1, 1352, 128])\n",
      "loss is nan\n",
      "tensor(False)\n",
      "Epoch [1/10], Loss: nan\n",
      "input mask shapetorch.Size([1, 1867])\n",
      "memory is tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]],\n",
      "       grad_fn=<NativeLayerNormBackward0>)\n",
      "tgt_emb in decoder is tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]], grad_fn=<TransposeBackward0>)\n",
      "Memory shape from within decoder function: torch.Size([1867, 1, 256])\n",
      "Tgt Embedding shape from within decoder function: torch.Size([1352, 1, 256])\n",
      "tensor(True) tensor(True)\n",
      "first output in decoder tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]],\n",
      "       grad_fn=<NativeLayerNormBackward0>)\n",
      "output after fc is tensor([[[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]],\n",
      "\n",
      "        [[nan, nan, nan,  ..., nan, nan, nan]]], grad_fn=<ViewBackward0>)\n",
      "output is tensor([[[nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         ...,\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan],\n",
      "         [nan, nan, nan,  ..., nan, nan, nan]]], grad_fn=<TransposeBackward0>)\n",
      "shape of output torch.Size([1, 1352, 128])\n",
      "shape of target_seq torch.Size([1, 1352, 128])\n",
      "loss is nan\n",
      "tensor(False)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[105], line 42\u001b[0m\n\u001b[1;32m     39\u001b[0m torch\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mutils\u001b[38;5;241m.\u001b[39mclip_grad_norm_(transformer_model\u001b[38;5;241m.\u001b[39mparameters(), max_norm\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1.0\u001b[39m)\n\u001b[1;32m     41\u001b[0m \u001b[38;5;66;03m# Backward pass and optimization\u001b[39;00m\n\u001b[0;32m---> 42\u001b[0m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     43\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m     44\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "File \u001b[0;32m~/Documents/Research/Posture2Melody/humanposturn_venv/lib/python3.12/site-packages/torch/_tensor.py:521\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    511\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    512\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    513\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    514\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    519\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    520\u001b[0m     )\n\u001b[0;32m--> 521\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    522\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    523\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/Research/Posture2Melody/humanposturn_venv/lib/python3.12/site-packages/torch/autograd/__init__.py:289\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    284\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    286\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    287\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    288\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 289\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    290\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    291\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    292\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    293\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    294\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    295\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    296\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    297\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/Research/Posture2Melody/humanposturn_venv/lib/python3.12/site-packages/torch/autograd/graph.py:769\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    767\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[1;32m    768\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 769\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    770\u001b[0m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    771\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    772\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    773\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "num_epochs = 10  # Define number of epochs\n",
    "transformer_model.train()  # Set model to training mode\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss = 0  # Variable to accumulate loss\n",
    "\n",
    "    for batch in dataloader:\n",
    "        # Get padded inputs and targets\n",
    "        input_seq = batch['input'].to(device)  # Shape: (batch_size, max_input_len, input_dim)\n",
    "        target_seq = batch['target'].to(device)  # Shape: (batch_size, max_output_len, output_dim)\n",
    "        input_mask = batch['input_mask'].to(device)  # Input mask\n",
    "        target_mask = batch['target_mask'].to(device)  # Output mask\n",
    "        print(f\"input mask shape{input_mask.shape}\")\n",
    "\n",
    "        # Move tensors to the appropriate device (CPU/GPU)\n",
    "        input_seq = input_seq.to(device)\n",
    "        target_seq = target_seq.to(device)\n",
    "\n",
    "        # Zero the gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Forward pass through the model\n",
    "        memory = encoder(input_seq, src_key_padding_mask = input_mask)  # Encoder output\n",
    "        print(f\"memory is {memory}\")\n",
    "        # print(f\"memory shape {memory.shape}\")\n",
    "        # print(f\"target shape {target_seq.shape}\")\n",
    "        output = decoder(target_seq, memory, tgt_key_padding_mask = target_mask, memory_key_padding_mask = input_mask)  # Decoder output\n",
    "        print(f\"output is {output}\")\n",
    "        print(f\"shape of output {output.shape}\")\n",
    "        print(f\"shape of target_seq {target_seq.shape}\")\n",
    "\n",
    "        # Calculate loss\n",
    "        loss = criterion(output, target_seq)  # Modify if needed based on your output processing\n",
    "        print(f\"loss is {loss}\")\n",
    "        print(torch.isinf(loss))\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # Clip gradients\n",
    "        torch.nn.utils.clip_grad_norm_(transformer_model.parameters(), max_norm=1.0)\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        break\n",
    "\n",
    "    # Print average loss for the epoch\n",
    "    avg_loss = total_loss / len(dataloader)\n",
    "    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "humanposturn_venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}