Ziyu Li's Homepage

Agent Memory Research In MMLab (Part 4)

Thu, 06 Aug 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

文档说明

本文档记载一次MLE的预训练和极小的微调

实验数据

从生成的长对话里面提取出来的树，总共100颗，88样本内，12val，树比较大，大概每棵树有150个fact+50个aspect，大概共200个节点

实验过程

首先是用chatgpt-4o做标注，也就是对每个snapshot的closure，让gpt去决定应该做什么动作，这就得到了所有的(X,y)，做监督学习

这些(X,y)做100个epoch的监督学习

结束之后，做online学习，但只做两个policy，online的成本非常高，2个epoch大概跑了10小时

评判结果

从结构上来说，teacher(gpt)和student(小模型)做出来的树的aspect都比较少，至少比数据集里面的原树少得多

用gpt5.6-sol当强模型，直接分别给ref, teacher, student树打分如下(理论上student的分不可能超过teacher)

盘踞大概是这样：假如一棵树有100个fact，LLM去逐个判断每个fact和他父亲的关系的合理性，得到一个百分比，比如说50个fact是合理的，那这个值就是0.5

接下来再做很多次(比如说10次)，这10次就把fact随机乱放，然后每次能得到一个百分比，这个百分比大概率很低，因为乱放可能只有5个fact合理，平均可能也是5%左右

然后再考虑超额，超额=树的合理百分比-随机乱放的平均百分比即可

当然了，因为随便乱放的baseline大概率趋于0，也可以不剪掉，直接看合理程度就行了

从这个图上可以看出来，teacher的质量其实也一般，或者说teacher在看不到ref_tree的情况下，应该不太可能做到比较好

生产结果

我们用一棵树，分别展示他的ref/teacher/student

# ref
# REFERENCE tree task_id=46 entity='Company: Voyage Vista Travel'
# facts=122 aspects=63 max_depth=6 avg_branching=2.92
# policy rollout used the first max_facts=0 facts of this tree (BFS order)

[A] Voyage Vista Travel
  [A] Corporate Programs & Partnerships
    [A] Corporate Retreats Division
      [A] Program Development & Strategy
        [F:0] Voyage Vista Travel is brainstorming proposals for corporate team-building retreats.
        [F:1] Voyage Vista Travel plans to present proposals for corporate team-building retreats to local tech co
      [A] Synaptic Solutions Retreat (Costa Rica)
        [F:2] Voyage Vista Travel organized an adventure retreat for Synaptic Solutions.
        [F:3] The adventure retreat organized by Voyage Vista Travel was for 40 employees of Synaptic Solutions.
        [F:4] The adventure retreat organized by Voyage Vista Travel is a 5-day multi-activity trip.
        [F:5] The adventure retreat organized by Voyage Vista Travel is to the Costa Rican rainforest.
        [F:6] The adventure retreat organized by Voyage Vista Travel is scheduled for August 2023.
        [F:7] The adventure retreat organized by Voyage Vista Travel is worth over $75,000.
      [A] Syntel Systems Retreat (Redwood NP)
        [A] Program Details & Scope
          [F:8] Voyage Vista landed a massive corporate retreat program with Syntel Systems.
          [F:9] The corporate retreat program for Syntel Systems is for 50 employees.
          [F:10] The corporate retreat program for Syntel Systems is focused on team-building and leadership developm
          [F:11] The corporate retreat program for Syntel Systems is a four-day, three-night program.
          [F:12] The corporate retreat program for Syntel Systems is planned for Redwood National Park.
          [F:13] The corporate retreat program for Syntel Systems is scheduled for August 15th to August 18th.
          [F:14] The corporate retreat program for Syntel Systems includes guided hikes.
          [F:15] The corporate retreat program for Syntel Systems includes a conservation project.
        [A] Business Impact
          [F:16] The contract with Syntel Systems is a significant revenue boost for Voyage Vista.
          [F:17] The contract with Syntel Systems positions Voyage Vista for more corporate work.
      [A] TechSolutions Inc. Retreat (Whistler)
        [A] Proposal Phase
          [A] Initial Scope & Activities
            [F:18] Voyage Vista Travel landed its first big proposal for corporate team-building retreats.
            [F:19] Voyage Vista Travel's first big proposal is with TechSolutions Inc.
            [F:20] Voyage Vista Travel is proposing a week-long leadership and team-building retreat for TechSolutions 
            [F:21] Voyage Vista Travel is planning guided backcountry skiing for the TechSolutions Inc. retreat.
            [F:22] Voyage Vista Travel is planning snowshoeing for the TechSolutions Inc. retreat.
            [F:23] Voyage Vista Travel is planning a 'build your own igloo' survival challenge for the TechSolutions In
          [A] Proposal Timeline
            [F:24] Voyage Vista Travel is working on the final proposal for TechSolutions Inc. during the week of Novem
            [F:25] Voyage Vista Travel is aiming for a client meeting with TechSolutions Inc. on November 17, 2023.
        [A] Execution & Logistics
          [A] Contract & Financials
            [F:26] Voyage Vista Travel landed the TechSolutions Inc. contract for their corporate retreat in Whistler.
            [F:27] The meeting for the TechSolutions Inc. contract happened on November 17, 2023.
            [F:28] Voyage Vista Travel received an initial deposit of $45,000 for the TechSolutions Inc. corporate retr
            [F:29] The total project value for the TechSolutions Inc. corporate retreat is estimated at $120,000.
          [A] Vendor Management & Challenges
            [F:30] The TechSolutions Inc. Whistler trip is scheduled for January 15, 2024.
            [F:31] Voyage Vista Travel hit a minor snag with ski and snowboard rentals for the TechSolutions Inc. Whist
            [F:32] Voyage Vista Travel's original vendor for ski and snowboard rentals was Alpine Peak Sports.
            [F:33] Alpine Peak Sports could not provide 45 sets of premium gear for the TechSolutions Inc. Whistler tri
            [F:34] Voyage Vista Travel secured a bulk deal with Whistler Blackcomb Gear Rentals for the ski and snowboa
            [F:35] The deal with Whistler Blackcomb Gear Rentals was at a 5% higher cost than the original vendor.
            [F:36] Voyage Vista Travel finalized Whistler Blackcomb Gear Rentals on December 28, 2023.
    [A] Strategic Partnerships
      [A] General Corporate Deals
        [F:37] Voyage Vista Travel closed a significant deal.
        [F:38] The deal confirms Voyage Vista Travel's niche in high-end corporate adventure travel.
      [A] EcoTrek Adventures Partnership (Sustainable Safaris)
        [A] Partnership Overview & Goals
          [F:39] The 'Sustainable Safaris' partnership with EcoTrek Adventures closed on January 31, 2024.
          [F:40] The 'Sustainable Safaris' partnership with EcoTrek Adventures met its $150,000 first-year goal.
        [A] Peruvian Highlands Trek Program
          [A] Program Schedule & Bookings
            [F:41] The first 'Peruvian Highlands Trek' group for Voyage Vista Travel's EcoTrek Adventures partnership i
            [F:42] Voyage Vista Travel has secured 8 bookings out of 12 available spots for the first 'Peruvian Highlan
            [F:43] The 'Peruvian Highlands Trek' group departs on April 10, 2024.
            [F:44] The 'Peruvian Highlands Trek' group's departure is scheduled for April 10, 2024.
          [A] Marketing & Preparation
            [F:45] The promotional materials for the 'Peruvian Highlands Trek' need to be ready by February 21, 2024.
  [A] Client Acquisition & Key Accounts
    [A] Sales & Lead Generation
      [A] Silicon Valley Outreach Events
        [F:46] Voyage Vista Travel hosted a successful tasting event on June 01, 2023.
        [F:47] The tasting event was for ten potential clients from various tech startups in Silicon Valley.
        [F:48] The tasting event featured ceviche and Pisco Sours.
      [A] Prospective Clients & Pipeline
        [F:49] Voyage Vista Travel secured a follow-up meeting with Zenith Innovations.
        [F:50] The follow-up meeting with Zenith Innovations is for a team-building retreat.
        [F:51] Securing Salesforce as a client would be huge for Voyage Vista Travel.
        [F:52] Securing Google as a client would be huge for Voyage Vista Travel.
    [A] Major Client Engagements
      [A] The Sterling Group (Patagonia Trek)
        [F:53] Liam O'Connell landed The Sterling Group as a major new client for Voyage Vista Travel.
        [F:54] Voyage Vista Travel is developing a 14-day Andes Mountain Trek in Patagonia for The Sterling Group's
        [F:55] The Andes Mountain Trek for The Sterling Group is scheduled for May 2024.
      [A] Chen Family (Kenya Safari)
        [F:56] Voyage Vista Travel finalized a major booking with the Chen Family.
        [F:57] The Chen Family is looking for a 3-week luxury safari in Kenya for July 2024.
        [F:58] The safari for the Chen Family is a high-end, bespoke itinerary.
        [F:59] The safari for the Chen Family includes stays at Cottar's 1920s Camp.
        [F:60] The safari for the Chen Family includes stays at Angama Mara.
        [F:61] The safari for the Chen Family focuses on ethical wildlife viewing.
        [F:62] The safari for the Chen Family focuses on immersive cultural experiences.
        [F:63] The budget for the Chen Family's safari is approaching $75,000.
  [A] Destination & Product Development
    [A] Costa Rica Programs
      [A] Eco-Lodge Partnerships
        [A] Finca Verde (Prospective)
          [F:64] Voyage Vista Travel is looking at a partnership with a luxury eco-lodge in Costa Rica.
          [F:65] Voyage Vista Travel hopes to finalize a deal with Finca Verde for a launch by early 2024.
          [F:66] The partnership Voyage Vista Travel is looking at is for exclusive travel packages.
        [A] Casa Verde Eco-Resort (Finalized)
          [F:67] Voyage Vista finalized its partnership with Casa Verde Eco-Resort.
          [F:68] The partnership with Casa Verde Eco-Resort took months to finalize.
      [A] New Adventure Package Development
        [A] Package Design & Scope
          [F:69] Voyage Vista Travel plans to offer unique week-long excursions through the partnership with Finca Ve
          [F:70] Voyage Vista developed five exclusive new adventure packages.
        [A] Activity Inclusions
          [F:71] The new adventure packages include volcano trekking around Arenal.
          [F:72] The new adventure packages include white-water rafting on the Pacuare River.
          [F:73] The new adventure packages include cloud forest canopy tours.
          [F:74] The new adventure packages include sustainable coffee farm visits.
          [F:75] The new adventure packages include indigenous Bribri culture immersion.
      [A] Package Launch & Marketing
        [F:76] Voyage Vista is launching the new Costa Rica packages on its website on June 26, 2023.
        [F:77] Voyage Vista is doing a targeted social media campaign for the new Costa Rica packages.
    [A] Future & Diverse Destinations
      [F:78] Voyage Vista Travel has upcoming European tours planned for summer 2023.
      [F:79] Voyage Vista Travel is planning potential new South American itineraries.
  [A] Marketing & Brand Strategy
    [A] Digital Presence & Website
      [A] Website Redesign Project
        [A] Project Scope & Team
          [F:80] Liam O'Connell started working on a complete website redesign for Voyage Vista.
          [F:81] The Voyage Vista website redesign needs to better reflect expanded offerings.
          [F:82] The Voyage Vista website redesign needs to better reflect their new target audience.
          [F:83] Liam O'Connell brought in a freelance designer named Maria for the Voyage Vista website redesign.
        [A] Design & Launch Milestones
          [F:84] Maria specializes in travel industry sites.
          [F:85] Liam O'Connell is aiming for a beta launch of the Voyage Vista website redesign by mid-August 2023.
    [A] Content & Campaign Development
      [A] Sustainable Travel Guide
        [F:86] Liam O'Connell finished developing a new 'Sustainable Travel Guide' for Voyage Vista's clients.
        [F:87] The 'Sustainable Travel Guide' focuses on immersive local experiences.
        [F:88] The 'Sustainable Travel Guide' focuses on minimizing environmental impact.
        [F:89] The 'Sustainable Travel Guide' highlights destinations like Costa Rica's cloud forests.
        [F:90] The 'Sustainable Travel Guide' highlights destinations like Patagonia's conservation efforts.
      [A] Digital Marketing Campaigns
        [A] Instagram Campaign: Hidden Gems
          [F:91] Liam O'Connell is planning a new Instagram campaign for Voyage Vista called 'Hidden Gems of Costa Ri
          [F:92] The 'Hidden Gems of Costa Rica' campaign is planned for a Q1 2024 launch.
        [A] Luxury Eco-Tourism Campaign
          [F:93] Liam O'Connell and David Miller finalized a new marketing campaign for Voyage Vista Travel.
          [F:94] The new marketing campaign for Voyage Vista Travel targets luxury eco-tourism packages for the sprin
          [F:95] The new marketing campaign for Voyage Vista Travel is launching in early January 2024.
          [F:96] Voyage Vista Travel allocated $15,000 for digital ads for the new marketing campaign.
          [F:97] Voyage Vista Travel is launching a new digital marketing campaign for luxury eco-tourism on January 
          [F:98] Voyage Vista Travel has a digital marketing launch scheduled for January 8, 2024.
    [A] Public Relations & Media
      [F:99] Voyage Vista Travel received a mention on the 'Bay Area Wanderlust' travel blog.
      [F:100] The mention on the 'Bay Area Wanderlust' travel blog drove a lot of traffic to Voyage Vista Travel's
  [A] Company Operations & Development
    [A] General Company Status
      [F:101] Things at Voyage Vista Travel are good.
    [A] Office & Infrastructure
      [A] New SOMA Office
        [A] Space Requirements
          [F:102] The new office space for Voyage Vista is intended to be in the SOMA district.
          [F:103] The new office space for Voyage Vista is hoped to be around 1500 sq ft.
          [F:104] The new office space for Voyage Vista is hoped to accommodate 8-10 people.
        [A] Lease & Occupancy Details
          [F:105] Liam O'Connell signed the lease for Voyage Vista Travel's 1600 sq ft SOMA office on Mission Street o
          [F:106] Voyage Vista Travel will get the keys to its new SOMA office on Mission Street on September 1, 2023.
    [A] Team & Personnel
      [A] Sarah Chen (Client Operations)
        [F:107] Sarah Chen will start working at Voyage Vista Travel on October 2, 2023.
        [F:108] Sarah Chen streamlined Voyage Vista Travel's client intake forms within her first 9 days of work.
        [F:109] Sarah Chen reorganized Voyage Vista Travel's client database within her first 9 days of work.
      [A] Chloe Davies (Consulting)
        [F:110] Chloe Davies is interested in helping with Voyage Vista Travel's social media strategy.
        [F:111] Chloe Davies might start as a consultant for Voyage Vista Travel by March 15, 2024.
    [A] Key Company Events
      [A] Grand Opening Event
        [A] Planning & Invitations
          [F:112] Voyage Vista Travel's grand opening is planned for October 12, 2023.
          [F:113] Voyage Vista Travel has sent out about 75 invitations for its grand opening.
          [F:114] The Voyage Vista Travel grand opening is scheduled for October 12, 2023.
          [F:115] Catering by Celeste confirmed their setup for 75 guests for the Voyage Vista Travel grand opening.
        [A] Event Execution & Impact
          [F:116] The Voyage Vista Travel grand opening was on October 12, 2023.
          [F:117] The Voyage Vista Travel grand opening was a huge success.
          [F:118] Over 120 people attended the Voyage Vista Travel grand opening throughout the day.
          [F:119] Catering by Celeste provided catering for the Voyage Vista Travel grand opening.
          [F:120] Voyage Vista Travel booked three new client consultations directly from its grand opening event.
          [F:121] Sarah Chen collected over 50 solid leads at the Voyage Vista Travel grand opening.

# teacher
# TEACHER-ENDPOINT tree task_id=46 entity='Company: Voyage Vista Travel' (oracle cache replay)
# facts=122 aspects=24 max_depth=6 avg_branching=5.84
# max_facts=0 replay_time=9.4s legacy_cache_key=auto persona=Company_Voyage_Vista_Travel
# cache: {"s1": {"hits": 313, "misses": 0, "puts": 0, "legacy_hits": 0}, "s2": {"hits": 132, "misses": 0, "puts": 0, "legacy_hits": 0}} oracle_fail_events=0
# teacher endpoint valid ONLY relative to cache_dir=/mnt/algo_code/zyli/2S_RL/runs/forest_v2/prod100_cache_5mini (audit R10: stale-teacher — regenerate after any prompt/payload change); misses below were answered by the LIVE vLLM (current teacher)

[ROOT] Voyage Vista Travel
  [F:101] Things at Voyage Vista Travel are good.
  [A] Program and Itinerary Planning
    [F:78] Voyage Vista Travel has upcoming European tours planned for summer 2023.
    [F:79] Voyage Vista Travel is planning potential new South American itineraries.
    [A] Synaptic Solutions adventure retreat
      [F:2] Voyage Vista Travel organized an adventure retreat for Synaptic Solutions.
      [F:3] The adventure retreat organized by Voyage Vista Travel was for 40 employees of Synaptic Solutions.
      [F:7] The adventure retreat organized by Voyage Vista Travel is worth over $75,000.
      [A] Trip duration, destination & dates
        [F:55] The Andes Mountain Trek for The Sterling Group is scheduled for May 2024.
        [F:57] The Chen Family is looking for a 3-week luxury safari in Kenya for July 2024.
        [F:71] The new adventure packages include volcano trekking around Arenal.
        [A] Synaptic Solutions retreat schedule & location
          [F:4] The adventure retreat organized by Voyage Vista Travel is a 5-day multi-activity trip.
          [F:5] The adventure retreat organized by Voyage Vista Travel is to the Costa Rican rainforest.
          [F:6] The adventure retreat organized by Voyage Vista Travel is scheduled for August 2023.
          [F:72] The new adventure packages include white-water rafting on the Pacuare River.
        [F:73] The new adventure packages include cloud forest canopy tours.
      [F:58] The safari for the Chen Family is a high-end, bespoke itinerary.
      [F:87] The 'Sustainable Travel Guide' focuses on immersive local experiences.
    [A] Tech-company retreat outreach
      [A] Proposal development and client meetings
        [F:1] Voyage Vista Travel plans to present proposals for corporate team-building retreats to local tech co
        [F:50] The follow-up meeting with Zenith Innovations is for a team-building retreat.
        [F:0] Voyage Vista Travel is brainstorming proposals for corporate team-building retreats.
        [F:47] The tasting event was for ten potential clients from various tech startups in Silicon Valley.
        [F:54] Voyage Vista Travel is developing a 14-day Andes Mountain Trek in Patagonia for The Sterling Group's
        [F:20] Voyage Vista Travel is proposing a week-long leadership and team-building retreat for TechSolutions 
        [A] Syntel Systems corporate retreat program
          [F:14] The corporate retreat program for Syntel Systems includes guided hikes.
          [F:9] The corporate retreat program for Syntel Systems is for 50 employees.
          [F:10] The corporate retreat program for Syntel Systems is focused on team-building and leadership developm
          [F:23] Voyage Vista Travel is planning a 'build your own igloo' survival challenge for the TechSolutions In
          [F:29] The total project value for the TechSolutions Inc. corporate retreat is estimated at $120,000.
      [A] Schedule and location
        [F:30] The TechSolutions Inc. Whistler trip is scheduled for January 15, 2024.
        [F:31] Voyage Vista Travel hit a minor snag with ski and snowboard rentals for the TechSolutions Inc. Whist
        [A] Syntel Systems retreat (Aug 15-18, Redwood National Park)
          [F:11] The corporate retreat program for Syntel Systems is a four-day, three-night program.
          [F:12] The corporate retreat program for Syntel Systems is planned for Redwood National Park.
          [F:13] The corporate retreat program for Syntel Systems is scheduled for August 15th to August 18th.
          [F:15] The corporate retreat program for Syntel Systems includes a conservation project.
      [F:21] Voyage Vista Travel is planning guided backcountry skiing for the TechSolutions Inc. retreat.
      [F:22] Voyage Vista Travel is planning snowshoeing for the TechSolutions Inc. retreat.
    [A] Chen Family safari itinerary
      [F:59] The safari for the Chen Family includes stays at Cottar's 1920s Camp.
      [F:60] The safari for the Chen Family includes stays at Angama Mara.
      [F:61] The safari for the Chen Family focuses on ethical wildlife viewing.
      [F:62] The safari for the Chen Family focuses on immersive cultural experiences.
      [F:63] The budget for the Chen Family's safari is approaching $75,000.
    [F:70] Voyage Vista developed five exclusive new adventure packages.
    [A] Sustainable Travel Guide
      [F:86] Liam O'Connell finished developing a new 'Sustainable Travel Guide' for Voyage Vista's clients.
      [F:91] Liam O'Connell is planning a new Instagram campaign for Voyage Vista called 'Hidden Gems of Costa Ri
      [F:92] The 'Hidden Gems of Costa Rica' campaign is planned for a Q1 2024 launch.
      [A] Environmental focus and featured destinations
        [F:88] The 'Sustainable Travel Guide' focuses on minimizing environmental impact.
        [F:89] The 'Sustainable Travel Guide' highlights destinations like Costa Rica's cloud forests.
        [F:90] The 'Sustainable Travel Guide' highlights destinations like Patagonia's conservation efforts.
      [F:69] Voyage Vista Travel plans to offer unique week-long excursions through the partnership with Finca Ve
      [F:74] The new adventure packages include sustainable coffee farm visits.
      [F:75] The new adventure packages include indigenous Bribri culture immersion.
      [F:41] The first 'Peruvian Highlands Trek' group for Voyage Vista Travel's EcoTrek Adventures partnership i
    [F:43] The 'Peruvian Highlands Trek' group departs on April 10, 2024.
    [F:44] The 'Peruvian Highlands Trek' group's departure is scheduled for April 10, 2024.
  [A] Sales and Partnerships
    [A] Client wins and closed opportunities
      [A] Closed client wins
        [A] Syntel Systems corporate contract
          [F:17] The contract with Syntel Systems positions Voyage Vista for more corporate work.
          [F:8] Voyage Vista landed a massive corporate retreat program with Syntel Systems.
          [F:16] The contract with Syntel Systems is a significant revenue boost for Voyage Vista.
        [A] Major client and partner closures
          [F:67] Voyage Vista finalized its partnership with Casa Verde Eco-Resort.
          [F:68] The partnership with Casa Verde Eco-Resort took months to finalize.
          [F:39] The 'Sustainable Safaris' partnership with EcoTrek Adventures closed on January 31, 2024.
          [F:40] The 'Sustainable Safaris' partnership with EcoTrek Adventures met its $150,000 first-year goal.
          [A] Notable client contracts
            [F:26] Voyage Vista Travel landed the TechSolutions Inc. contract for their corporate retreat in Whistler.
            [F:53] Liam O'Connell landed The Sterling Group as a major new client for Voyage Vista Travel.
            [F:56] Voyage Vista Travel finalized a major booking with the Chen Family.
          [F:28] Voyage Vista Travel received an initial deposit of $45,000 for the TechSolutions Inc. corporate retr
          [F:37] Voyage Vista Travel closed a significant deal.
          [F:36] Voyage Vista Travel finalized Whistler Blackcomb Gear Rentals on December 28, 2023.
        [F:120] Voyage Vista Travel booked three new client consultations directly from its grand opening event.
        [F:18] Voyage Vista Travel landed its first big proposal for corporate team-building retreats.
        [F:42] Voyage Vista Travel has secured 8 bookings out of 12 available spots for the first 'Peruvian Highlan
      [F:64] Voyage Vista Travel is looking at a partnership with a luxury eco-lodge in Costa Rica.
      [F:65] Voyage Vista Travel hopes to finalize a deal with Finca Verde for a launch by early 2024.
      [F:108] Sarah Chen streamlined Voyage Vista Travel's client intake forms within her first 9 days of work.
      [F:109] Sarah Chen reorganized Voyage Vista Travel's client database within her first 9 days of work.
      [A] Prospective enterprise clients
        [F:49] Voyage Vista Travel secured a follow-up meeting with Zenith Innovations.
        [F:51] Securing Salesforce as a client would be huge for Voyage Vista Travel.
        [F:52] Securing Google as a client would be huge for Voyage Vista Travel.
        [A] TechSolutions Inc. opportunity
          [F:33] Alpine Peak Sports could not provide 45 sets of premium gear for the TechSolutions Inc. Whistler tri
          [F:19] Voyage Vista Travel's first big proposal is with TechSolutions Inc.
          [F:27] The meeting for the TechSolutions Inc. contract happened on November 17, 2023.
      [F:24] Voyage Vista Travel is working on the final proposal for TechSolutions Inc. during the week of Novem
      [F:25] Voyage Vista Travel is aiming for a client meeting with TechSolutions Inc. on November 17, 2023.
      [F:38] The deal confirms Voyage Vista Travel's niche in high-end corporate adventure travel.
      [F:34] Voyage Vista Travel secured a bulk deal with Whistler Blackcomb Gear Rentals for the ski and snowboa
      [F:32] Voyage Vista Travel's original vendor for ski and snowboard rentals was Alpine Peak Sports.
      [F:35] The deal with Whistler Blackcomb Gear Rentals was at a 5% higher cost than the original vendor.
    [A] Marketing and Promotion
      [F:66] The partnership Voyage Vista Travel is looking at is for exclusive travel packages.
      [A] Social media strategy and staffing
        [F:77] Voyage Vista is doing a targeted social media campaign for the new Costa Rica packages.
        [F:110] Chloe Davies is interested in helping with Voyage Vista Travel's social media strategy.
        [F:111] Chloe Davies might start as a consultant for Voyage Vista Travel by March 15, 2024.
        [F:99] Voyage Vista Travel received a mention on the 'Bay Area Wanderlust' travel blog.
        [F:100] The mention on the 'Bay Area Wanderlust' travel blog drove a lot of traffic to Voyage Vista Travel's
        [F:96] Voyage Vista Travel allocated $15,000 for digital ads for the new marketing campaign.
        [A] Campaign planning and targets
          [F:95] The new marketing campaign for Voyage Vista Travel is launching in early January 2024.
          [F:93] Liam O'Connell and David Miller finalized a new marketing campaign for Voyage Vista Travel.
          [F:94] The new marketing campaign for Voyage Vista Travel targets luxury eco-tourism packages for the sprin
          [F:97] Voyage Vista Travel is launching a new digital marketing campaign for luxury eco-tourism on January 
          [F:98] Voyage Vista Travel has a digital marketing launch scheduled for January 8, 2024.
        [F:113] Voyage Vista Travel has sent out about 75 invitations for its grand opening.
        [F:45] The promotional materials for the 'Peruvian Highlands Trek' need to be ready by February 21, 2024.
      [A] Website redesign and specialists
        [F:85] Liam O'Connell is aiming for a beta launch of the Voyage Vista website redesign by mid-August 2023.
        [F:82] The Voyage Vista website redesign needs to better reflect their new target audience.
        [F:84] Maria specializes in travel industry sites.
        [F:76] Voyage Vista is launching the new Costa Rica packages on its website on June 26, 2023.
      [F:117] The Voyage Vista Travel grand opening was a huge success.
      [F:118] Over 120 people attended the Voyage Vista Travel grand opening throughout the day.
      [F:119] Catering by Celeste provided catering for the Voyage Vista Travel grand opening.
      [F:121] Sarah Chen collected over 50 solid leads at the Voyage Vista Travel grand opening.
  [F:107] Sarah Chen will start working at Voyage Vista Travel on October 2, 2023.
  [F:80] Liam O'Connell started working on a complete website redesign for Voyage Vista.
  [F:83] Liam O'Connell brought in a freelance designer named Maria for the Voyage Vista website redesign.
  [F:102] The new office space for Voyage Vista is intended to be in the SOMA district.
  [F:103] The new office space for Voyage Vista is hoped to be around 1500 sq ft.
  [F:104] The new office space for Voyage Vista is hoped to accommodate 8-10 people.
  [F:105] Liam O'Connell signed the lease for Voyage Vista Travel's 1600 sq ft SOMA office on Mission Street o
  [F:106] Voyage Vista Travel will get the keys to its new SOMA office on Mission Street on September 1, 2023.
  [F:46] Voyage Vista Travel hosted a successful tasting event on June 01, 2023.
  [F:48] The tasting event featured ceviche and Pisco Sours.
  [F:81] The Voyage Vista website redesign needs to better reflect expanded offerings.
  [F:115] Catering by Celeste confirmed their setup for 75 guests for the Voyage Vista Travel grand opening.
  [F:116] The Voyage Vista Travel grand opening was on October 12, 2023.
  [F:112] Voyage Vista Travel's grand opening is planned for October 12, 2023.
  [F:114] The Voyage Vista Travel grand opening is scheduled for October 12, 2023.

# POLICY-GROWN tree task_id=46 entity='Company: Voyage Vista Travel' (S1=ckpt_epoch39.pt, S2=ckpt_epoch39.pt, decode_prior_tau=0.0)
# facts=122 aspects=24 max_depth=8 avg_branching=5.84
# s2_rounds=122 decode_rows=960 | GROUP_rows=114 (group_events=24) DOWN=0 UP=0 IDLE=846 (idle_frac=0.881) | s1_navs=422 apply_errors=0
# max_facts=0 rollout_time=10.4s reference: tree_46_reference.txt
# R10 pairF1(policy,teacher)=0.1385 [flat-null=0.0915] n_common=122 teacher(aspects=24 depth=6) teacher_vs_ref_pairF1=0.2737
# R10 pairF1(policy,reference)=0.3477 [flat-null=0.0569] n_common=122
# R10 policy depth_hist={1: 2, 2: 6, 3: 15, 4: 46, 5: 41, 6: 3, 7: 5, 8: 4} longest_1child_chain=0 aspect_edges=24 same_name_ratio=0.000 high_sim(cos>0.9)_ratio=0.000

[ROOT] Voyage Vista Travel
  [F:78] Voyage Vista Travel has upcoming European tours planned for summer 2023.
  [F:0] Voyage Vista Travel is brainstorming proposals for corporate team-building retreats.
  [A] Positive Business Momentum
    [F:101] Things at Voyage Vista Travel are good.
    [F:1] Voyage Vista Travel plans to present proposals for corporate team-building retreats to local tech co
    [A] Media Coverage and Itinerary Expansion
      [F:7] The adventure retreat organized by Voyage Vista Travel is worth over $75,000.
      [F:47] The tasting event was for ten potential clients from various tech startups in Silicon Valley.
      [A] Blog Mention and Retreats
        [F:79] Voyage Vista Travel is planning potential new South American itineraries.
        [F:99] Voyage Vista Travel received a mention on the 'Bay Area Wanderlust' travel blog.
        [F:100] The mention on the 'Bay Area Wanderlust' travel blog drove a lot of traffic to Voyage Vista Travel's
        [F:6] The adventure retreat organized by Voyage Vista Travel is scheduled for August 2023.
      [F:49] Voyage Vista Travel secured a follow-up meeting with Zenith Innovations.
      [F:52] Securing Google as a client would be huge for Voyage Vista Travel.
      [A] Client Acquisition and Retreats
        [F:50] The follow-up meeting with Zenith Innovations is for a team-building retreat.
        [F:51] Securing Salesforce as a client would be huge for Voyage Vista Travel.
        [F:27] The meeting for the TechSolutions Inc. contract happened on November 17, 2023.
        [F:29] The total project value for the TechSolutions Inc. corporate retreat is estimated at $120,000.
        [F:30] The TechSolutions Inc. Whistler trip is scheduled for January 15, 2024.
      [F:55] The Andes Mountain Trek for The Sterling Group is scheduled for May 2024.
    [F:5] The adventure retreat organized by Voyage Vista Travel is to the Costa Rican rainforest.
    [A] Corporate Adventure Retreats
      [F:2] Voyage Vista Travel organized an adventure retreat for Synaptic Solutions.
      [F:3] The adventure retreat organized by Voyage Vista Travel was for 40 employees of Synaptic Solutions.
      [F:4] The adventure retreat organized by Voyage Vista Travel is a 5-day multi-activity trip.
      [F:48] The tasting event featured ceviche and Pisco Sours.
    [F:37] Voyage Vista Travel closed a significant deal.
    [F:38] The deal confirms Voyage Vista Travel's niche in high-end corporate adventure travel.
    [F:60] The safari for the Chen Family includes stays at Angama Mara.
    [A] Luxury Bespoke Itineraries
      [F:46] Voyage Vista Travel hosted a successful tasting event on June 01, 2023.
      [F:53] Liam O'Connell landed The Sterling Group as a major new client for Voyage Vista Travel.
      [F:86] Liam O'Connell finished developing a new 'Sustainable Travel Guide' for Voyage Vista's clients.
      [A] Kenya Luxury Safari
        [F:57] The Chen Family is looking for a 3-week luxury safari in Kenya for July 2024.
        [F:58] The safari for the Chen Family is a high-end, bespoke itinerary.
        [F:59] The safari for the Chen Family includes stays at Cottar's 1920s Camp.
        [F:61] The safari for the Chen Family focuses on ethical wildlife viewing.
        [F:62] The safari for the Chen Family focuses on immersive cultural experiences.
        [F:63] The budget for the Chen Family's safari is approaching $75,000.
        [F:76] Voyage Vista is launching the new Costa Rica packages on its website on June 26, 2023.
        [F:77] Voyage Vista is doing a targeted social media campaign for the new Costa Rica packages.
      [A] Patagonia Executive Trek
        [F:54] Voyage Vista Travel is developing a 14-day Andes Mountain Trek in Patagonia for The Sterling Group's
        [F:56] Voyage Vista Travel finalized a major booking with the Chen Family.
      [F:111] Chloe Davies might start as a consultant for Voyage Vista Travel by March 15, 2024.
      [F:8] Voyage Vista landed a massive corporate retreat program with Syntel Systems.
      [A] Sustainable Travel Operations
        [F:87] The 'Sustainable Travel Guide' focuses on immersive local experiences.
        [F:88] The 'Sustainable Travel Guide' focuses on minimizing environmental impact.
        [F:89] The 'Sustainable Travel Guide' highlights destinations like Costa Rica's cloud forests.
        [F:90] The 'Sustainable Travel Guide' highlights destinations like Patagonia's conservation efforts.
        [A] New Hire Onboarding
          [F:107] Sarah Chen will start working at Voyage Vista Travel on October 2, 2023.
          [F:108] Sarah Chen streamlined Voyage Vista Travel's client intake forms within her first 9 days of work.
          [F:72] The new adventure packages include white-water rafting on the Pacuare River.
        [A] Cultural Immersion Marketing
          [F:110] Chloe Davies is interested in helping with Voyage Vista Travel's social media strategy.
          [F:75] The new adventure packages include indigenous Bribri culture immersion.
        [F:84] Maria specializes in travel industry sites.
        [A] Digital Platform Expansion
          [F:109] Sarah Chen reorganized Voyage Vista Travel's client database within her first 9 days of work.
          [F:73] The new adventure packages include cloud forest canopy tours.
          [F:74] The new adventure packages include sustainable coffee farm visits.
          [F:80] Liam O'Connell started working on a complete website redesign for Voyage Vista.
          [F:81] The Voyage Vista website redesign needs to better reflect expanded offerings.
          [F:82] The Voyage Vista website redesign needs to better reflect their new target audience.
          [F:83] Liam O'Connell brought in a freelance designer named Maria for the Voyage Vista website redesign.
        [A] Voyage Vista Marketing Initiatives
          [F:94] The new marketing campaign for Voyage Vista Travel targets luxury eco-tourism packages for the sprin
          [F:98] Voyage Vista Travel has a digital marketing launch scheduled for January 8, 2024.
          [A] Liam O'Connell Projects
            [F:114] The Voyage Vista Travel grand opening is scheduled for October 12, 2023.
            [A] Voyage Vista Launch Plans
              [F:112] Voyage Vista Travel's grand opening is planned for October 12, 2023.
              [F:35] The deal with Whistler Blackcomb Gear Rentals was at a 5% higher cost than the original vendor.
              [A] Marketing Campaigns
                [F:85] Liam O'Connell is aiming for a beta launch of the Voyage Vista website redesign by mid-August 2023.
                [F:91] Liam O'Connell is planning a new Instagram campaign for Voyage Vista called 'Hidden Gems of Costa Ri
                [F:92] The 'Hidden Gems of Costa Rica' campaign is planned for a Q1 2024 launch.
                [F:93] Liam O'Connell and David Miller finalized a new marketing campaign for Voyage Vista Travel.
            [F:116] The Voyage Vista Travel grand opening was on October 12, 2023.
            [F:26] Voyage Vista Travel landed the TechSolutions Inc. contract for their corporate retreat in Whistler.
            [A] TechSolutions Inc. Proposal
              [F:18] Voyage Vista Travel landed its first big proposal for corporate team-building retreats.
              [F:24] Voyage Vista Travel is working on the final proposal for TechSolutions Inc. during the week of Novem
              [F:25] Voyage Vista Travel is aiming for a client meeting with TechSolutions Inc. on November 17, 2023.
          [F:120] Voyage Vista Travel booked three new client consultations directly from its grand opening event.
          [F:31] Voyage Vista Travel hit a minor snag with ski and snowboard rentals for the TechSolutions Inc. Whist
          [F:32] Voyage Vista Travel's original vendor for ski and snowboard rentals was Alpine Peak Sports.
          [F:34] Voyage Vista Travel secured a bulk deal with Whistler Blackcomb Gear Rentals for the ski and snowboa
          [F:36] Voyage Vista Travel finalized Whistler Blackcomb Gear Rentals on December 28, 2023.
          [F:42] Voyage Vista Travel has secured 8 bookings out of 12 available spots for the first 'Peruvian Highlan
        [A] Grand Opening Logistics
          [F:95] The new marketing campaign for Voyage Vista Travel is launching in early January 2024.
          [F:96] Voyage Vista Travel allocated $15,000 for digital ads for the new marketing campaign.
          [F:97] Voyage Vista Travel is launching a new digital marketing campaign for luxury eco-tourism on January 
          [F:113] Voyage Vista Travel has sent out about 75 invitations for its grand opening.
          [F:115] Catering by Celeste confirmed their setup for 75 guests for the Voyage Vista Travel grand opening.
        [F:23] Voyage Vista Travel is planning a 'build your own igloo' survival challenge for the TechSolutions In
        [A] Corporate Retreat Planning
          [F:117] The Voyage Vista Travel grand opening was a huge success.
          [F:118] Over 120 people attended the Voyage Vista Travel grand opening throughout the day.
          [F:119] Catering by Celeste provided catering for the Voyage Vista Travel grand opening.
          [F:121] Sarah Chen collected over 50 solid leads at the Voyage Vista Travel grand opening.
          [F:20] Voyage Vista Travel is proposing a week-long leadership and team-building retreat for TechSolutions 
          [F:21] Voyage Vista Travel is planning guided backcountry skiing for the TechSolutions Inc. retreat.
          [F:22] Voyage Vista Travel is planning snowshoeing for the TechSolutions Inc. retreat.
          [F:33] Alpine Peak Sports could not provide 45 sets of premium gear for the TechSolutions Inc. Whistler tri
        [F:41] The first 'Peruvian Highlands Trek' group for Voyage Vista Travel's EcoTrek Adventures partnership i
        [F:43] The 'Peruvian Highlands Trek' group departs on April 10, 2024.
        [F:44] The 'Peruvian Highlands Trek' group's departure is scheduled for April 10, 2024.
        [F:45] The promotional materials for the 'Peruvian Highlands Trek' need to be ready by February 21, 2024.
      [A] Syntel Leadership Retreat
        [F:9] The corporate retreat program for Syntel Systems is for 50 employees.
        [F:10] The corporate retreat program for Syntel Systems is focused on team-building and leadership developm
        [F:11] The corporate retreat program for Syntel Systems is a four-day, three-night program.
        [F:12] The corporate retreat program for Syntel Systems is planned for Redwood National Park.
        [F:13] The corporate retreat program for Syntel Systems is scheduled for August 15th to August 18th.
        [F:14] The corporate retreat program for Syntel Systems includes guided hikes.
      [A] Syntel and EcoTrek Partnerships
        [F:102] The new office space for Voyage Vista is intended to be in the SOMA district.
        [A] Syntel and EcoTrek Outcomes
          [F:15] The corporate retreat program for Syntel Systems includes a conservation project.
          [F:16] The contract with Syntel Systems is a significant revenue boost for Voyage Vista.
          [F:17] The contract with Syntel Systems positions Voyage Vista for more corporate work.
          [F:39] The 'Sustainable Safaris' partnership with EcoTrek Adventures closed on January 31, 2024.
          [F:40] The 'Sustainable Safaris' partnership with EcoTrek Adventures met its $150,000 first-year goal.
        [F:103] The new office space for Voyage Vista is hoped to be around 1500 sq ft.
        [F:104] The new office space for Voyage Vista is hoped to accommodate 8-10 people.
        [F:19] Voyage Vista Travel's first big proposal is with TechSolutions Inc.
        [A] Voyage Vista Office Lease
          [F:105] Liam O'Connell signed the lease for Voyage Vista Travel's 1600 sq ft SOMA office on Mission Street o
          [F:106] Voyage Vista Travel will get the keys to its new SOMA office on Mission Street on September 1, 2023.
          [F:28] Voyage Vista Travel received an initial deposit of $45,000 for the TechSolutions Inc. corporate retr
      [F:71] The new adventure packages include volcano trekking around Arenal.
      [A] Eco-Resort Partnerships
        [F:64] Voyage Vista Travel is looking at a partnership with a luxury eco-lodge in Costa Rica.
        [F:65] Voyage Vista Travel hopes to finalize a deal with Finca Verde for a launch by early 2024.
        [F:66] The partnership Voyage Vista Travel is looking at is for exclusive travel packages.
        [F:67] Voyage Vista finalized its partnership with Casa Verde Eco-Resort.
        [F:68] The partnership with Casa Verde Eco-Resort took months to finalize.
        [F:69] Voyage Vista Travel plans to offer unique week-long excursions through the partnership with Finca Ve
        [F:70] Voyage Vista developed five exclusive new adventure packages.

Agent Memory Research In MMLab (Part 1)

Tue, 04 Aug 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

项目说明

本项目是我在HKU MMLab实习期间, 由Dr. Xinyu Pan和Prof.Bo Dai指导的Agent Memory的科研项目的研究与实验记录, 在此之前我们已经做了大概4个月(From April 2026), 所以我会先写一些之前的设计和问题, 以及工程经验, 然后到了后面应该每次写的就都是最新的问题和研究成果了，对于踩坑了的地方，也许都会专门开一篇来写

文档说明

本文档概述训练的全流程，

研究问题

通俗的来说, 现在业界用的memory有两个问题:

结构过于平坦, 类似数据库, 不能很好的反应记忆元素之间的上下级关系
在调整记忆之间的关系的时候, 总是以来agent自身去调整, 成本高+时间慢

解决办法

我们针对这两个问题各自提出一个办法,

提出一个叫做membrain的树记忆结构, 每个记忆就算做是树的一个节点, 这样天然提供上下级和其他的语义关系
用一个小模型结合自己设计的动作空间来决策如何调整记忆结构

理想的话, 这应该可以大幅度降低开销

研究中的数据结构

Membrain

Membrain是一个记忆系统, 其中每个元素是一棵树，根代表记忆归属的主体, 比如说所有和电脑相关的记忆都归属于电脑这个根节点

Aspect节点代表一种子类，比如笔记本电脑就是电脑根节点的一个子类

Fact代表事实/记忆本身，比如说笔记本电脑在xxxx年被发明，这就是一个记忆

简单理解就是一个N叉树，上下级关系表示语义

一颗完整的树可能很宽很深，因为语义可以大分叉

闭包

闭包是一个专有名词，他必须由一个fact节点生成，规定他包含fact自己，兄弟，父亲，父亲的兄弟，以及祖父

当然有可能深度不够，导致没有祖父，那就是父亲和兄弟

这是我们定义的一个重要数据结构，以后的操作都在这上面做

两层子树

用这个专有名词来指代字面意思，某个node和他的所有直接孩子，叫做两层子树

两阶段决策拆分

一阶段stage1

假设你有一颗membrain树，现在来了个新节点，需要挂上去，这总可以用这样一个递归算法决策：

cur = root # 从root开始
new_node # 新node
while True:
    subtree = [cur + child for child in cur.children] # 两层子树
    decision_index = stage1_network(new_node, subtree) # 输出决策，必然是cur，或者cur的孩子中的aspect节点
    if subtree[decision_index] == cur: # 如果是决定挂载到cur上
        cur.child.append(new_node)
        return subtree[decision_index] # 决策完毕，挂上去

    else: # 如果是决定挂载到cur的孩子aspect
        cur = subtree[decision_index] # 又从那个孩子aspect重新决策

因为树高有限，每次往下走一步，总是会结束的，所以不会无限循环，到最后一层决策完了就跳出去了

stage1 network

讲简单点，这个网络就负责吃一个二层子树，然后输出一串logits，最后argmax一下把index选出来

可以理解为就是一个transformer_block，然后接一个Linear做成logits

二阶段Stage2

现在新的fact已经放到tree上了，但我假设这个fact周围的一些node排列还不是很合理，所以我希望另一个网络可以输出一些树的可执行动作目的就是让附近的node排列变得更合理

closure = build_closure(fact) # 根据这个fact生成closure
Action_Series = stage2_network(closure) # 生成做的动作
tree = apply_action(Action_Series) # 把动作做上去

显然这里就没有stage1那么容易了，因为这里的action_space比较复杂，不是选个index就完事了的

二阶段stage2的动作

Promote 把一个fact节点提到他的grandparent下

def Promote(fact):
    fact.parent.children.remove(fact)
    fact.parent.parent.children.append(fact)
    fact.parent = fact.parent.parent

Demote 把一个fact节点放到他的sibling aspect下

def Demote(fact, sibling_aspect):
    fact.parent.children.remove(fact)
    sibling_aspect.children.append(fact)
    fact.parnet = sibling_aspect

Group 把一堆fact放一起，做一个新的aspect为他们的parent，然后这个新aspect的parent是他们的老parent

def Group(facts):
    # Assert all fact in facts have same parent
    old_parent = facts[0].parent
    new_aspect = create_new_aspect(facts)
    new_aspect.parent = old_parent
    for fact in facts:
        old_parent.children.remove(fact)
        new_aspect.children.append(fact)

Idle 什么都不做

def idle():
    return

stage2 network

这里比较复杂，这个网络吃进去一个闭包，输出三个logits，分别代表:每个node做什么类型动作，Demote动作的对象，Group动作的对象

整个过程要借助一个额外的概率图模型和势函数，以及一个规划优化器来完成，这里不展开，先默认通过这三个logits可以得到一系列动作就行

训练

采用MLE训练，简单来说，这是一个在线的训练，每当模型需要做决策的时候，比如说stage1挂在哪里，stage2怎么调整，就去问问llm怎么做，然后用模型产出的logits和llm给的ground truth做CE loss训练就好了

具体场景

想象一下有一堆fact node是零散的，现在我要把它拼出一棵树来，那实际上就是来一个fact走一遍stage1和stage2，然后树长大一截，在这个过程当中训练就是自然而然的事情了

ongoing_tree = root
for fact in all_facts:
    ongoing_tree = update_tree(ongoing_tree, fact, stage1_model)
    stage1_model.step()

    ongoing_tree = update_tree(ongoing_tree, fact, stage2_model)
    stage2_model.step()

上面要讲两点，第一就是我可以不执行我自己model作出的决策，而执行LLM给的ground truth让ongoing_tree长大，其二就是可以不每次都更新模型，可以把梯度攒起来到了batch_size再更新

Agent Memory Research In MMLab (Part 2)

Tue, 04 Aug 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

文档说明

本文档说明数据集的制作方式

长程对话到树

假设我们有一个Locomo10那样的双人长程对话，那么LLM可以很轻松的从其中抽取出fact和树来，他只要自行判断对话当中的事实属于什么对象就行了，通俗理解就是replay session，等到session重放了一遍树也就做出来了，

树的分布不均

实验证明一个长程对话大概能抽取出5-6颗树，但是如果不加限制(如LocoMo那样)，容易出现不均，即对话的双方各占了一颗树，然后这两棵树可能总共占了90%以上的fact，剩下的树基本上是空的

上面这个图是从真实的一个locomo10数据集里面的一个长程对话(的一部分)抽出来的4颗树，明显这个对话就是James和John之间的对话，但是LLM可能识别到其他两个项目也算主体，所以又弄了两棵树，但是fact很少，约等于没有，所以如果我们照着Locomo去造对话的话，大概率会有这个问题

对话的要求

所以我们需要这样的长程对话：这个长程对话由若干个session组成，最后LLM可以从这个对话里面抽出若干棵树，然后每棵树的fact数量要差不多

如此一来不妨用一个跨session的配平算法来达到这个目的，对话的长度不需要先验的设定，fact不够就让他继续顺下去说就可以了

如此一来的话，对话和树都有了，但是作为训练数据集我们不需要对话，把树存下来就好了

树的例子

来看两件事情：第一是LLM能从一个多轮对话里面抽出几颗什么样的树，第二就是随机的抽一棵树看树的结构如何

我们希望的是：树和树之间，深度和fact数量要差不多，树内上下级关系合理，不要出现极端深度/flat的情况

最下面两颗树是副产物，不列入正式训练的，可以看到在这一场Layla和Anh的长对话当中，弄出了5颗fact和深度都差不多的树，达到了我们的要求

再来挑一棵树看看结构

# 树结构样例(gemini 产,production100 / dlg00017)
# root: Sofia Ramirez  (Person)
# facts=185  目录=aspect(保留名字便于读结构),叶子=fact(仅编号不展开内容)
#──────────────────────────────────────────────────────────────────────

Sofia Ramirez/   [ROOT]
├── Professional Life/
│   ├── Medical Training & Certifications/
│   │   └── Advanced Trauma Life Support (ATLS) Program/
│   │       ├── Program Overview & Enrollment/
│   │       │   ├── Program Details/
│   │       │   │   ├── fact_001
│   │       │   │   ├── fact_002
│   │       │   │   └── fact_003
│   │       │   └── Enrollment & Schedule/
│   │       │       ├── fact_004
│   │       │       ├── fact_005
│   │       │       ├── fact_006
│   │       │       └── fact_007
│   │       ├── Pre-Course Preparation/
│   │       │   ├── fact_008
│   │       │   ├── fact_009
│   │       │   └── fact_010
│   │       ├── Course Modules & Progress/
│   │       │   ├── Early Modules (Assessment & Airway)/
│   │       │   │   ├── fact_011
│   │       │   │   ├── fact_012
│   │       │   │   ├── fact_013
│   │       │   │   ├── fact_014
│   │       │   │   ├── fact_015
│   │       │   │   ├── fact_016
│   │       │   │   └── fact_017
│   │       │   └── Later Modules (Shock & Trauma Assessment)/
│   │       │       ├── fact_018
│   │       │       └── fact_019
│   │       ├── Practicals & Assessments/
│   │       │   ├── fact_020
│   │       │   ├── fact_021
│   │       │   ├── fact_022
│   │       │   ├── fact_023
│   │       │   └── fact_024
│   │       └── Course Completion/
│   │           ├── fact_025
│   │           ├── fact_026
│   │           └── fact_027
│   └── Emergency Medical Services (EMS) Role/
│       ├── Coastline EMS Employment/
│       │   ├── fact_028
│       │   └── fact_029
│       ├── Incident Responses/
│       │   ├── February 2024 Incidents/
│       │   │   └── fact_030
│       │   ├── March 2024 Incidents/
│       │   │   ├── fact_031
│       │   │   ├── fact_032
│       │   │   ├── fact_033
│       │   │   ├── fact_034
│       │   │   ├── fact_035
│       │   │   └── fact_036
│       │   └── May 2024 Incidents/
│       │       └── fact_037
│       └── Recognition & Awards/
│           ├── fact_038
│           └── fact_039
└── Personal Pursuits/
    ├── Fitness & Athletics/
    │   ├── Boxing Training/
    │   │   ├── Training Progress & General Outcomes/
    │   │   │   ├── fact_040
    │   │   │   ├── fact_041
    │   │   │   ├── fact_042
    │   │   │   ├── fact_043
    │   │   │   ├── fact_044
    │   │   │   └── fact_045
    │   │   ├── Sparring Sessions/
    │   │   │   ├── Past Sparring Sessions/
    │   │   │   │   ├── fact_046
    │   │   │   │   └── fact_047
    │   │   │   └── Upcoming Sparring Sessions/
    │   │   │       ├── fact_048
    │   │   │       ├── fact_049
    │   │   │       ├── fact_050
    │   │   │       └── fact_051
    │   │   ├── Competition Results & Specific Victories/
    │   │   │   ├── Past Competitions (Pre-October 2024)/
    │   │   │   │   ├── fact_052
    │   │   │   │   ├── fact_053
    │   │   │   │   ├── fact_054
    │   │   │   │   ├── fact_055
    │   │   │   │   ├── fact_056
    │   │   │   │   ├── fact_057
    │   │   │   │   └── fact_058
    │   │   │   └── October 2024 Golden Gloves Bout/
    │   │   │       ├── Bout Details & Outcome/
    │   │   │       │   ├── fact_059
    │   │   │       │   ├── fact_060
    │   │   │       │   ├── fact_061
    │   │   │       │   ├── fact_062
    │   │   │       │   ├── fact_063
    │   │   │       │   ├── fact_064
    │   │   │       │   └── fact_065
    │   │   │       └── Post-Bout Recovery/
    │   │   │           ├── fact_066
    │   │   │           └── fact_067
    │   │   └── Coach & Training Regimen/
    │   │       ├── fact_068
    │   │       ├── fact_069
    │   │       ├── fact_070
    │   │       ├── fact_071
    │   │       ├── fact_072
    │   │       ├── fact_073
    │   │       ├── fact_074
    │   │       └── fact_075
    │   ├── Incline Climbing/
    │   │   ├── fact_076
    │   │   └── fact_077
    │   └── Strength & Plyometrics Training/
    │       ├── Overall Goals & Achievements/
    │       │   ├── fact_078
    │       │   ├── fact_079
    │       │   ├── fact_080
    │       │   ├── fact_081
    │       │   └── fact_082
    │       └── Plyometrics Training Details/
    │           ├── fact_083
    │           ├── fact_084
    │           ├── fact_085
    │           └── fact_086
    ├── The Summit Seekers Project/
    │   ├── Project Foundation & Vision/
    │   │   ├── Project Naming & Core Vision/
    │   │   │   ├── fact_087
    │   │   │   └── fact_088
    │   │   ├── Charity & Educational Focus/
    │   │   │   ├── fact_089
    │   │   │   ├── fact_090
    │   │   │   ├── fact_091
    │   │   │   └── fact_092
    │   │   └── Team Members/
    │   │       ├── fact_093
    │   │       └── fact_094
    │   ├── Project Planning & Development/
    │   │   ├── Website & Initial Setup/
    │   │   │   └── fact_095
    │   │   ├── Coordination & Meetings/
    │   │   │   ├── fact_096
    │   │   │   ├── fact_097
    │   │   │   ├── fact_098
    │   │   │   └── fact_099
    │   │   ├── Workshops & Community Engagement/
    │   │   │   ├── fact_100
    │   │   │   ├── fact_101
    │   │   │   ├── fact_102
    │   │   │   └── fact_103
    │   │   └── Expert Feedback & Revisions/
    │   │       ├── fact_104
    │   │       ├── fact_105
    │   │       ├── fact_106
    │   │       ├── fact_107
    │   │       └── fact_108
    │   ├── Expeditions & Treks/
    │   │   ├── Appalachian Trail Trek/
    │   │   │   ├── fact_109
    │   │   │   └── fact_110
    │   │   ├── Patagonia Expedition Planning/
    │   │   │   ├── Route & Permits/
    │   │   │   │   ├── fact_111
    │   │   │   │   ├── fact_112
    │   │   │   │   ├── fact_113
    │   │   │   │   ├── fact_114
    │   │   │   │   ├── fact_115
    │   │   │   │   └── fact_116
    │   │   │   └── Equipment Budget/
    │   │   │       ├── fact_117
    │   │   │       └── fact_118
    │   │   └── Chile Expedition/
    │   │       ├── Travel Logistics & Preparations/
    │   │       │   ├── fact_119
    │   │       │   ├── fact_120
    │   │       │   └── fact_121
    │   │       ├── Equipment & Resources/
    │   │       │   ├── fact_122
    │   │       │   ├── fact_123
    │   │       │   ├── fact_124
    │   │       │   ├── fact_125
    │   │       │   ├── fact_126
    │   │       │   └── fact_127
    │   │       └── Community Engagement & Research/
    │   │           ├── fact_128
    │   │           ├── fact_129
    │   │           ├── fact_130
    │   │           ├── fact_131
    │   │           └── fact_132
    │   ├── Fundraising & Outreach/
    │   │   ├── Grant Applications/
    │   │   │   ├── fact_133
    │   │   │   ├── fact_134
    │   │   │   └── fact_135
    │   │   ├── Sponsorships & Partnerships/
    │   │   │   ├── fact_136
    │   │   │   ├── fact_137
    │   │   │   ├── fact_138
    │   │   │   └── fact_139
    │   │   ├── Fundraising Events/
    │   │   │   ├── fact_140
    │   │   │   ├── fact_141
    │   │   │   ├── fact_142
    │   │   │   └── fact_143
    │   │   └── Donor Relations & Community Outreach/
    │   │       ├── fact_144
    │   │       └── fact_145
    │   ├── Social Media Campaigns/
    │   │   └── Virtual Trail Clean-up Campaign/
    │   │       ├── Campaign Strategy & Launch/
    │   │       │   ├── fact_146
    │   │       │   └── fact_147
    │   │       └── Video Montage Production/
    │   │           ├── Initial Development & Review/
    │   │           │   ├── fact_148
    │   │           │   ├── fact_149
    │   │           │   ├── fact_150
    │   │           │   └── fact_151
    │   │           └── Black Forest Segment/
    │   │               ├── Segment Development/
    │   │               │   └── fact_152
    │   │               ├── Footage Acquisition Challenges/
    │   │               │   ├── fact_153
    │   │               │   ├── fact_154
    │   │               │   ├── fact_155
    │   │               │   ├── fact_156
    │   │               │   └── fact_157
    │   │               ├── SkyView Productions Engagement/
    │   │               │   ├── fact_158
    │   │               │   ├── fact_159
    │   │               │   ├── fact_160
    │   │               │   ├── fact_161
    │   │               │   └── fact_162
    │   │               └── Segment Outcome/
    │   │                   └── fact_163
    │   └── Project Milestones & Future/
    │       ├── Key Events & Celebrations/
    │       │   ├── fact_164
    │       │   └── fact_165
    │       └── Summit Seekers 2.0/
    │           ├── Vision & Concepts/
    │           │   ├── fact_166
    │           │   ├── fact_167
    │           │   └── fact_168
    │           └── Planning & Collaboration/
    │               └── fact_169
    └── Social Life & Leisure/
        ├── Concerts & Events/
        │   ├── fact_170
        │   ├── fact_171
        │   └── fact_172
        └── Personal Affairs/
            ├── Personal Finances/
            │   └── fact_173
            ├── Personal Tasks & Possessions/
            │   ├── fact_174
            │   └── fact_175
            ├── Gifts & Personal Gear/
            │   ├── fact_176
            │   ├── fact_177
            │   ├── fact_178
            │   ├── fact_179
            │   ├── fact_180
            │   ├── fact_181
            │   └── fact_182
            ├── Property & Housing/
            │   ├── fact_183
            │   └── fact_184
            └── Activities with Miguel Ramirez/
                └── fact_185

类似于一个一开N的树，没什么大问题

Locomo就只有两个Person的树，为了让训练样本多一点，特意让他多造了一些树

警告

这里有个工程经验，用LLM造树最好不要让他从上到下开始造，不然的话，很可能他会直接在你上面的aspect上扩写变成下面的fact，这会导致两个文本相似度过高，有可能模型分不开

Agent Memory Research In MMLab (Part 3)

Tue, 04 Aug 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

文档说明

本文档说明训练的动作空间，设计，损失，概率图模型和优化器

一阶段动作空间

一阶段的决策非常简单，模型只要在一个两层子树上面，找到一个节点把新fact挂上去就行了，鉴于二层子树中第二层的aspect是不确定的所以用transformer来支持变长输入，假设有1个root和K个第二层的aspect，那么输出的logits的长度显然是1+K

得到logits后只要argmax一下就行了

Ground Truth也很简单，因为LLM他也是选一个node挂上去，所以Ground Truth就是一个one-hot的向量，这里的loss设计为他们的CE loss

二阶段动作空间

这里复杂得多，我们假设一些前提：

闭包(Part1中说了闭包是什么)只有第三层，即启动这个闭包的fact的兄弟fact才能被操作
每个可被操作的fact只能被操作一次，不能既参与A操作又参与B操作
记可操作fact的数量为K
记最后一层的aspect的数量为M，显然，M可以为0

有了上面这些假设，就可以做如下的设计了

网络输出三个logits, 如下:

type_logits # [K,4]
down_logits # [K,M]
group_logits # [K,K]

这些logits语义是很明确的，对于K个fact中的每一个，有4种动作决策

对于K个fact当中的每一个，如果要执行demote操作，那么要在M个demote的对象aspect当中选一个

对于K个fact当中的每一个，他可以和其他的fact做group

Ground Truth空间没这么简单，显然对于type_logits，是行one-hot的，demote_logits也是，然后group_logits是对称的0/1矩阵

这样一来的话，一系列的动作就可以用这三个矩阵来描述了

条件loss设计

现在的loss衡量的是：我们模型的输出logits和ground truth差多少，这个差用element-wise来衡量，但这里有一个conditional，说简单点就是如果对于一个fact，GroundTruth不认为他参与了Demote操作，就不应该考虑他的Demote Logits之间的loss，Group同理

所以这个element-wise的MLE loss应该设计为：

$$ \begin{aligned} L = \frac{1}{K} \sum_{i=1}^{K} \Bigg[ & \text{CE}\left( z^{\text{type}}i + \tau \cdot \log \pi,\ t_i \right) \ &+ \mathbf{1}[t_i = \text{DOWN}] \cdot \text{CE}\left( z^{\text{down}}i,\ m_i \right) \ &+ \mathbf{1}[t_i = \text{GROUP}] \cdot \frac{1}{|V_i|} \sum{j \in V_i} \text{BCE}\left( z^{\text{grp}}{ij},\ \mathbf{1}[j \in G_i] \right) \Bigg] \end{aligned} $$

其中 $V_i = {j : j≠i}$

概率图模型，优化器

前面已经说过了，logits空间和ground truth空间不是一个空间，所以说如果要执行logits所代表的操作，需要有一个类似编译器的东西把logits送到一个可执行动作，这就是概率图模型

简单来说是这样的，概率图模型是一个势函数$\phi$，给定一组logits:$\theta$，对于这个logits空间所对应的那个ground truth空间里的任意的动作$\sigma$，总有一个概率$\phi_{\theta}(\sigma)$, 接下来要找使得前面这个概率最大的$\sigma$

所以我们需要:

设计势函数$\phi$
找到优化器, 使得不穷举ground truth空间里面所有的$\sigma$, 而是通过逐步优化的方式提高概率，最后达到OPTIMAL

幸运的是优化器不需要我们自己写，因为ground truth空间是一个闭式描述空间，然后目标函数$\phi$可以被设计成线性的或者对数线性的，所以实际上被转化为一个规划问题

$$ \begin{aligned} \phi_\theta(\sigma) = \prod_{i=1}^{K} &, p_\theta(t_i) \ \cdot \prod_{\substack{i=1 \ t_i = \text{DOWN}}}^{K} &, p_\theta(m_i \mid \text{DOWN}) \ \cdot \prod_{\substack{i=1 \ t_i = \text{GROUP}}}^{K} & \prod_{j \neq i} q_{ij}^{\mathbf{1}[j \in G_i]} \cdot (1 - q_{ij})^{\mathbf{1}[j \notin G_i]} \end{aligned} $$

其中

$$ \begin{aligned} & p_\theta(t_i) = \text{softmax}\bigl( z^{\text{type}}i \cdot \log \pi \bigr)[t_i] \quad \ & p\theta(m_i \mid \text{DOWN}) = \text{softmax}( z^{\text{down}}i )[m_i] \ & q{ij} = \text{sigmoid}( z^{\text{grp}}_{ij} ) \end{aligned} $$

优化器用的是OR-Tools CP-SAT

RL Course Notes (Part 3)

Wed, 06 May 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user';

RL Course Notes (Part 3)

最优状态值和贝尔曼最优方程

上次讲了状态值和动作值的计算方式:

通过一个简单的例子来算一算, 只要注意状态值是自迭代的, 动作值是依赖于状态值的即可

直观上可以这么理解, 对于状态值, 我从现在这个状态出发, 首先有第一个概率决定我下一步走到哪个状态, 这是第一层期望。当我走到这众多状态的某一个后, 我又有一个概率会获得许多不同的reward, 这是第二层期望的第一部分。同时我还有从这个状态出发再走到下一个状态, 所以这是第二层期望的第二部分。

对于动作值, 由于目前所做的动作是确定的了, 所以就没有第一层期望了, 但是注意到就算动作确定了, reward仍然是一个分布, 再下次走到的状态也是一个分布, 所以仍然是有一层期望的两部分。

最优状态值和最优策略

不管你用什么策略, 状态是已经被环境确定了的, 所以可以通过每个状态的状态值来衡量策略的好坏, 若满足下列条件:

$$ v_{\pi}(s) \geq v_{\pi'}(s) \quad \forall s \in \mathcal{S} $$

那么就说策略$\pi$比策略$\pi'$要好, 反之亦然。如果存在一个策略$\pi^*$使得他比其他所有策略好, 那么就说这个策略是最优策略。

对于这个最优策略$\pi^*$, 把他的状态值(这是一个集合, 因为有很多状态)叫做最优状态值

书上给了关于最优策略的四个问题, 分别是:

是否存在
是否唯一
是确定性策略还是随机策略
有没有一个算法去找到最优策略

贝尔曼最优方程

根据之前的贝尔曼最优方程,最大化右侧得到:

这是贝尔曼最优方程, 注意他不是一个恒等式, 还是一个关于$v(s)$和$\pi(a|s)$的方程, 所以以上的四个问题都变成了解这个方程的问题

多元约束条件

书上给了一个多元数量函数的例子

这个例子非常平凡, 首先你要让右边在这个$max ,, y$的条件下达到最大, 这个条件和x无关, 所以根据二次函数性质当然应该取$y=0$, 从而x=1

回到这个最优方程上来, 同样的我们也需要先找一个$\pi(s)$使得右边达到最大, 然后再来解这个方程, 看书上给的第二个初等例子:

题目给了一个比较trick的做法, 当然这里是个连续场景下的条件极值问题, 如果用Lagrange乘数法也可以得到一样的答案。

这个问题本质上和我们的贝尔曼最优方程也是一样的, 因为所有的$\pi(a|s)$求和为1, 所以:

RL Course Notes (Part 2)

Fri, 24 Apr 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user';

RL Course Notes (Part 2)

第二章: 状态值(State Value)和贝尔曼方程(Bellman Equation)

回报与策略

衡量一个策略好与坏的直接方式就是写出从起点到Target的轨迹然后计算回报, 以书上的一个简单例子来说明

图上已经写明了三种策略和Reward了, 注意这三条轨迹都是无穷轨迹, 因为走到了Target之后他们都会不定的原地打转, 所以算回报的时候要算无穷级数的和

计算过程就省略了(不记得了可以去翻翻数学分析算级数的那部分), 结果是

$$ Return_{1} = \frac{\gamma}{1 - \gamma} \ Return_{2} = -1 + \frac{\gamma}{1 - \gamma} \ Return_{3} = -0.5 + \frac{\gamma}{1 - \gamma} $$

所以 $$ Return_{1} > Return_{3} > Return_{2} $$

所以策略1是最好的, 策略2是最差的

计算回报的方法

书上给了一个比较特殊的例子:

以$v_{i}$记状态$s_{i}$的回报, 根据之前的公式"回报为之后的所有奖励求无穷和"可以得到

每一行里面又可以做迭代(注意前面已经证明了在有折扣因子$\gamma$存在的条件下, 级数是绝对收敛的, 所以这个迭代可以进行)

这样一来, 得到一个可解的线性非其次方程组, 可以通过Gauss消元或者矩阵代数的方法去解这个方程组:

实际上就是矩阵方程

$$ v = r + \gamma Pv \ v = (E - \gamma P)^{-1}r $$

这就是这个简单例子的贝尔曼方程, 实际上针对任意的情况, 思路都是一样的, 也就是状态的回报是相互依赖的, 所以总能写成一个可解的矩阵方程的形式

状态值(State Value)

我们先引入一种notation, 用来描述状态之间的转换, 用自然语言描述就是: 在状态$S_t$执行了策略$\pi$输出的动作$A_t$, 然后状态转移到了$S_{t+1}$, 并获得了立即奖励$R_{t+1}$

$$ S_t \xrightarrow{{A_t}} S_{t+1}, R_{t+1} $$

不严格的证明, 这里的$S_{t}, S{t+1}, A_{t}, R_{t+1}$都是随机变量, 因为他们都是从一系列的状态空间当中抽取出来的

所以对于任意的t, 我们可以得到一条从$S_t$出发的轨迹:

$$ S_t \xrightarrow{{A_t}} S_{t+1}, R_{t+1} \xrightarrow{{A_{t+1}}} S_{t+2}, R_{t+2} \xrightarrow{{A_{t+2}}} ... $$

根据前面的定义, t状态下(或者说从t状态出发)的折扣回报(沿着这条轨迹)为

$$ G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} + ... $$

$G_t$也是随机变量, 所以可以对他求期望, 同时注意到$G_t$依赖于$S_t$, 所以可以用以下的条件期望来表示:"在某种状态下执行某种策略所得到的折扣回报的期望"

$$ v_{\pi}(s) = \mathbb{E}[G_t|S_t = s] $$

之所以有期望这个说法是因为就算给定了策略和状态t, 每次的轨迹是一个采样而不是一个确定性的东西(因为策略可能会以概率输出不同的动作)

注意他是不依赖于时间t的, 实际上可以这么想, 无论在什么时间, 只要给定策略, 走到s这个状态, $v_{\pi}(s)$就已经确定了, 实际上他只依赖于策略和状态

贝尔曼方程

前面讲了一个简单例子的贝尔曼方程, 现在进行一般情况下的抽象推导:

$$ G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} + ... \ = R_{t+1} + \gamma G_{t+1} $$

第二个等号的迭代来自于级数的绝对收敛性

再考虑状态值的定义

$$ v_{\pi}(s) = \mathbb{E}[G_t|S_t = s] \ = \mathbb{E}[R_{t+1} + \gamma G_{t+1}|S_t = s] = \mathbb{E}[R_{t+1}|S_t = s] + \gamma \mathbb{E}[G_{t+1}|S_t = s] $$

最后一个等号来自于期望的线性性质, 接下来可以对两项进行拆分:

$$ \mathbb{E}[R_{t+1}|S_t = s] = \sum_{a \in \mathcal{A}} \pi(a|s) \mathbb{E}[R_{t + 1} | S_t = s, A_t = a] \ = \sum_{a \in \mathcal{A}} \pi(a|s) \sum_{r \in \mathcal{R}}p(r|s, a)r $$

这里其实就是两次应用期望的定义, 要搞清楚获得这个$R_{t+1}$其实需要双重的概率, 一个是"做什么动作"的概率, 一个是"这个动作获得什么奖励"的概率, 所以只需要进行两次求和即可

对于第二项也是同样的道理, 先按照状态拆分, 这里需要一点trick, 我们可以给这个期望的条件"加上一个":

$$ \mathbb{E}[G_{t+1}|S_t = s] = \sum_{s' \in \mathcal{S}} \mathbb{E}[G_{t+1}|S_t = s, S_{t+1} = s']p(s'|s) \ = \sum_{s' \in \mathcal{S}} \mathbb{E}[G_{t+1}|S_{t+1} = s']p(s'|s) \ = \sum_{s' \in \mathcal{S}} v_{\pi}(s')p(s'|s) \ = \sum_{s' \in \mathcal{S}} v_{\pi}(s') \sum_{a \in \mathcal{A}} \pi(a|s')p(s'|s, a) $$

套路都是一样的, 只要注意第二个等号是由于马尔可夫性质就好

综合以上两项然后提取出公因子, 就得到一般形式的贝尔曼方程:

接下来要明确贝尔曼方程里什么是已知的, 什么是要求的

$v_{\pi}(s)$和$v_{\pi}(s')$是未知的, 也是最后要解的, 这必须通过联立所有的状态的贝尔曼方程来解决(类似前面例子里面的矩阵代数求解)

$\pi(a|s)$是已知的, 因为他是策略, 是给定的

$p(r|s, a)$和$p(s'|s, a)$是已知的, 因为他是环境模型, 是给定的

$r$是已知的, 因为他是奖励函数, 是给定的

贝尔曼方程和例子当中的直觉

我们通过一些例子, 用直觉来感知贝尔曼方程的理论

上图例子当中, 策略是确定性的, 我们用一些条件概率来刻画策略的行为和状态转移的行为:

以上概率刻画了从$s_1$出发到$s_3$的所有变量, Recall一下贝尔曼方程

取$s = s_1, , s' = s_3$, 等式变为: $$ v_{\pi}(s_1) = 0 + \gamma \cdot v_{\pi}(s_3) $$

光这一个方程是结不出两个状态值的, 所以需要联立所有的贝尔曼方程来解:

$$ v_{\pi}(s_2) = 1 + \gamma \cdot v_{\pi}(s_4) \ v_{\pi}(s_3) = 1 + \gamma \cdot v_{\pi}(s_4) \ v_{\pi}(s_4) = 1 + \gamma \cdot v_{\pi}(s_4) \ $$

注意为什么右侧式子总是如此简单, 因为这里的概率分布非常容易, 所有求和号下总是一项为1其他为0, 所以每个求和号下其实都只剩下一项了

根据上面的四元方程组可以解出四个状态值, 过程略

考虑稍微复杂一点的情况, 如果策略不是确定性的, 而是带概率的, 那么此时一个求和号下可能会有若干项, 但是本质上都是一样的

概率刻画为:

贝尔曼方程组为:

可得解

以上的简单例子说明了从直觉角度去理解贝尔曼方程的方式: s处的状态值取决于现在立刻得到的一个奖励加上未来状态的折扣回报, 即期奖励是动作概率和奖励概率双重加权的求和, 远期奖励是动作概率和状态转移概率双重加权的求和

矩阵形式的贝尔曼方程组

首先引入如下的notation:

$$ r_{\pi}(s) = \sum_{a \in \mathcal{A}} \pi(a|s) \sum_{r \in \mathcal{R}}p(r|s, a)r \ p_{\pi}(s'|s) = \sum_{a \in \mathcal{A}} \pi(a|s)p(s'|s, a) $$

也许这个记法有一点难理解, 其实只要记住一点就好, 就是把求和式子写成变量的函数的时候, 函数只取决于不在求和号底下的那些变量

第一行的变量有$a,s,r$, 而$a,r$都被求和了, 所以是$s$的函数, 第二行同理

所以贝尔曼方程被改写为:

$$ v_{\pi}(s) = r_{\pi}(s) + \gamma \sum_{s' \in \mathcal{S}} p_{\pi}(s'|s)v_{\pi}(s') $$

当左边的$s$取遍状态空间$\mathcal{S}$的时候, 就得到了一个方程组:

$$ v_{\pi}(s_i) = r_{\pi}(s_i) + \gamma \sum_{s_j \in \mathcal{S}} p_{\pi}(s_j|s_i)v_{\pi}(s_j) $$

接下来只要引入向量记号即可:

$$ v_{\pi} = (v_{\pi}(s_1), v_{\pi}(s_2), ..., v_{\pi}(s_{|\mathcal{S}|}))^T \ r_{\pi} = (r_{\pi}(s_1), r_{\pi}(s_2), ..., r_{\pi}(s_{|\mathcal{S}|}))^T \ [P_{\pi}]{ij} = p{\pi}(s_j|s_i) $$

于是贝尔曼方程组被改写为:

$$ v_{\pi} = r_{\pi} + \gamma P_{\pi}v_{\pi} $$

这个方程组是一个线性方程组, 可以解出$v_{\pi}$

前面那个随机策略的例子也可以用这个向量形式一步写出方程组:

贝尔曼方程组的解

最直接的方式是矩阵代数给出的闭式解

$$ v_{\pi} = (I - \gamma P_{\pi})^{-1}r_{\pi} $$

关于为什么上面这个矩阵一定可逆, 书上给出了证明, 或者也可以当作一个高代习题去做

还有一种是可以用于离散场合下的递推法, 给出一个初始值$v_{\pi}^0$, 然后迭代:

$$ v_{\pi}^{k+1} = r_{\pi} + \gamma P_{\pi}v_{\pi}^k $$

最后会有 $$ v_k \rightarrow v_{\pi} ,, as ,, k \rightarrow \infty $$

动作值(Action Value)

回忆一下状态值的定义:

$$ v_{\pi}(s) = \mathbb{E}[G_t|S_t = s] $$

即: 从状态$s$出发, 按照策略$\pi$所得到的折扣回报的期望

动作值的定义与之类似, 只不过是多了一个动作$a$的维度:

$$ q_{\pi}(s, a) = \mathbb{E}[G_t|S_t = s, A_t = a] $$

注意条件概率的定义, 动作值是状态和动作的函数, 实际上动作值和状态值之间有以下关系:

$$ v_{\pi}(s) = \mathbb{E}[G_t|S_t = s] = \sum_{a \in \mathcal{A}} \pi(a|s) \mathbb{E}[G_t|S_t = s, A_t = a] = \sum_{a \in \mathcal{A}} \pi(a|s) q_{\pi}(s, a) $$

注意状态值的另一种表达

所以对比一下就可以得到:

$$ q_{\pi}(s, a) = \sum_{r \in \mathcal{R}} p(r|s, a)r + \gamma \sum_{s' \in \mathcal{S}} p(s'|s, a)v_{\pi}(s') $$

所以本质上动作值就是状态值脱掉了外面那层动作的权重, 其余的求和都是一样的, 这也就意味着如果我们知道了状态值, 可以轻松的求到动作值

动作值的贝尔曼方程

将动作值与状态值的方程中的状态值替换成动作值的表达方式:

$$ q_{\pi}(s,a) = \sum_{r \in \mathcal{R}} p(r|s,a)r + \gamma \sum_{s' \in \mathcal{S}} p(s'|s,a) \sum_{a' \in \mathcal{A(s')}} \pi(a'|s')q_{\pi}(s',a') $$

注意$(s,a)$的取值范围是状态空间和动作空间的笛卡尔积, 所以这里换成向量形式的时候, 都是$Card(\mathcal{S}) \times Card(\mathcal{A})$维的向量

$$ q_{\pi} = \tilde{r} + \gamma P \prod{q_{\pi}} $$

其中 $q_{\pi}$ 是动作值向量，下标$(s,a)$表示状态动作对；$\tilde{r}$ 是即时奖励向量，$ \tilde{r} = \sum_{r \in \mathcal{R}} p(r|s,a)r$；$P$ 是状态转移矩阵，$[P]{(s,a),s'} = p(s'|s,a)$；$\Pi$ 是块对角矩阵，每块为对应状态的策略向量 $[\Pi]{s',(s',a')} = \pi(a'|s')$。

RL Course Notes (Part 1)

Tue, 21 Apr 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user';

RL Course Notes (Part 1)

课程简介

本课程使用的是西湖大学赵世钰老师的 "强化学习的数学原理" 课程, B站链接为官方课程链接, 书的官方仓库为课本官方repo

本课程为理论课, 几乎无代码, 实际上现在学RL的coding比理论也快得多, 毕竟天天让Agent写RL的各种组件, 看多了总归是"熟读唐诗三百首, 不会吟诗也会吟"

学习动机

最近在FeelingAi做Research Intern, 第一篇文章的Idea(向我的带教潘薪宇博士致敬)就是关于如何使用RL让Agent自己决策Agent Memory的"增删查改", 事实上我觉得很多算法的idea都是比较直观的, 但是在设计RL的过程中, 我对其背后原理没什么理解, 于是过来补课

之前我有看过一些别人写的notes, 其实RL用到的数学知识并不难(无非工科一二年级三件套), 问题就是notation写的太垃圾让人看不懂. 我感觉工科Researcher并不在乎他们写的equation 是否真的能让人轻松看懂(可能确实也没什么人去看毕竟paper里面放一个equation会减少50%的读者, 再放一个会再减少50%), 或者因为本身equation就不多, 大不了别人看不懂的时候他再去解释, 我觉得这是个非常不好的习惯, 尤其是对于Junior Researcher和想学习算法背后数学原理的人来说, 会让学习效率大打折扣。

最后不得不说, 能把RL的math notation写的看起来比PDE还复杂, 确实是有点艺术细胞在身上的。

第一章: 基本概念

例子: 机器人走格子

RL当中有一些概念, 用一个经典的例子:grid world example(机器人走格子)来解释

机器人从Start出发, 每次走一格, 那显然在每一格上都有很多的选择, 比如往上下左右走一格, 或者原地保持不动, 在RL中把这个可以自己做出决策(动作)的机器人叫做Agent

有一个初始状态(Start), 一个最终目标(Target)和一些我们不希望Agent(用来指代这个机器人)走进的地方

我们的RL算法就是要找到, 或者说推理出一个好的动作序列, 让Agent从Start走到Target

在这个例子里面, "好"是比较好衡量的, 明显Agent如果不撞墙, 不走到Forbidden Area里面并且尽快的走到Target就是好的。这个例子里的格子叫做Environment(环境), 显然Agent是知道环境的全貌的, 从而可以轻松的弄出一个好的算法, 但如果不知道环境的全部, 而只是能在和环境交互(比如走到某个新的格子里面)的时候得到一个反馈, 那要找一个最优算法也许就不容易了。

概念: State(状态)和Action(动作)

抽象的来说, Agent会在环境中处于不同的"位置", 我们可以认为总存在一个环境的有限划分, 然后每一个划分中的元素就是一个State(状态)

所有的State合起来就是集合$\mathcal{S}$, 比如说在这里9个位置的$\mathcal{S} = {s_1 ... s_9}$

动作集合指的是Agent能干的事情, 比如说可以上下左右的移动, 或者原地不动, 把这五个动作编号一下就得到$\mathcal{A} = {a_1 ... a_5}$

注意这里的notation, 显然动作是状态的映射, 在不同的状态上能做的动作是不同的, 比如在这个case下Agent在最下面一行他就不能再往下走了, 又比如说$\mathcal{A(s_1)} = {a_2, a_3, a_5}$

状态之间的变换

显然在一个状态下如果做出一个动作, 状态会改变(当然也有可能维持原有的状态, 比如动作是原地不动)

如果动作和状态的变化是确定的, 即对于任意的状态a执行任意的(动作空间里面的)动作b后状态变为c, 那么非常容易的可以画出以下的变化矩阵:

可惜不是所有的状态/动作变换都是确定的, 有这么一种可能, 在状态a执行动作b后有p的概率状态变为c, 1-p的概率状态变为d, 在RL中我们用以下的notation

$$ \mathcal{P}(c|a,b) = p \ \mathcal{P}(d|a,b) = 1 - p $$

这个notation给人感觉不是那么的直观, 所以要牢牢地记住

策略(Policy)

策略就是一个状态的函数$$\pi$$, 给定一个状态s, 策略会给出在这个状态下应该采取的动作a, 即$$\pi(s) = a$$, 或者至少给出采取动作a的概率$$\pi(a|s)$$

上图是一个确定性的策略, 非常直观, 当我们知道我们现在在哪个格子的时候, 策略就告诉我们下一步往哪里走

一个随机的策略如下图所示, 和状态那里一样, 同样用条件概率来表示采取各个动作的概率

实际上就是说, 在位于左上角那个格子的状态的时候, 两种可能的动作各占一半可能, 在这个机器人爬格子的例子下, 因为我们对所有的状态统一了动作空间(上下左右和不动), 所以可以用一个二位矩阵来表示每个状态下可能的动作和他的概率

奖励(Reward)

奖励是人为设计的一个, 对已执行的动作产生反馈的函数, 他是用来评估Agent在某个状态下做出的某个动作的好坏的

$$ R = R(s,a) $$

在例子里面可以随便设计符合逻辑的奖励, 比如说在边界状态下试图跨越边界, 给-1作为消极奖励, 如果从某个状态走到了Target, 给1作为积极奖励

同样的, 我们用条件概率去表示奖励, 比如说我想表达在在状态a执行动作b后得到奖励r, 可以写成:

$$ \mathcal{P}(R = r|a,b) = 1 $$

同样的可以画出矩阵, 来表达在状态s下执行动作a得到什么奖励

轨迹(Trajectory)

轨迹是一个(状态, 动作, 奖励)的序列, 用自然语言描述就是:

在状态$$A_{i}$$执行动作$$B_{i}$$并且获得奖励$$R_{i}$$,其中i是轨迹的每一步的index

对于一个确定性的策略, 即 $$ \forall i , ,P(\pi(s_{j}|a_{i})) = 1 \ P(\pi(s_{k}|a_{i})) = 0 , ,, \forall k \neq j \ \forall i $$

给定一个起始的状态, 轨迹就是确定的了

回报(Return)

回报指的是沿着一条轨迹从头到尾走完, 把每一步的Reward求和得到的最后数值, 一般来说用回报来衡量一个策略的好坏

有个数学问题是, 如果轨迹无限长, 有可能部分和是发散的, 为了确保回报收敛需要引入一个折价因子$$\gamma \in (0,1)$$, 例如

$$ Reward = \sum_{i=0}^{\infty} \gamma^i R(s_{i}, a_{i}) $$

这个数项级数显然是绝对收敛的, 因为只要取 $$ M = max_{i}|R(s_{i}, a_{i})| $$

那么

$$ {\sum_{i=0}^{\infty} |\gamma^i R(s_{i}, a_{i})|} \leq \sum_{i=0}^{\infty} |\gamma^i| \cdot M = \frac{M}{1-\gamma} $$

证明完毕

马尔可夫决策过程(Markov decision processes)

一个马尔可夫决策过程(MDP)是上面这些概念的大杂烩:

状态空间: $\mathcal{S}$
动作空间: $\mathcal{A}(s), 是某个状态的函数$
奖励集合: $\mathcal{R}(s,a), 是某个状态和动作的函数$

除此之外还有两个模型(用条件概率表示):

第一个条件概率度量在状态s执行动作a后变化到其他状态s'的概率$p(s'|s,a)$, 显然有以下概率之和的恒等式成立:

$$ \sum_{s' \in \mathcal{S}} p(s'|s,a) = 1 $$

第二个条件概率度量在状态s执行动作a后获得奖励r的概率$p(r|s,a)$, 显然有以下概率之和的恒等式成立:

$$ \sum_{r \in \mathcal{R}} p(r|s,a) = 1 $$

MDP还具备一种性质,即当前状态$$s_t$$只和前一个状态$$s_{t-1}$$以及前一个动作$$a_{t-1}$$有关, 而和更早的状态和动作无关, 即:

$$ P(s_t|s_{t-1},a_{t-1}) = P(s_t|s_{t-1},a_{t-1},s_{t-2},a_{t-2},...,s_1,a_1) \ P(r_t|s_t,a_t) = P(r_t|s_t,a_t,s_{t-1},a_{t-1},...,s_1,a_1) $$

注意一个事实, 在RL中模型即概率(条件概率), 所以以上两个条件概率 $$ p(s'|s,a) ,, p(r|s,a) $$ 其实就是模型本身

马尔可夫过程(Markov Process)

当MDP里面的策略被固定后, 就退化成一个马尔可夫过程(MP), 此时所有的状态转移概率和奖励-动作概率都已经确定, 下面这个图是一个例子:

图中状态之间的转移概率已经确定, 因为策略已经确定了, 所以每个状态上做动作的分布也确定了, 即一切的概率都确定了

如何从零轴之下重回水上---外语类保送生转行指南

Thu, 26 Mar 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user' import { Aside } from 'astro-pure/user';

写在前面

这篇文章是我个人从2019年上本科以来, 如何爬出外语这个大坑的经历总结, 在此过程中我有很多失误的地方, 现在归纳出来希望大家能够在每一次选择的时候都尽量选对。

当然, 有些同学是非常适合学习外语的, 社会也非常需要学外语学的好的同学, 所以我也不希望有本身非常适合走这条路的同学被这篇文章误导。

前置知识

我假设这篇文章的读者对外语保送制度有一定的了解, 并且对国内高等教育的体系及升学/就业也略懂, 对于本科生常见的出路: 就业/推免硕士/海外硕士/国内直博/海外直博/就业/统一考试硕士至少了解基本概念

书写顺序

我会通过时间顺序来写这篇文章, 换言之, 我会从保送生资格考试到大四的时间线来描述我大概做了什么, 我当时是什么想法, 如果站在今天的角度回看我应该做什么不该做什么, 也许大家可以从这些过往经历当中有个参考, 从而知道自己该干什么

个人情况简介

过去(中学)

中学就读于NCFLS, 从进这个学校的第一天开始就不怎么学习(沉迷于LOL)并且成绩一直不好(到后面几乎变成重点班倒数30%), 最后大概认真学了五六个星期, 以年级前10%左右的成绩入围保送(大概20%可以入围)

成绩上比较偏科, 数理化生比较好, 语文英语非常烂(基本不学, 英语感觉到高中毕业还是小学六年级的课外班新概念二水平)

如果是初中生看到这篇文章(虽然这不太可能), 我会建议你好好学习然后去高考/海本, 外保有点类似期货交易, 在高考分数上允许你用一定的杠杆(量化的来说可能是以普通985的分比肩华五), 但是非常容易爆仓变成路边一条

过去(校考)

一共考了三所学校:

UIBE对外经贸, 2019年1月6日, Rej, 这个是家里帮报名的, 我自己根本不想去(无法忍受北方的宿舍)
UESTC电子科大, 2019年1月10日, Ad, 最后去处, 当时这个算是比较拉的985, 考的原因是他可以"送"你一张CS的学位证书
ZJU浙江大学, 2019年1月16日, Rej, 没什么好说的, 考的人太多, 和强的人差距太大, 面试的时候和南外杭外的人比一下感觉是原始人

当年的这个选校策略基本上是完全错误的, 我们从宏观和微观的两个角度来讨论.

宏观

我们只讨论>=2015年以后的事情, 因为以往是可以转专业的, 完全没有可比性(想想TOP校的外院和CS/EE就知道了)

交大巴院是我认为T0级别的项目, 因为他真的是让你去学工科, 而且交大牌子够硬, 从身边统计学来说有这个项目的学生去Citadel工作, TOP2外院(尤其是什么N年招一次的小语种)大概率没出过这样的人

TOP2就属于各凭本事, 因为资源太好, 但是具体能不能争取到, 那又是另一回事

其他华五(科大不招生)都差不多, 能去则去, 年年都会冒出一些特异样本点, 比如ZJU我知道有从外院直博到数院去的

剩下的985里面没啥可挑的, 哪里给学双学位的机会就去哪里, 2019年明确能给的有HIT/UESTC/HUST/NKU, 当时北航应该也有(记不清了), 最后在HIT/UESTC/北航三选一选了UESTC(都是同一天考试)

对于喜欢外语并且乐于从事和外语相关的工作的同学, 除了TOP2和华五, 无脑北外上外就行了

微观

我认为最该去考的是交大巴院, 当时之所以考浙大而没有报名交大的原因仅仅是考试时间比较相近, 但其实就算是开个汽车从上海到杭州也绝对够, 虽然说交大不见得考得上(但怎么说考英语+数学也比浙大的语数英好得多), 但至少这个决策逻辑是绝对错误的

关于浙大, 那年考浙大的人特别多, 而且浙大的招生改革从以前两个学院(外院+竺院)改成单外院招生并且不递补, 导致录取率急剧下降, 其实当时报名的时候看到那个报名号的人数就不该去考试, 只能说钻空子的本领还是没学到位(学校并不会公布报考人数, 但可以从报名号看出来)

关于最后的决赛圈三选一(电子科大/北航/哈工)感觉就没啥了, 我觉得去哪个都差不多, 还是那句话, 985都有资源, 但是进去了你能不能拿得到就不好说了

过去(本科)

本科主修的是CS和外语, 当然外语是几乎没学, 成绩平平, 最后发动黑科技上了一堆数院的数学课拿到学分, 考研来的UTongji

总的来说就没做对过几次正确的选择, 我觉得最好的选择是一年级把基础打好, 二年级开始找学校的老板去lab里面打杂, 三四年级继续做RA/科研类的实习拿Connection最后去直博海外, 还考个勾八的硕士。

总结就是一句话: 非必要不考研, 如果那个硕士对你直接就业没什么帮助又不是什么强组强Connection对申博有直接影响, 对于985本的人来说本科毕业赖在学校实验室里做RA都比你读研好得多, 浪费时间读三年硕士还不如做两年RA直接申phd弯道超车, 等其他人35岁被迫下机, 求着老板让自己每天工作12个小时都没机会的时候你已经NIW到手在Google美美拿股权了

现在(硕士)

硕士阶段主要还是学CS和Math, 也做过一些量化的实习, 但是最后还是决定转LLM并且申请Ph.D

未来(Ph.D)

Nobody Knows. 未完待续

本科阶段

一年级

一年级主要是学习微积分, C语言和英语精读。感谢那个(大概自己都没明白指针和内存地址)的老师, C语言学的并不好, 除此之外没有做太多其他的事情。

第一学期其实可以学习CS61A和离散数学以及微积分, 下学期可以学习CS61B。可惜我这个倒霉蛋没有在某种机缘巧合之下点击到csdiy.wiki这个网站, 不然天天刷lab早点接触科研也许Career就会不痛了

二年级

二年级一遍应付各种拖后腿的外语课一边学习数据结构、网络和组成原理, 在第二学期的时候做了一件(也许是四年唯一一件)有意义的事情, 就是到数学学院通过求神拜佛的方式让他们允许我去参加数学系的课程并获得考试资格, 要知道在此之前我的成绩单上只有一个微积分(文科)的成绩, 从此我每个学期都会上两三门数学课, 并最终修读完了数学系的绝大部分本科核心课程(从数学分析到实变函数), 虽然成绩一般般但是也总是比没有好了。

在大学里碰到贵人的概率约等于你的大学里每年直博到USNews前50的人数/该级本科生总人数, 所以你基本上可以认为没人会帮到你, 实际上你想到一件有意义的事情(比如说你的数学课少了, 你就去补修)就要赶快去做, 不要去指望有哪个老师会主动邀请你或者暗示你去做这个事情, 大家都是混口饭吃, 多一事不如少一事, 你脑子里想的是如果我上了这个课有了学分对于我以后申请/升学会有帮助, 他脑子里想的是如果你挂科了还影响毕业和他的业绩挂钩

当然了, 这种事情也分人, 如果你是CS选EE/数学选CS这种当然是不太有所谓, 大家应该也能猜到在工科大学里面外语学生是什么地位, 当一个外语学生走到数学学院的教务科并且告诉staff他认为他可以修读并且通过数学分析考试的时候, 他大概率会觉得这个学生疯了(后面我发现这种数学课一般挂科率还是不低于5%的, 他们的担忧或许也可以理解)

二年级可以当作一个新的开始, 如果是愿意读phd的同学, 我觉得可以大胆的给学校里的老板(尤其是年轻的AP)去发邮件寻求科研机会, 因为实际上没什么人会拒绝打白工的请求, 何况还是在同一个Campus里面。如果是愿意直接工作的同学, 这时候就应该去寻找一些实习机会, 并不一定是非要大厂, 对于算法的实习可以从一些Startup开始做起, 开发的实习在学完了OOP之后就可以去小厂开始干活了

主要的差距都是在大二拉开的, 其实除了那些CMO/NOI的选手, 大部分人不过就是普通人, 一个从大二开始做Research的人, 可能四年级的时候就已经有了一篇会议一作, 一个从大二开始写Cpp工程的人, 可能四年级的时候就能独立Handle一个后端项目了, 但是由于信息差的缘故, 总是看到他们最后的那个大佬结局而忽略了之前的成长过程

实际上大部分这种人只不过是有那么一点小小的狗运, 这种狗运的来源是多种多样的, 可能是家里有个亲戚在当AP/读Phd或是从事互联网行业, 甚至只是鼠标不小心点错了点到一个经验贴里面去开启了新的世界线, 如果你能知道这些路径, 你也可以

LLama2 Baby SFT (Part1)

Sun, 01 Mar 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user';

LLama2 Baby SFT

项目简介

这个项目是关于一个llama2模型的预训练和微调的, 代码仓库在:baby-llama2-chinese

目前只有两步, 预训练和SFT指令微调, 是比较好的入门实验

Part1的内容是对数据集的理解, 清理数据的工程和分词的调用

硬件配置

单卡 32G显存
内存90G

文件结构

baby-llama2-chinese/
├── model.py            # LLaMA2 模型结构与推理实现
├── pretrain.py         # 预训练脚本（百度百科 / 维基语料）
├── sft.py              # 有监督微调（SFT）脚本
├── eval.py             # 使用微调后模型的推理 / 简单评测脚本
├── eval_pretrain.py    # 预训练阶段模型的评估脚本
├── dataset.py          # 预训练数据集定义与加载
├── dataset_sft.py      # SFT 数据集定义与加载
├── data_process.py     # 预训练语料预处理脚本
├── sft_data_process.py # SFT 数据预处理脚本
├── chatglm_tokenizer/  # ChatGLM 分词器相关文件
├── data/               # 预训练语料与中间数据
├── data_clean/         # 数据清洗、日志等辅助工具
├── sft_data/           # SFT 微调所用的指令/对话数据
├── out/                # 训练输出目录（checkpoint、日志等）
├── requirements.txt    # Python 依赖列表
├── README.md           # 项目说明文档

数据集

数据集分为两块, 预训练数据和微调数据

预训练数据

我们对两个数据集分别举例看一看, 一个是百度百科, 一个是Wiki中文百科, 前一个是.bin的已经分词处理之后的语料, 后一个是.json的原始文本, 下载链接可以在原始仓库的readme里面找到

对于被分词后的.bin语料, 要用np.uint16格式读取, 读出来是Vocab ID, 即嵌入前的Token ID



# 读取为 int32 类型的 token ID 序列
data = np.fromfile('./../data/baidubaike_563w_1.bin', dtype=np.uint16)

# 使用
print(data[:100])  # 前100个token
print(f"Total tokens: {len(data)}")

[30910 32632 32357 31211 32632 32357 33662 32357 54541 32632 31201 57600
 32632 54746 57897 42290 32357 31155 34699 31813 31123 36776 54727 32632
 32357 54568 32883 36640 31155 32632 32357 54536 54997 57818 56194 31201
 44385 31201 38146 31201 54997 54581 54902 57409 31201 54997 54722 31301
 47755 32217 54997 33461 31201 44578 31201 44016 31201 54729 40964 31201
 54997 55055 31201 53565 54609 31155 35284 31965 55478 55210 54642 43505
 54542 34177 37746 33403 31123 54570 43647 54609 33739 32222 54536 34406
 31827 31123 35400 54548 35635 54706 31123 54536 31917 40968 31123 54853
 55822 33052 31123 32316]
Total tokens: 721068248

我们还可以用项目里给出的tokenizer来把Token ID解码成自然语言, 注意一下路径问题

import sys
sys.path.append('./..')  # 添加上级目录到路径
from chatglm_tokenizer.tokenization_chatglm import ChatGLMTokenizer



# 加载 tokenizer
tokenizer = ChatGLMTokenizer(vocab_file='./../chatglm_tokenizer/tokenizer.model')

# 取前 100 个 token 试试
token_ids = data[:100].tolist()

# 解码成文本
text = tokenizer.decode(token_ids)
print(text)

红色食品：红色食品是指食品为红色、橙红色或棕红色的食品。科学家认为，多吃些红色食品可预防感冒。
红色食品有红柿椒、西红柿、胡萝卜、红心白薯、红果（山楂）、红苹果、草莓、红枣、老南瓜、红米、柿子等。 
有治疗缺铁性贫血和缓解疲劳的作用，对乳腺癌等肿瘤疾病有防治作用，给人以兴奋感，有增加食欲，光洁皮肤，增强

同理去读取.json的原始文本

import json

with open('./../data/wikipedia-cn-20230720-filtered.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

print(data[:3] if isinstance(data, list) else data)  # 打印前几条看看结构

[{'completion': '昭通机场（ZPZT）是位于中国云南昭通的民用机场，始建于1935年，1960年3月开通往返航班“昆明－昭通”，
原来属军民合用机场。1986年机场停止使用。1991年11月扩建，于1994年2月恢复通航。是西南地区「文明机场」，通航城市昆明。 
机场占地1957亩，飞行区等级为4C，有一条跑道，长2720米，宽48米，可供波音737及以下机型起降。机坪面积6600平方米，停机位2个，
航站楼面积1900平方米。位于城东6公里处，民航路与金鹰大道交叉处。\n航点\n客服电话\n昭通机场客服电话：0870-2830004', 
'source': 'wikipedia.zh2307'}, {'completion': '我的英雄学院：英雄新世纪\n《我的英雄学院剧场版：英雄新世纪》（仆のヒーローアカデミア THE MOVIE ヒーローズ:ライジング）
是一部于2019年12月20日上映的日本动画电影，由长崎健司执导、黑田洋介编剧，改编自日本漫画家堀越耕平创作的漫画系列《我的英雄学院》，同时也是其系列第二部电影版。
\n概要\n本作电影内容同样为作者堀越耕平监修的原创故事，并表示「这部电影版某种意义上可以说是《我的英雄学院》的结局了」。\n登场角色\n制作人员\n主题曲\n 「ハイヤーグラウンド」\n 
作词：片冈健太，作曲：黑田隼之介，主唱：sumika\n跨媒体展开\n集英社亦推出该片的小说版和文库版，由小说家誉司アンリ执笔著作，小说版于2019年12月20日上市。
\n 仆のヒーローアカデミア THE MOVIE ヒーローズ:ライジング\n 仆のヒーローアカデミア\u3000THE\u3000MOVIE\u3000ヒーローズ 
: ライジング\u3000ノベライズ\u3000みらい文库版', 'source': 'wikipedia.zh2307'}, {'completion': '黄大仙文化公园（Wong Tai Sin Culture Park）是香港一个公园，
位于九龙黄大仙摩士公园，门牌编号为香港黄大仙区竹园大成街8号，为一个公园中的公园，入口设于东头村道及大成街。公园原址为摩士二号公园的苗圃，后来由黄大仙区民政事务总署拨款1,370万港元，
改建为一个以中国文化为主题的公园，于2008年10月落成启用。\n设施\n* 无极广场：铺砌《易经》、八卦及太极符号\n* 诗墙：刻有书法诗篇\n* 百年古井：原为石鼓垄村的水井', 
'source': 'wikipedia.zh2307'}]

可以看到他的结构如下:

{'completion': '文本内容', 'source': '数据来源'}

知道原始数据的结构我们才能在后面把它处理成prompt的形式

微调数据

微调数据都是.json格式, 打印出来看一下就行

import json

# 方法1：一次性读取（适合小文件）
with open('./../sft_data/alpaca_gpt4_data_zh.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

print(data[:3] if isinstance(data, list) else data)  # 打印前几条看看结构

[{'instruction': '保持健康的三个提示。', 'input': '', 'output': '以下是保持健康的三个提示：\n\n1. 保持身体活动。
每天做适当的身体运动，如散步、跑步或游泳，能促进心血管健康，增强肌肉力量，并有助于减少体重。\n\n2. 均衡饮食。
每天食用新鲜的蔬菜、水果、全谷物和脂肪含量低的蛋白质食物，避免高糖、高脂肪和加工食品，以保持健康的饮食习惯。\n\n3. 
睡眠充足。睡眠对人体健康至关重要，成年人每天应保证 7-8 小时的睡眠。良好的睡眠有助于减轻压力，促进身体恢复，并提高注意力和记忆力。'}, 
{'instruction': '三原色是什么？', 'input': '', 'output': '三原色通常指的是红色、绿色和蓝色（RGB）。它们是通过加色混合原理创建色彩的三种基础颜色。
在以发光为基础的显示设备中（如电视、计算机显示器、智能手机和平板电脑显示屏）, 三原色可混合产生大量色彩。其中红色和绿色可以混合生成黄色，红色和蓝色可以混合生成品红色，
蓝色和绿色可以混合生成青色。当红色、绿色和蓝色按相等比例混合时，可以产生白色或灰色。\n\n此外，在印刷和绘画中，三原色指的是以颜料为基础的红、黄和蓝颜色（RYB）。
这三种颜色用以通过减色混合原理来创建色彩。不过，三原色的具体定义并不唯一，不同的颜色系统可能会采用不同的三原色。'}, 
{'instruction': '描述原子的结构。', 'input': '', 'output': '原子是物质的基本单位，它由三种基本粒子组成：质子、中子和电子。质子和中子形成原子核，
位于原子中心，核外的电子围绕着原子核运动。\n\n原子结构具有层次性。原子核中，质子带正电，中子不带电（中性）。原子核非常小且致密，占据了原子总质量的绝大部分。
电子带负电，通常围绕核运动，形成若干层次，称为壳层或电子层。电子数量与质子数量相等，使原子呈电中性。\n\n电子在每个壳层中都呈规律分布，
并且不同壳层所能容纳的电子数也不同。在最里面的壳层一般只能容纳2个电子，其次一层最多可容纳8个电子，再往外的壳层可容纳的电子数逐层递增。
\n\n原子核主要受到两种相互作用力的影响：强力和电磁力。强力的作用范围非常小，主要限制在原子核内，具有极强的吸引作用，使核子（质子和中子）紧密结合在一起。
电磁力的作用范围较大，主要通过核外的电子与原子核相互作用，发挥作用。\n\n这就是原子的基本结构。原子内部结构复杂多样，不同元素的原子核中质子、中子数量不同，
核外电子排布分布也不同，形成了丰富多彩的化学世界。'}]

可以看到他的结构如下:

{'instruction': '提示内容', 'input': '输入内容', 'output': '输出内容'}

无input因为这不是一个续写问题, 这里需要的只是模型直接回答问题

数据清洗(Optional)

/data_clean/clear.py里面有一大堆数据清理的函数, 用来清理原始的.json语料, 不过并不是每个函数都用到了, 我们这里详细讲一下process_baike()函数对百度百科格式数据的处理

储存格式

处理完毕后, 会储存为.parquet格式, 等待进一步被Tokenizer处理

删除输出目录

如果输出目录存在, 则先询问用户是否要删除

# clear.py
def delete_file(file: str)-> bool:
    '''
    询问删除文件
    '''
    if exists(file):
        ans = input('delete file: {} ? Yes (y) or No (n)'.format(file))
        ans = ans.lower()
        if ans in ('yes', 'y'):
            remove(file)
            print('deleted.')
            return True
    return False

再从process_baike()函数里面调用这个函数

# clear.py
def process_baike(response_less_word: int=15) -> None:
    file_names = [
        '../data/563w_baidubaike/563w_baidubaike.json',
    ]
    save_file_name = '../data/563w_baidubaike/baike.parquet'
    if exists(save_file_name): 
        assert delete_file(save_file_name)

去除重复的标点符号

# clear.py
def remove_duplicate_punctuation(sentence: str) -> str:
    '''
    删除句子中重复的标点符号、重复的空格，同时将换行变为特殊字符'\n'
    '''
    # 将空格（全角空格）替换为逗号, 可能会有重复的空格，下面删除重复标点会删除
    sentence = re.sub(' |　', '，', sentence) 

    ans = ''
    n = len(sentence)
    p = 0
    while p < n:
        ans += sentence[p]

        while p + 1 < n and sentence[p] in punctuation and sentence[p + 1] in punctuation:
            p += 1
        p += 1

    return ans

此函数把逗号变为空格, 同时删除重复的标点符号

接下来还要根据原始json数据的格式去应用这个函数

{
    "title": "文章标题",
    "summary": "文章简介/摘要",
    "sections": [
        {
            "title": "章节标题1",
            "content": "章节内容..."
        },
        {
            "title": "章节标题2", 
            "content": "章节内容..."
        }
    ],
    "tags": ["标签1", "标签2", ...],
    "url": "原始链接"
}

# process_baike()函数中
def process_function(line: str) -> dict:
    
    item = ujson.loads(line)
    item_title = item['title']
    item_sections = item ['sections']
    for data in item_sections:
        #print(item['completion'])
        # 数据清洗
        response = data['content'].replace('\r','')
        response = remove_duplicate_punctuation(response)
        # 剔除短数据
        if len(response) < response_less_word:
            return None
        response = data['title']+data['content']
        write_dict = {
                "response": response,
            }
        return write_dict

这个函数的作用是:

读取原始json数据
提取标题和内容(内容是一个嵌套结构, 所以需要遍历)
应用remove_duplicate_punctuation函数去把内容当中的每一个元素做处理
剔除短数据
返回处理后的数据

逐个处理文件

首先需要集成之前实现的对每个文件的处理函数, 把它当作一个回调函数传入read_and_write_template_baike()函数里面去, 这样这个函数就会一次处理一个文件并且写入.parquet

# clear.py
def read_and_write_template_baike(read_file: str, write_to_file: str, call_back: object, group_cnt: int=10000) -> None:
    '''
    处理数据读写模板，需要提供一个回调函数call_back，
    read_file: 原始数据文件
    write_to_file：处理后的要保存数据文件
    call_back：函数输入一个字符串，输出一个处理后的字典dict，如果输入的字符串为无效数据，请返回None
    group_cnt: parquet file分割行数
    如：
    >>> def call_back(inputs: str) -> dict:
    >>>     if check(inputs) not valid:
    >>>         return None
    ...    
    ...    do something for inputs
    ...
    >>>     my_dict = {
    >>>             'prompt': inputs['p'],
    >>>             'response': inputs['a1'] + inputs['a2'],
    >>>             ...
    >>>         }
    >>>     return my_dict
    '''

    log.info('process file:{}'.format(read_file), save_to_file=True)
    start = time.time()
    
    raw_line_cnt = 0
    keep_line_cnt = 0
    with progress.open(read_file, 'r', encoding='utf-8') as f_read:
        cur_rows = []
        append = cur_rows.append
        for line in f_read:
            try:
                #print(line)
                raw_line_cnt += 1
                write_dict = call_back(line)
                if write_dict is None: continue
                keep_line_cnt += 1
                append(write_dict)
                if len(cur_rows) >= group_cnt:
                    df = pd.DataFrame(cur_rows)
                    write_single_parquet_file(write_to_file, df)
                    cur_rows = []
                    append = cur_rows.append
            except Exception as e:
                # log.error('处理文件异常：{}, content:{}'.format(str(e), line))
                print(line)
                raise e
            # end for
            # 处理末尾部分
        if len(cur_rows) > 0:
            df = pd.DataFrame(cur_rows)
            write_single_parquet_file(write_to_file, df)
            cur_rows = []
        end = time.time()
        log.info('原始文件:{}，共{}行，处理后剩余{}行，保存到文件：{}。耗时：{:.6}s'\
                    .format(read_file, raw_line_cnt, keep_line_cnt, write_to_file, end - start), save_to_file=True)

简而言之就是:

输入文件 (JSON/文本)
    │
    ▼ 逐行读取
┌─────────────────┐
│  call_back()   │ ← 回调函数处理每一行
└────────┬────────┘
         │
    ┌────┴────┐
    ▼         ▼
  None     返回 dict
 (跳过)     (收集保存)
    │         │
    └────┬────┘
         ▼
   累计达到 group_cnt 行
         │
         ▼
   写入 Parquet 文件

然后用一个for循环应用即可

# process_baike()函数
for file_name in file_names:
    read_file = file_name
    read_and_write_template_baike(read_file, save_file_name, process_function)

分词与编码

本项目用的是预训练好的chatglm_tokenizer, 首先理解一下一个分词器能干什么

tokenize(text): 分词
encode(text): 文本转ID
decode(ids): ID转文本

这三个是最核心的功能了, 注意分词是把一串文本分成一些子词, encode是把子词转成ID, decode是把ID转成子词

接下来这个函数就实现如何去编码这个.json语料, 还是以百度百科为例, 首先把title,summary和sections里面的内容拼接成一个str, 然后去encode这个字符串就好了

# data_process.py
def process_baidu():
    BATCH_SIZE = 1000000

    cnt=0
    batch_cnt=0
    token=0
    doc_ids=[]

    f1=open('./data/563w_baidubaike/563w_baidubaike.json','r',encoding='utf-8')
    
    while True:
        line = f1.readline()
        if not line:
            break
        line=json.loads(line)
        text=''
        try:
            text+=line['title']+'：'+line['summary']
        except:
            pass
        for per in line['sections']:
            text+=per['title']+'：'+per['content']+'。'
        text_id=tokenizer.encode(text,add_special_tokens=False)
        text_id.append(tokenizer.special_tokens['<eos>'])
        if len(text_id)>5:
            doc_ids+=text_id
        cnt+=1
        if cnt%BATCH_SIZE==0:
            batch_cnt+=1
            arr = np.array(doc_ids,dtype=np.uint16)
            doc_ids=[]
            print('cnt:',cnt,'arr_shape:',arr.shape)
            with open('./data/baidubaike_563w_{}.bin'.format(batch_cnt),'wb') as f2:
                f2.write(arr.tobytes())
            del arr

    if not doc_ids:
        batch_cnt+=1
        arr = np.array(doc_ids,dtype=np.uint16)
        print('cnt:',cnt,'arr_shape:',arr.shape)
        with open('./data/baidubaike_563w_{}.bin'.format(batch_cnt),'wb') as f:
            f.write(arr.tobytes())

最后保存成.bin的格式

UC Berkeley CS189 Assignment 5(Part 1)

Sat, 28 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 5

项目介绍

这个lab是关于语言模型微调的, 回忆一下在Assignment4的Part2里面大致介绍了微调, 那里的微调和这里的Qwen-0.5B模型微调主要有以下区别:

前面用的是ConvNeXt模型, 参数比较少, 微调比较方便
前面的微调数据是.wav化为的二维频谱图, 这里是自然语言, 处理起来要复杂得多
前面的任务是分类, 这里是在baseline上的QA检测
这里微调的操作空间大得多, 除了冻结参数, 还有很多超参数可以调整

模型配置与加载

基础模型配置

# ============================================================================
# === CONFIGURATION - ALL SETTINGS IN ONE PLACE ===
# ============================================================================

# --- Model Configuration ---
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" # YOU CANNOT CHANGE THIS

# --- Dataset Configuration ---
#TODO: REPLACE WITH YOUR OWN PATH
MCQ_CSV_PATH = "hw5_sample_eval.csv"  # Path to CS189 MCQ sample eval dataset

# --- Training Configuration (feel free to adjust!) ---
TRAIN_BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 4
WARMUP_STEPS = 5
MAX_STEPS = 50  # or set num_train_epochs instead
LEARNING_RATE = 1e-5
WEIGHT_DECAY = 0.01
LR_SCHEDULER_TYPE = "linear"
OPTIM = "adamw_8bit"  # requires bitsandbytes
SEED = 189

# --- Evaluation Configuration ---
EVAL_MAX_NEW_TOKENS = 64  # How many tokens to generate for inference
OUTPUT_DIR = "./mcq_finetuned_model"

首先需要加载一个预训练好的模型, 准备一份评测集去算出绩效的baseline, 再设置各种超参数, 比如batch_size, 学习率等等

评测集一览

所谓评测集其实就是一些单选题, 我们通过模型选对/选错就能评测出绩效了

id,question,A,B,C,D,E,answer
mcq_1,"Peanut wants to train a model to accurately classify different types of animals from images. After training and testing his model, he observes that the model has high training error and high test error. What can we most confidently say about the bias/variance characteristics of Peanut’s model?",High bias.,Low bias.,High variance.,Low variance.,none of the above,A
mcq_2,"Consider a binary classification data set with 9000 positively labelled examples and 1000 negatively labelled examples. What is the area under the ROC curve (AUC-ROC) of a random classifier that classifies any example as positive with probability π and as negative with probability 1 − π? Here, the probability π is a hyperparameter.",Close to zero.,Close to 0.1.,Close to 0.5.,Close to 0.9.,Close to one.,C
mcq_3,"Again, consider a binary classification data set with 9000 positively labelled examples and 1000 negatively labelled examples. What is the precision and the recall of a classifier that always classifies any example as positive?","The precision is 0.1, and the recall is 0.9.","The precision is 0.9, and the recall is 0.1.","The precision is 1.0, and the recall is 0.9.","The precision is 0.9, and the recall is 1.0.","The precision is 0.1, and the recall is 1.0.",D
mcq_4,"Assume we are given X ∈ Rn×d and y ∈ Rn for n > d. The Ridge regression estimator with regularization coefficient λ estimates the weight vector to be  2 2 wˆ=argmin y−Xw2+λ∥w∥2 . (1) w The Ridge regression estimator is equivalent to the ordinary least squares estimator on which of the following modified version of X and y? Id denotes the d × d identity matrix. 0d denotes the all-zero d-dimensional vector, and 1d denotes the all-one d-dimensional vector.","y′ = [ y; 0d ], X′ = [ X; √λ Id ]","y′ = [ y; 1d ], X′ = [ X; √λ Id ]","y′ = [ y; 0d ], X′ = [ X; λ Id ]","y′ = [ y; 1d ], X′ = [ X; λ Id ]",none of the above,A
mcq_5,Which of the following statements are TRUE regarding positive semi-definite and positive-definite matrices?,“Every entry of a matrix is non-negative” is a necessary but insufficient condition for a matrix to be positive definite.,The singular values of a positive semi-definite matrix are the same as its eigenvalues.,"If a matrix A is positive semi-definite, then there exists a matrix B such that BT B = A. (heuristic)",The covariance matrix of any distribution is positive semi-definite and invertible.,"If the Jacobian of a function is positive semi-definite, then the function is convex.",B

这些问题涵盖了很多方面, 比如说数学/物理/生活常识等等

Tokenizer

除此之外, 我们还需要一个Tokenizer, 这同样也是预训练好的, 不然没法把自然语言变成Token

# === Load base model & tokenizer ===
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Ensure we have a pad token for training
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
)
model.resize_token_embeddings(len(tokenizer))
model.to(device)
model.eval()
print('Model loaded.')

提一下上面的Pad Token, 这是为了在训练的时候把不同长度的序列补齐到相同长度

准备微调所需的数据

格式调整

首先我们需要建立Prompt, 在这个QA体系下就是Question + Options的组合, 课程组已经给出了代码

# === MCQ helpers ===
LETTER_SET = set(list("ABCDE"))

def load_mcq_dataset(csv_path: str = MCQ_CSV_PATH):
    """Load the CS189 MCQ dataset.

    Expected columns:
        - question
        - A, B, C, D, E
        - answer (single letter A-E)
    """
    df = pd.read_csv(csv_path)
    required = ["question", "A", "B", "C", "D", "E", "answer"]
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns in MCQ CSV: {missing}")

    df = df.copy()
    df["answer"] = (
        df["answer"]
        .astype(str)
        .str.strip()
        .str.upper()
    )
    df = df[df["answer"].isin(LETTER_SET)].reset_index(drop=True)
    return df

def build_mcq_prompt(row):
    """Prompt for inference: instruction + question + options.

    The model is expected to answer with the correct letter in \\boxed{} format.
    """
    q = str(row["question"]).strip()
    options = "\n".join([
        f"A. {row['A']}",
        f"B. {row['B']}",
        f"C. {row['C']}",
        f"D. {row['D']}",
        f"E. {row['E']}",
    ])
    prompt = (
        "Choose exactly one correct option from A, B, C, D, and E.\n"
        "Return your answer inside a LaTeX box.\n\n"
        f"{q}\n\n{options}\n\nAnswer:"
    )
    return prompt

实际上就是一些字符串的处理, 因为评测集给的格式非常好, 其实选项和答案都已经给出来了, 所以只要把他们装填成这个特定的数据结构就行了

还给了个辅助函数从选项框里面把答案提取出来:

def parse_choice_from_boxed(text: str):
    """Parse an MCQ choice A–E from the model output.

    We first look for a literal '\\boxed{X}' pattern. If not found, we
    fallback to the last standalone A-E in the decoded text.
    """
    if text is None:
        return None
    # Direct \\boxed{A} ... \\boxed{E}
    m = re.search(r"\\boxed\{\s*([A-E])\s*\}", text)
    if m:
        return m.group(1)
    # Fallback: last standalone A–E
    letters = re.findall(r"\b([A-E])\b", text.upper())
    if letters:
        return letters[-1]
    return None

QA示例

# === Load MCQ CSV (Evaluation Data) ===
try:
    mcq_df = load_mcq_dataset(MCQ_CSV_PATH)
    print(f"Loaded MCQ dataset with {len(mcq_df)} rows from {MCQ_CSV_PATH}.")
except Exception as e:
    mcq_df = None
    print("Error loading MCQ CSV — check MCQ_CSV_PATH.")
    raise e
mcq_df

OpenAI Style的Prompt

我们需要把input prompt调整成如下的格式:

{"role": "user", "content": "..."}

模型给我们返回

{"role": "assistant", "content": "..."}

在这个MCQ体系下, 具体为:

{"role": "user", "content": "Choose exactly one correct option... [Question] ... [Options]"}
{"role": "assistant", "content": "\boxed{A}"}

加载微调数据集

我们使用的是MMLU数据集, 主要使用里面的机器学习问题部分

# === MMLU Helper Functions ===
def load_mmlu_dataset(subset: str = "machine_learning", split: str = "test"):
    """Load a subset of the MMLU dataset from Hugging Face."""
    print(f"Loading MMLU dataset (subset={subset}, split={split})...")
    ds = load_dataset("cais/mmlu", subset, split=split)
    return ds

def build_mmlu_prompt(row):
    """Prompt for inference: instruction + question + options."""
    q = str(row["question"]).strip()
    choices = row["choices"]

    options_list = []
    for i, choice in enumerate(choices):
        letter = chr(ord("A") + i)
        options_list.append(f"{letter}. {choice}")
    options_str = "\n".join(options_list)

    prompt = (
        "Choose exactly one correct option from the choices provided.\n"
        "Return your answer inside a LaTeX box.\n\n"
        f"{q}\n\n{options_str}\n\nAnswer:"
    )
    return prompt

def build_mmlu_sft_text(row, tokenizer):
    """Build properly formatted chat template text for training."""
    user_content = build_mmlu_prompt(row)

    answer_int = row["answer"]
    answer_letter = chr(ord("A") + answer_int)
    assistant_content = f"\\boxed{{{answer_letter}}}"

    messages = [
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": assistant_content}
    ]

    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )

# === Load MMLU Machine Learning Dataset ===
mmlu_ds = load_mmlu_dataset("machine_learning", split="test")
mmlu_text_ds = mmlu_ds.map(lambda x: {"text": build_mmlu_sft_text(x, tokenizer)})
print("Loaded MMLU ML dataset with", len(mmlu_text_ds), "rows")

# Set the training dataset - you can mix and match datasets here
train_dataset = mmlu_text_ds

在真实的生产(非学习)环境下, 实际上只要看一下数据集的数据格式, 再确定好Prompt的格式, 然后用字符串处理一步步转过来就行了, 当然这东西感觉自己手写也并不太容易, 因为(至少对我来说)总是忘记字符串处理函数是什么

MMLU示例

<|im_start|>system
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
<|im_start|>user
Choose exactly one correct option from the choices provided.
Return your answer inside a LaTeX box.

Statement 1| Linear regression estimator has the smallest variance among all unbiased estimators. Statement 2| The coefficients α assigned to the classifiers assembled by AdaBoost are always non-negative.

A. True, True
B. False, False
C. True, False
D. False, True

Answer:<|im_end|>
<|im_start|>assistant
\boxed{D}<|im_end|>

Baseline计算

让预训练好的模型去对测试集做出回答, 计算正确率

def eval_mcq_accuracy(
    curr_model,
    curr_tokenizer,
    df,
    max_new_tokens: int = 64,
    return_details: bool = False,
):
    """Evaluate a model on the MCQ dataset using greedy decoding.

    If return_details=True, also return a pandas DataFrame with
    [idx, question, A, B, C, D, E, gold, decoded, parsed, correct].
    """
    curr_model.eval()
    n = len(df)
    correct = 0
    total = 0
    records = []

    for idx in range(n):
        row = df.iloc[idx]
        user_content = build_mcq_prompt(row)

        # Apply chat template for inference
        messages = [{"role": "user", "content": user_content}]
        prompt = curr_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = curr_tokenizer(prompt, return_tensors="pt").to(device)

        with torch.no_grad():
            outputs = curr_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
            )

        gen_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        decoded = curr_tokenizer.decode(gen_tokens, skip_special_tokens=True)

        pred = parse_choice_from_boxed(decoded)
        is_correct = (pred is not None and pred == row["answer"])
        if is_correct:
            correct += 1
        total += 1

        records.append({
            "idx": idx,
            "question": row["question"],
            "A": row["A"],
            "B": row["B"],
            "C": row["C"],
            "D": row["D"],
            "E": row["E"],
            "gold": row["answer"],
            "prompt": prompt,
            "decoded": decoded,
            "parsed": pred,
            "correct": is_correct,
        })

        if (idx + 1) % 20 == 0:
            print(f"Processed {idx + 1}/{n} questions...")

    acc = correct / max(total, 1)
    print(f"MCQ accuracy: {acc * 100:.2f}% ({correct}/{total})")

    details_df = pd.DataFrame(records)
    if return_details:
        return acc, details_df
    return acc

# === Baseline MCQ accuracy before fine-tuning ===
print("Evaluating baseline model on MCQ dataset...")
baseline_acc, baseline_details = eval_mcq_accuracy(
    model,
    tokenizer,
    mcq_df,
    max_new_tokens=EVAL_MAX_NEW_TOKENS,
    return_details=True,
)
baseline_details.head()

The following generation flags are not valid and may be ignored: ['temperature', 'top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
Evaluating baseline model on MCQ dataset...
Processed 20/25 questions...
MCQ accuracy: 28.00% (7/25)

微调

我们用trl(transformer reinforcement learning)来指定微调的参数和启动训练

# === Set up SFTTrainer ===
sft_config = SFTConfig(
    dataset_text_field="text",
    per_device_train_batch_size=TRAIN_BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
    warmup_steps=WARMUP_STEPS,
    max_steps=MAX_STEPS,
    learning_rate=LEARNING_RATE,
    logging_steps=1,
    optim=OPTIM,
    weight_decay=WEIGHT_DECAY,
    lr_scheduler_type=LR_SCHEDULER_TYPE,
    seed=SEED,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    eval_dataset=None,
    processing_class=tokenizer,
)

trainer

# === Fine-tune the model ===
model.train()
trainer.train()
model.eval()

微调后的绩效计算

直接把微调后的模型放到测试集上评估一次就好了

# === Evaluate MCQ accuracy after fine-tuning ===
print("Evaluating fine-tuned model on MCQ dataset...")
ft_acc, ft_details = eval_mcq_accuracy(
    model,
    tokenizer,
    mcq_df,
    max_new_tokens=EVAL_MAX_NEW_TOKENS,
    return_details=True,
)
ft_details.head()
print(f"Baseline acc: {baseline_acc:.4f}, Fine-tuned acc: {ft_acc:.4f}")

Evaluating fine-tuned model on MCQ dataset...
Processed 20/25 questions...
MCQ accuracy: 28.00% (7/25)
Baseline acc: 0.2800, Fine-tuned acc: 0.2800

UC Berkeley CS189 Assignment 5(Part 2)

Sat, 28 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 5

项目简介

这一次我们把微调的训练集从mmlu换成Ceval数据集, 其他的和Part1保持一致, 所以我们要改写Prompt工程的代码以适配Ceval数据集

Ceval数据集

制作微调数据集的Prompt工程

加载


def load_ceval_dataset(subset,split="test"):
    datasets=[]
    print(f"Loading ceval dataset (subset={subset}, split={split})...")

    for subset in subset:
        ds = load_dataset("ceval/ceval-exam", subset, split=split)
        datasets.append(ds)
        
    return datasets

def build_ceval_prompt(row):
    q=str(row["question"]).strip()
    choices=row["choices"]

    options_list=[]
    for i,choice in enumerate(choices):
        letter=chr(ord("A")+i)
        options_list.append(f"{letter}. {choice}")

    options_str="\n".join(options_list)

    prompt = (
        "Choose exactly one correct option from the choices provided.\n"
        "Return your answer inside a LaTeX box.\n\n"
        f"{q}\n\n{options_str}\n\nAnswer:"
    )
    return prompt

def build_ceval_sft_text(row,tokenizer):
    user_content = build_ceval_prompt(row)

    answer_int = row["answer"]
    answer_letter = chr(ord("A") + answer_int)
    assistant_content = f"\\boxed{{{answer_letter}}}"

    messages = [
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": assistant_content}
    ]

    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )

science_subsets = [
    "middle_school_physics","high_school_physics","college_physics",
    "middle_school_chemistry","high_school_chemistry","college_chemistry",
    "middle_school_biology","high_school_biology",
    "middle_school_mathematics","high_school_mathematics","advanced_mathematics",
    "probability_and_statistics",
    "middle_school_geography","high_school_geography",
    "basic_medicine","clinical_medicine","physician",
    "plant_protection","veterinary_medicine",
    "environmental_impact_assessment_engineer",
]

cs_subsets = [
    "college_programming", "computer_architecture", "computer_network",
    "operating_system", "discrete_mathematics", "logic"
]
subset=science_subsets+cs_subsets


# === Load CEVAL Machine Learning Dataset ===
CEVAL_ds=load_ceval_dataset(subset=subset,split='test')

CEVAL_ds = concatenate_datasets(CEVAL_ds)

大体上和之前差不多, 只不过现在是4选1而不是5选1

查看Prompt基本信息

# 基本信息
print("size:", len(CEVAL_ds))
print("columns:", CEVAL_ds.column_names)
print("features:", CEVAL_ds.features)

# 看前 5 条（注意：若很大不要全部 to_pandas）
for i in range(5):
    print(i, CEVAL_ds[i])   # 打印字典形式的单条样本

size: 5195
columns: ['id', 'question', 'A', 'B', 'C', 'D', 'answer', 'explanation']
features: {'id': Value('int32'), 'question': Value('string'), 'A': Value('string'), 'B': Value('string'), 'C': Value('string'), 'D': Value('string'), 'answer': Value('string'), 'explanation': Value('string')}
0 {'id': 0, 'question': '关于信息的传递，下列说法正确的是____', 'A': '北斗卫星定位系统可提供全天候即时定位服务', 'B': '5G网络通信主要是利用光导纤维传递信息的', 'C': '手机话筒的主要作用是把声音信号变成恒定电流', 'D': '电磁波只能传递声音信号，不能传递图像信号', 'answer': 'A', 'explanation': ''}
1 {'id': 1, 'question': '下列说法符合实际情况的是____', 'A': '人的正常体温约为39℃', 'B': '成年人步行的速度约为1.1m/s', 'C': '中学生的体重约为50N', 'D': '一个篮球的体积约为1m3', 'answer': 'B', 'explanation': ''}
2 {'id': 2, 'question': '下列关于测量仪器的分析正确的是____', 'A': '水银温度计利用了液体热胀冷缩的原理', 'B': '托盘天平利用了省力杠杆的原理', 'C': '电能表利用电流的热效应工作', 'D': '液体压强计利用了连通器的原理', 'answer': 'A', 'explanation': ''}
3 {'id': 3, 'question': '关于分子动理论，下列说法中不正确的是____', 'A': '物质是由大量分子组成的', 'B': '温度越高，分子的运动越剧烈', 'C': '分子是组成物质的最小微粒', 'D': '固体很难被压缩，说明分子间存在斥力', 'answer': 'C', 'explanation': ''}
4 {'id': 4, 'question': '有关安全用电，下列做法错误的是____', 'A': '使用验电笔时手要接触验电笔后端金属部分', 'B': '为了使用方便，可以剪掉三脚插头中保护接地线的插脚', 'C': '接入漏电保护器，可以在导线漏电、电器短路等故障时断开电路起保护作用', 'D': '高压线发生断线落地时，人不能靠近落地处', 'answer': 'B', 'explanation': ''}

调整成最终的格式

from datasets import concatenate_datasets
import re

# 如果 CEVAL_ds 是列表先合并
if isinstance(CEVAL_ds, list):
    CEVAL_ds = concatenate_datasets(CEVAL_ds)

# 先把原始 CEVAL 映射为标准的 A/B/C/D/E 字段（如果你之前已有 mcq_ds 可跳过这步）
def to_mcq_format(example):
    return {
        "id": example.get("id", None),
        "question": str(example.get("question", "") or "").strip(),
        "A": example.get("A", "") or "",
        "B": example.get("B", "") or "",
        "C": example.get("C", "") or "",
        "D": example.get("D", "") or "",
        "E": example.get("E", "") or ""  # 保留 E 列（若无则为空）
    }

mcq_ds = CEVAL_ds.map(to_mcq_format)

# 把 A-D(可选E) 合并成 choices 列，并把字母答案转为数值 label（0..）
def add_choices_and_numeric_answer(example):
    choices = [example.get("A",""), example.get("B",""), example.get("C",""), example.get("D","")]
    if example.get("E"):
        choices.append(example.get("E",""))
    m = re.search(r"([A-E])", str(example.get("answer","")).upper())
    label = (ord(m.group(1)) - ord("A")) if m else None
    return {"choices": choices, "answer": label}

mcq_ds = mcq_ds.map(add_choices_and_numeric_answer)

# 过滤掉无法解析到 label 的样本（可选）
mcq_ds = mcq_ds.filter(lambda x: x["answer"] is not None)

# 现在直接调用现有的 build_mmlu_sft_text（它期望 choices 列和 numeric answer）
CEVAL_text_ds = mcq_ds.map(lambda x: {"text": build_ceval_sft_text(x, tokenizer)})

print("Loaded CEVAL (converted) with", len(CEVAL_text_ds), "rows")
print("example text:\n", CEVAL_text_ds[0]["text"])

Loaded CEVAL (converted) with 5195 rows
example text:
 <|im_start|>system
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
<|im_start|>user
Choose exactly one correct option from the choices provided.
Return your answer inside a LaTeX box.

关于信息的传递，下列说法正确的是____

A. 北斗卫星定位系统可提供全天候即时定位服务
B. 5G网络通信主要是利用光导纤维传递信息的
C. 手机话筒的主要作用是把声音信号变成恒定电流
D. 电磁波只能传递声音信号，不能传递图像信号

Answer:<|im_end|>
<|im_start|>assistant
\boxed{A}<|im_end|>

可以看到这已经和之前的格式一样了

计算训练baseline

和之前一样, 直接用他给的eval_mcq_accuracy就行了, 不再赘述

调整训练参数

# === Set up SFTTrainer ===
sft_config = SFTConfig(
    dataset_text_field="text",
    per_device_train_batch_size=TRAIN_BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
    warmup_steps=WARMUP_STEPS,
    # num_train_epochs=NUM_TRAIN_EPOCHS,
    max_steps=MAX_STEPS,
    learning_rate=LEARNING_RATE,
    logging_steps=1,
    optim=OPTIM,
    weight_decay=WEIGHT_DECAY,
    lr_scheduler_type=LR_SCHEDULER_TYPE,
    seed=SEED,
    report_to="none",

    eval_strategy=EVALUATION_STRATEGY,
    eval_steps=EVAL_STEPS,                    # 每50步评估一次
    save_steps=SAVE_STEPS,
    load_best_model_at_end=LOAD_BEST_MODEL_AT_END,
    metric_for_best_model=METRIC_FOR_BEST_MODEL,
    greater_is_better=GREATER_IS_BETTER,
    save_total_limit=SAVE_TOTAL_LIMIT,
)

# early_stopping_callback = EarlyStoppingCallback(
#     early_stopping_patience=10,
# )

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    processing_class=tokenizer,
    # callbacks=[early_stopping_callback],
)

trainer

| 参数 | 说明 |
|------|------|
| `dataset_text_field` | 指定数据集中存储训练文本的字段名，SFTTrainer 会自动从中提取数据进行训练 |
| `per_device_train_batch_size` | 每个 GPU/CPU 设备上的 batch 大小，越大训练越快但需要更多显存 |
| `gradient_accumulation_steps` | **梯度累积**：当显存不够时，用小 batch 模拟大 batch。实际 batch = `batch_size × accumulation_steps` |
| `warmup_steps` | 学习率预热步数，训练初期学习率从 0 逐渐增加到设定值，稳定训练 |
| `max_steps` | 最大训练步数，与 `num_train_epochs` 二选一 |
| `learning_rate` | 学习率，推荐 1e-5 ~ 5e-5，过大容易不收敛 |
| `logging_steps` | 每几步记录一次日志（loss、学习率等） |
| `optim` | 优化器，`adamw_8bit` 是 8-bit AdamW，节省显存 |
| `weight_decay` | 权重衰减（L2 正则化），防止过拟合，通常 0.01~0.1 |
| `lr_scheduler_type` | 学习率调度器：`linear`、`cosine`、`constant` 等 |
| `seed` | 随机种子，保证实验可复现 |
| `report_to` | 设为 `"none"` 关闭 wandb/tensorboard 等远程记录 |
| `eval_strategy` | 评估策略：`"no"`、`"steps"`、`"epoch"` |
| `eval_steps` | 每多少步评估一次 |
| `save_steps` | 每多少步保存一次检查点 |
| `load_best_model_at_end` | 训练结束后自动加载验证集上表现最好的模型 |
| `metric_for_best_model` | 用于选择最佳模型的指标名 |
| `greater_is_better` | 该指标是否越大越好（如 accuracy 是，loss 否） |
| `save_total_limit` | 最多保存几个检查点，超过会删除旧的 |

调整参数是一个经验性的问题, 这里我调整的并不多, 因为主要还是把这个实验跑通而不是在测试集上做的特别完美

# ============================================================================
# === CONFIGURATION - ALL SETTINGS IN ONE PLACE ===
# ============================================================================

# --- Model Configuration ---
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" # YOU CANNOT CHANGE THIS

# --- Dataset Configuration ---
#TODO: REPLACE WITH YOUR OWN PATH
MCQ_CSV_PATH = "hw5_sample_eval.csv"  # Path to CS189 MCQ sample eval dataset

# --- Training Configuration (feel free to adjust!) ---
TRAIN_BATCH_SIZE = 2
GRADIENT_ACCUMULATION_STEPS = 2
WARMUP_STEPS = 5
MAX_STEPS = 500  # or set num_train_epochs instead
NUM_TRAIN_EPOCHS=10
LEARNING_RATE = 5e-5
WEIGHT_DECAY = 0.01
LR_SCHEDULER_TYPE = "cosine"
OPTIM = "adamw_8bit"  # requires bitsandbytes

SEED = 189

# --- Evaluation Configuration ---
EVAL_MAX_NEW_TOKENS = 64  # How many tokens to generate for inference
OUTPUT_DIR = "./mcq_finetuned_model"
EVALUATION_STRATEGY="steps"
EVAL_STEPS=50
SAVE_STEPS=50
LOAD_BEST_MODEL_AT_END=True
METRIC_FOR_BEST_MODEL="eval_loss"
GREATER_IS_BETTER=False
SAVE_TOTAL_LIMIT=3

启动训练

# === Fine-tune the model ===
model.train()
trainer.train()
model.eval()

最终绩效

# === Evaluate MCQ accuracy after fine-tuning ===
print("Evaluating fine-tuned model on MCQ dataset...")
ft_acc, ft_details = eval_mcq_accuracy(
    model,
    tokenizer,
    mcq_df,
    max_new_tokens=EVAL_MAX_NEW_TOKENS,
    return_details=True,
)
ft_details.head()
print(f"Baseline acc: {baseline_acc:.4f}, Fine-tuned acc: {ft_acc:.4f}")

Evaluating fine-tuned model on MCQ dataset...
Processed 20/25 questions...
MCQ accuracy: 32.00% (8/25)
Baseline acc: 0.2800, Fine-tuned acc: 0.3200

UC Berkeley CS189 Assignment 4 (Part 2)

Thu, 26 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 4

Happy Chinese New Year!

项目介绍

本项目是Part1的延续, 利用BERT做一些任务, 同时这次的数据集是语音, 不是之前常见的数据点矩阵

BERT和Tokenizer

什么是BERT

Bi-Directional Encoder Representations from Transformers, 双向的Transformer Encoder, 这个模型是双向的

BERT 中的双向编码器意味着模型一次性处理整个输入序列，同时考虑每个 token 左侧和右侧的上下文。传统的语言模型是从左到右处理序列：当预测位置 $t$ 的 token 时, Transformer 的解码器只能看到位置 1 到 $t-1$ 的 token(这就是为什么我们需要在解码器的自注意力中使用掩码的原因)。与传统的只从左到右读取文本的语言模型不同, BERT 的编码器使用自注意力机制同时查看所有 token, 使其能够学习更丰富的上下文感知表示。因此，每个 token 的表示都受到序列中所有其他 token 的影响，无论这些 token 位于它之前还是之后。这对于模型理解 DNA 或文本序列的完整上下文至关重要

什么是Tokenizers

分词器(Tokenizers)是把文本序列转化为数字序列的算法, 例如在CS336当中实现的BPE分词器就是一种

* Sentence: `"I love dinosaurs"`
* Tokens: `["[CLS]", "I", "love", "dinosaurs", "[SEP]"]`
* Token IDs: `[101, 146, 1567, 4083, 102]` (example values)

可以从上面看到这个Tokenizers的目的就是把每个Token映射到一个数字, 所以分词器的预训练是非常重要的, 显然我们不希望输入的文本当中出现不存在于分词器定义域的Token, 否则模型会无法处理

在CS336中我们自己训练的BPE分词器就容易出现上面这种情况, 因为训练用的语料极其有限(至多也就几个GB的故事集), 但是在实际的任务中我们常用的是预训练好的分词器

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("<name_of_awesome_bert_model>")
model = AutoModel.from_pretrained("<name_of_awesome_bert_model>")

DNABERT

在这个lab里面用DNABERT-6, 这个模型是针对DNA序列的

Architecture

12个Transformer Encoder layers
12个Attention Heads

有了Part1实验当中的经历, 图上的很多模块看起来比较熟悉了, 基本上就是如下序列处理:

输入序列 -> Tokenizers分成一个个的Token -> Token Embedding层化为向量 -> Positional Encoding -> Input Embedding -> Transformer Encoder layers -> 尾部额外的层

Problem 4a

加载实验给出的DNA数据, 数据列是一个sequence, 代表DNA序列, species代表物种, 这里有大猩猩, 狗和人, 我们要手动创建一个species到target的映射, 把物种映射到整数标签

还有一个class列, 代表该DNA序列在生物学上的分类, 但在本实验中我们还是预测target而非class

set_seed() # Resets the global seed state for reproducibility

# Map of species name to class ID
species_to_target = {
    'chimpanzee': 0,
    'dog': 1,
    'human': 2
}

# Map of targets (class IDs) to species name
target_to_species = {
    target: species for species, target in species_to_target.items()
}


# TODO: Create 3 dataframes of the chimpanzee, dog, and human DNA sequences
chimpanzee_df = pd.read_table("./chimpanzee_train.txt")
dog_df = pd.read_table("./dog_train.txt")
human_df = pd.read_table("./human_train.txt")

chimpanzee_df['species']='chimpanzee'
dog_df['species']='dog'
human_df['species']='human'



# TODO: Calculate the length of the dataframe that has the smallest number of DNA sequences
num_samples_per_species = min(len(chimpanzee_df),len(dog_df),len(human_df))

# TODO: Sample `num_samples_per_species` sequences from each of the 3 dataframes and combine them into 1 dataframe
combined_df = pd.concat([chimpanzee_df.sample(n=num_samples_per_species),dog_df.sample(n=num_samples_per_species),
human_df.sample(n=num_samples_per_species)])

combined_df['target'] = combined_df['species'].map(species_to_target)

# Check that combined_df has the expected length
assert len(combined_df) == 3 * 656, f"Length of the combined dataframe is not {3 * 656}"

查看每个物种(总共三个物种)的第一行的DNA序列的前50个碱基对

chimpanzee_dna = combined_df[combined_df['species']=='chimpanzee'].iloc[0]['sequence']
print(f"First 50 DNA tags of chimpanzee sequence: {chimpanzee_dna[:50]}...")

dog_dna = combined_df[combined_df['species']=='dog'].iloc[0]['sequence']
print(f"First 50 DNA tags of dog sequence: {dog_dna[:50]}...")

human_dna = combined_df[combined_df['species']=='human'].iloc[0]['sequence']
print(f"First 50 DNA tags of human sequence: {human_dna[:50]}...")

First 50 DNA tags of chimpanzee sequence: ATGGGCATGACACGGATGCTCCTGGAATGCAGTCTCAGTGACAAGTTGTG...
First 50 DNA tags of dog sequence: ATGGAGGTGCAGACAAAGAAAGTTCGAAAAGTTCCTCCAGGTTTGCCATC...
First 50 DNA tags of human sequence: ATGGAATCTGTGGTAAAGAACTGTGGCCAGACAGTTCATGATGAGGTGGC...

这个sequence和正常的语料是不一样的, 我们需要找一个分词器来把这个sequence切成一个个的token

Problem 4b

利用k-mers来进行切分, 这里k=6, 本质上就是个长度为k的滑动窗口

ATGCGTACTAAG
ATGCGT index 0
TGCGTA index 1
GCGTAC index 2
...

看起来是一种很奇怪的切分方式, 这里的token排列组合只有4^6种, 看起来不太适合自然语言的处理

def sequence_to_kmer(sequence, k=6):
    """
    Converts a sequence into a space-separated string of k-mers.

    A k-mer is a substring of length `k` from overlapping positions in
    the input sequence. Returns all k-mers as a single string separated by spaces.

    Args:
        sequence (str): The input sequence (e.g., DNA, RNA, or text) to tokenize into k-mers.
        k (int, optional): Length of each k-mer. Defaults to 6.

    Returns:
        str: Space-separated string of k-mers.

    Example:
        >>> kmer_tokenize("ATGCGT", k=3)
        'ATG TGC GCG CGT'
    """
    # TODO: Implement the kmer_tokenize function
    result=[]
    for i in range(0,len(sequence)-k+1):
        string=sequence[i:i+k]
        result.append(string)
    
    return ' '.join(result)

# TODO: Apply `kmer_tokenze` to the DNA sequences and save it into a dataframe column called `kmers`
combined_df['kmers'] = combined_df['sequence'].apply(sequence_to_kmer)

# Print the first 5 rows of the dataframe
print(combined_df.sample(5, random_state=SEED))

sequence  class     species  \
622   ATGTCTTTGGTGGACTTGGGGAAGAGGTTGCTAGAAGCAGCAAGAA...      6         dog   
2633  ATGGGAGGCCGCGTCTTTCTCGCATTCTGTGTCTGGCTGACTCTGC...      0       human   
1101  ATGAAAGCCCACCCCAAGGAGATGGTGCCTCTCATGGGCAAGAGAG...      5  chimpanzee   
294   ATGGTCAACGTCTTGAAAGGAGTGCTGATAGAATGTGACCCTGCCA...      6         dog   
48    ATGGGCTGTGTGTTCTGCAAGAAGTCGGAGCCGGGGCTCAAGGACG...      1         dog   

      target                                              kmers  
622        1  ATGTCT TGTCTT GTCTTT TCTTTG CTTTGG TTTGGT TTGG...  
2633       2  ATGGGA TGGGAG GGGAGG GGAGGC GAGGCC AGGCCG GGCC...  
1101       0  ATGAAA TGAAAG GAAAGC AAAGCC AAGCCC AGCCCA GCCC...  
294        1  ATGGTC TGGTCA GGTCAA GTCAAC TCAACG CAACGT AACG...  
48         1  ATGGGC TGGGCT GGGCTG GGCTGT GCTGTG CTGTGT TGTG...

这样就实现了切分

接下来加载一下预训练好的DNABERT-6模型和分词器

dnabert_tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True, revision="c56e67ea5827e0ddc67ef059addcf71569b1216e")
dnabert_model = AutoModel.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True, revision="c56e67ea5827e0ddc67ef059addcf71569b1216e")

注意前面的k-mers只不过是切分成小的字符串, 我们要把它转换成数字才能送进模型, 这一步就是用Tokenizer来实现的

DNA 序列:  ATGCGTACTAAG
            ↓
生成 k-mers: ATGCGT, TGCGTA, GCGTAC, ...
            ↓
Tokenizer:  [CLS] ATG CGT TGC ... [SEP] [PAD] ...
            ↓
input_ids:  [1, 25, 30, 28, ..., 2, 0, 0, ...]
attention_mask: [1, 1, 1, 1, ..., 1, 0, 0, ...]
            ↓
DNABERT 模型
            ↓
last_hidden_state: (1, 512, 768)
            ↓
取 [CLS] 嵌入 → 用于分类或其他任务

Step1: 拿出一行(做例子)k-mers之后了的数据送入Tokenizer里面

Step2: 用返回的input_ids和attention_mask送入模型, 得到last_hidden_state

last_hidden_state是一个特殊结构, 大概理解为模型的输出就好

Problem 4c

现在模型输出的是一个(1, 512, 768)的tensor, 也就是所谓的last_hidden_state, 但我们最终想要得到的是整数, 也就是完成这个分类任务

为了得到整数分类, 一个很Trivial的想法自然是用Linear层去把768维映射成num_classes维, 但这里要注意, 每个token都被送到一个768维的向量了, 我们只需要提取第0个token, 即[CLS]对应的那个嵌入向量

文档中说这是因为[CLS]的嵌入向量是整个序列的总结, 所以我们要提取这个向量来做分类, 但我觉得这并不是一个很平凡的结论, 也许这是长期实践当中归纳出的经验吧

class DNAClassifier(nn.Module):
  def __init__(self, num_classes):
    super().__init__()
    self.backbone = AutoModel.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True)
    self.classifier = nn.Linear(in_features=768, out_features=num_classes)

  def forward(self, input_ids, attention_mask):
    # TODO: Pass the input_ids and attention_mask
    last_hidden_state = self.backbone(input_ids=input_ids,attention_mask=attention_mask).last_hidden_state


    # TODO: Get the embedding of the [CLS] token
    cls_embeddings = last_hidden_state[:,0,:]

    # TODO: Pass the cls_embeddings into the classifier to get class predictions
    logits = self.classifier(cls_embeddings)
    return logits

  def print_params(self):
    for name, param in self.backbone.named_parameters():
      print(f"Name: {name}\tparameter shape: {param.shape}\trequires grad: {param.requires_grad}")
    for name, param in self.classifier.named_parameters():
      print(f"Name: {name}\tparameter shape: {param.shape}\trequires grad: {param.requires_grad}")

注意一下这个

cls_embeddings = last_hidden_state[:,0,:]

中间的0就代表第0个token

Problem 4d

把k-mers之后数据组装成Dataset

讲一下这个__getitem__方法, 这里拿到kmers里面对应的数据之后要去进行Tokenize做映射, 原则上要返回三个元素:

'input_ids': 数字序列
'attention_mask': 掩码序列
'target': 标签

前两个都是分词器返回的, 第三个是从类成员变量里拿的label, 注意分词器返回的要squeeze一下

其他没有什么难点, 跟着TODO走就好, 指导已经写的很明白了

class DNADataset(Dataset):
  def __init__(self, kmers: list[str], targets: list[int] = None, tokenizer=None, return_targets=True):
    super().__init__()
    self.kmers = kmers
    self.targets = targets
    self.return_targets = return_targets  # Whether to return targets or not
    self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True, revision="c56e67ea5827e0ddc67ef059addcf71569b1216e")

  def __len__(self):
    # TODO: Return the number of samples in the dataset
    return len(self.kmers)

  def __getitem__(self, idx):





    # TODO: Get the k-mer at the requested index
    k_mer=self.kmers[idx]
    # TODO: If return_targets is True, get the target at the requested index and convert to torch.Tensor with dtype=torch.long
    if self.return_targets:
      target=self.targets[idx]
      target=torch.tensor(target,dtype=torch.long)

    # TODO: Tokenize the k-mer using the tokenizer. Don't forget to specify return_tensors, max_length, truncation, and padding!
    tokens=self.tokenizer(k_mer,return_tensors='pt',max_length=512,truncation=True,padding='max_length')

    # TODO: Extract and reshape the input_ids
    input_ids=tokens['input_ids'].squeeze(0)

    # TODO: Extract and reshape the attention_mask
    attention_mask=tokens['attention_mask'].squeeze(0)

    # TODO: Return a dictionary with the input_ids, attention_mask, and target (if return_targets is True)
    result = {'input_ids':input_ids,'attention_mask':attention_mask}

    if self.return_targets:
      result['target']=target

    return result

Problem 4e

从原始数据当中切分出训练集和验证集, 并且装载进刚实现的类里面

注意在train_test_split之后要把得到的数据做to_list操作, 否则会报错, 因为刚刚写的类传入的kmers和target都是list

set_seed() # Resets the global seed state for reproducibility
# TODO: Sample 1000 DNA sequences from `combined_df`. Make sure to set the random_state!
combined_df=combined_df.sample(n=1000,random_state=SEED).reset_index()

# TODO: Use train_test_split to create training and validation splits
kmers_train,kmers_val,targets_train,targets_val=train_test_split(combined_df['kmers'],combined_df['target'],test_size=0.2,train_size=0.8,stratify=combined_df['target'],
random_state=SEED,shuffle=True)


# 在 train_test_split 之后，创建 dataset 前运行这段以规范化索引/类型

# 1) 把 kmers 转为普通 list（安全，索引语义不会影响）
kmers_train = kmers_train.tolist() if hasattr(kmers_train, "tolist") else list(kmers_train)
kmers_val   = kmers_val.tolist()   if hasattr(kmers_val, "tolist")   else list(kmers_val)

# 2) 把 targets 转为普通 list（或 numpy 数组）
targets_train = targets_train.tolist() if hasattr(targets_train, "tolist") else list(targets_train)
targets_val   = targets_val.tolist()   if hasattr(targets_val, "tolist")   else list(targets_val)

# 3) 重新创建 dataset 和 dataloader（确保 num_workers=0 便于调试）
train_set = DNADataset(kmers=kmers_train, targets=targets_train, return_targets=True)
val_set   = DNADataset(kmers=kmers_val,   targets=targets_val,   return_targets=True)

train_dataloader = DataLoader(train_set, batch_size=32, shuffle=True)
val_dataloader   = DataLoader(val_set,   batch_size=32, shuffle=False)

Problem 4f

实现训练循环, 无须多言, 都是公式化代码了

def train_dna_classifier(model, optimizer, criterion, device, num_epochs, train_dataloader, val_dataloader):
    """
    Args:
        model: the model to train
        optimizer: the optimizer to use
        criterion: the loss function to use
        num_epochs: the number of epochs to train for
        train_dataloader: the dataloader for the training set
        val_dataloader: the dataloader for the validation set
    Returns:
        train_losses: a list of training losses for each epoch
        val_losses: a list of validation losses for each epoch
        train_accuracies: a list of training accuracies for each epoch
        val_accuracies: a list of validation accuracies for each epoch
    """
    # === SETUP ===
    model.to(device)

    # Lists to store metrics across epochs
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []

    # === EPOCH LOOP ===
    for epoch in range(num_epochs):
        # === TRAINING PHASE ===
        model.train() # Set model to training mode

        # Initialize metrics for this epoch
        train_loss = 0.0
        train_correct = 0

        # === INNER LOOP (iterate over training batches) ===
        for batch in train_dataloader:
            # TODO 1: Move the data and the targest to device
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            targets = batch['target'].to(device)

            # TODO 2: Reset gradients
            optimizer.zero_grad()

            # TODO 3: Forward pass: pass inputs to model
            logits = model(input_ids=input_ids,attention_mask=attention_mask)

            # TODO 4: Compute loss
            loss = criterion(logits,targets)

            # TODO 5: Backward pass/compute gradients
            loss.backward()

            # TODO 6: Update parameters
            optimizer.step()

            # TODO 7: Track training metrics for this epoch
            train_loss+=loss.item()
            _, preds = torch.max(logits,dim=1)
            train_correct+=(preds==targets).sum().item()

        # === END OF INNER LOOP ===

        # Compute average training metrics for the epoch
        train_loss /= len(train_dataloader)
        train_acc = train_correct / len(train_dataloader.dataset)

        # Append this epoch's training metrics to history
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)

        print(f"Epoch {epoch + 1}: Training loss = {train_loss}\tTraining accuracy = {train_acc}")

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        model.eval() # set model to evaluation mode
        val_loss = 0.0
        val_correct = 0

        with torch.no_grad():
            # === INNER LOOP (iterate over validation batches) ===
            for batch in val_dataloader:
                # TODO 8: Move data to device
                input_ids = batch['input_ids'].to(device)
                attention_mask = batch['attention_mask'].to(device)
                targets = batch['target'].to(device)

                # TODO 9: Forward pass only
                logits = model(input_ids=input_ids,attention_mask=attention_mask)

                # TODO 10: Compute loss
                loss = criterion(logits,targets)

                # TODO 11: Track validation metrics
                val_loss+=loss.item()
                _, preds = torch.max(logits,dim=1)
                val_correct+=(preds==targets).sum().item()

            # === END OF INNER LOOP ===

            val_loss /= len(val_dataloader)
            val_acc = val_correct / len(val_dataloader.dataset)

        val_losses.append(val_loss)
        val_accuracies.append(val_acc)
        print(f"Epoch {epoch + 1}: Validation loss = {val_loss}\tValidation accuracy = {val_acc}")

        # === END OF VALIDATION PHASE ===

    # === END OF EPOCH LOOP ===

    # Return history
    return train_losses, val_losses, train_accuracies, val_accuracies

Problem 4g

指定模型, 优化器, 损失函数以及其他超参数, 启动训练即可

# TODO: Instantiate a DNAClassifier
dna_classifier = DNAClassifier(num_classes=3)

# TODO: Define the optimizer
optimizer = AdamW(params=dna_classifier.parameters(),lr=0.0001)

# TODO: Define the loss function
criterion = nn.CrossEntropyLoss()

# TODO: Train your DNA classifier for 5 epochs!
dna_train_losses, dna_val_losses, dna_train_accuracies, dna_val_accuracies = train_dna_classifier(
    model=dna_classifier,
    optimizer=optimizer,
    criterion=criterion,
    device=device,
    num_epochs=10,
    train_dataloader=train_dataloader,
    val_dataloader=val_dataloader
)

A new version of the following files was downloaded from https://huggingface.co/zhihan1996/DNA_bert_6:
- configuration_bert.py
. Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.
A new version of the following files was downloaded from https://huggingface.co/zhihan1996/DNA_bert_6:
- dnabert_layer.py
. Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.
Epoch 1: Training loss = 1.1392503881454468	Training accuracy = 0.37875
Epoch 1: Validation loss = 1.1413285732269287	Validation accuracy = 0.355
Epoch 2: Training loss = 1.1108909916877747	Training accuracy = 0.3875
Epoch 2: Validation loss = 1.1069353307996477	Validation accuracy = 0.375
Epoch 3: Training loss = 1.0065120267868042	Training accuracy = 0.495
Epoch 3: Validation loss = 1.089234403201512	Validation accuracy = 0.45
Epoch 4: Training loss = 0.8722425246238709	Training accuracy = 0.61875
Epoch 4: Validation loss = 1.138353475502559	Validation accuracy = 0.415
Epoch 5: Training loss = 0.7556110572814941	Training accuracy = 0.67625
Epoch 5: Validation loss = 1.2105737583977836	Validation accuracy = 0.435
Epoch 6: Training loss = 0.5860666036605835	Training accuracy = 0.79125
Epoch 6: Validation loss = 1.2875271865299769	Validation accuracy = 0.45
Epoch 7: Training loss = 0.4519534611701965	Training accuracy = 0.84125
Epoch 7: Validation loss = 1.397601638521467	Validation accuracy = 0.415
Epoch 8: Training loss = 0.37866723477840425	Training accuracy = 0.87125
Epoch 8: Validation loss = 1.5645615543637956	Validation accuracy = 0.46
Epoch 9: Training loss = 0.2757827317714691	Training accuracy = 0.90625
Epoch 9: Validation loss = 1.8461172580718994	Validation accuracy = 0.42
Epoch 10: Training loss = 0.24834500044584273	Training accuracy = 0.91625
Epoch 10: Validation loss = 1.8626392909458704	Validation accuracy = 0.445

Problem 4h

画损失曲线

def plot_metrics(train_losses, val_losses, train_accuracies=None, val_accuracies=None, num_epochs=None, title=""):
    """
    Plots the training loss, training accuracy, validation loss, and validation accuracy.
    Args:
        train_losses: list of training losses
        val_losses: list of validation losses
        train_accuracies: list of training accuracies
        val_accuracies: list of validation accuracies
        num_epochs: number of epochs
        title: title of the plot
    """

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # 训练损失
    axes[0, 0].plot(train_losses, label='Training Loss')
    axes[0, 0].set_title('Training Loss')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')

    axes[0, 1].plot(train_accuracies, label='Training Accuracy')
    axes[0, 1].set_title('Training Accuracy')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy')

    axes[1, 0].plot(val_losses, label='Validation Loss')
    axes[1, 0].set_title('Validation Loss')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('Loss')

    axes[1, 1].plot(val_accuracies, label='Validation Accuracy')
    axes[1, 1].set_title('Validation Accuracy')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('Accuracy')

    plt.tight_layout()
    plt.suptitle(title)
    plt.show()

# TODO: Plot your DNA classifier's loss and accuracy curves on the training data and the validation data.
# You should have 4 plots total!
plot_metrics(dna_train_losses, dna_val_losses, dna_train_accuracies, dna_val_accuracies)

Problem 4i

生成一份预测结果csv, 一键运行代码即可(只要前面的模型训练好了)

第二部分: 声音分类

wav数据

现在我们来处理一个更麻烦的任务, 给出一些.wav的文件, 这些声音来自不同的源, 比如说有些可能是空调的声音, 有些是狗叫, 目标是做这个分类

比较麻烦的点在于, 要先把这些.wav文件变成正常的数据格式, 这一点我们用torchaudio来实现

文件名里面有一个整数代表class_id, 这也是我们需要的label

本实验要用到交叉验证, 也就是把数据分为多个"折"(fold), 每次用n-1个折来训练, 用剩下的一个折来测试, 然后这个过程可以重复n次, 每次用一个不同的折来做测试

class_id_to_sound = {
   0: "air_conditioner",
   1: "car_horn",
   2: "children_playing",
   3: "dog_bark",
   4: "drilling",
   5: "engine_idling",
   6: "gun_shot",
   7: "jackhammer",
   8: "siren",
   9: "street_music"
}

课程组给了一些示例代码来播放.wav文件, 因为我是服务器环境所以听不到, 如果在本地应该是可以听的, 不过记住要配好pygame环境

# audio_file = '/content/drive/MyDrive/cs189/hw/hw4/data/fold1_train/101415-3-0-2.wav'
audio_file = './data/fold1_train/101415-3-0-2.wav'

os.environ['SDL_AUDIODRIVER'] = 'dummy'
# no sound on server

if IS_COLAB:
    from IPython.display import Audio, display
    display(Audio(audio_file, autoplay=False))
else:
    import pygame
    pygame.mixer.init()
    pygame.mixer.music.load(audio_file)
    pygame.mixer.music.play()
    while pygame.mixer.music.get_busy():
        pygame.time.Clock().tick(10)

频谱图

接下来要把这个音频变成频谱图, 频谱图X轴是时间, Y轴是频率, 我们只要实例化一个torchaudio.transforms.Spectrogram就行了, 然后把这个变换应用到原始文件上去

spectrogram_transform = torchaudio.transforms.Spectrogram(n_fft=1024, normalized=True)

class_id_to_sound = {
    0: "air_conditioner",
    1: "car_horn",
    2: "children_playing",
    3: "dog_bark",
    4: "drilling",
    5: "engine_idling",
    6: "gun_shot",
    7: "jackhammer",
    8: "siren",
    9: "street_music"
}

# audio_file = '/content/drive/MyDrive/cs189/hw/hw4/data/fold1/101415-3-0-2.wav'
audio_file = './data/fold1_train/101415-3-0-2.wav'

try:
    # Read WAV file
    waveform, sample_rate = torchaudio.load(audio_file)

    # Converts stero to mono by averaging channels. Shape becomes [1, num_samples]
    waveform = waveform.mean(dim=0, keepdim=True)

    # Transform the waveform into a spectrogram
    spec = spectrogram_transform(waveform) # torch.Tensor of shape [1, freq_bins, time_bins]

    # Print the shape of the spectrogram
    print(f"Spectrogram shape: {spec.shape}") # shape [num_frequencies, num_time_bins]

    # Visualize the spectrogram
    plt.figure()
    plt.imshow(spec.log2()[0, :, :].numpy(), aspect='auto', origin='lower') # origin = 'lower' sets input[0, 0] at bottom left
    plt.title(f"Spectrogram: {audio_file.split('/')[-1]}")
    plt.xlabel("Time bins")
    plt.ylabel("Frequency")
    plt.show()
except Exception as e:
    print(f"Error processing {audio_file}: {e}")

运行他给的代码可以看到频谱图

这本质上就是个二维图像了

Problem 5a

组建Dataset, 几个要注意的点:

从文件名里面把正确的分类拿出来(label)返回
把频谱图的尺寸(1, H, W)通过repeat变成(3, H, W)
如果传入了额外的transform参数, 记得在repeat之后应用

这里repeat的作用就是沿着第0维复制3次把通道数变成3, 以便于适配后面的输入

class SpectrogramDataset(Dataset):
    def __init__(self, audio_dir, transforms=None, spectrogram_transform=None, return_targets=True):
        super().__init__()
        self.audio_dir = audio_dir # Folder where the audio files will be located
        self.transforms = transforms # Additional image transforms
        self.spectrogram_transform = spectrogram_transform or torchaudio.transforms.Spectrogram(n_fft=1024, normalized=True)
        self.return_targets = return_targets  # Whether to return targets or not when __get_item__ is called
        try:
            # Get file paths of all .wav files in the provided directory
            self.file_paths = [f for f in os.listdir(audio_dir) if f.endswith('.wav')]
        except Exception as e:
            print(f"Error accessing files at {audio_dir}: {e}")
            self.file_paths = []

    def __len__(self):
        # TODO: 1. Return the number of samples in the dataset
        return len(self.file_paths)

    def __getitem__(self, index):
        # Get the file name by index
        file_name = self.file_paths[index]

        # Construct the path to the requested audio file
        filepath = os.path.join(self.audio_dir, file_name)

        # TODO: 2. Extract the target class_id from the file path if return_targets is True
        # Hint: Files are stored in the format [freesoundID]-[classID]-[occurrenceID]-[sliceID].wav
        # Don't forget to cast the class_id to dtype=torch.long for loss functions!
        if self.return_targets:
            target = int(file_name.split('-')[1])
            target = torch.tensor(target,dtype=torch.long)

        try:
            # TODO: 3. Load the audio file located at `filepath`
            waveform, sample_rate = torchaudio.load(filepath)
            # TODO: 4. Convert stereo to mono by averaging channels.
            waveform = waveform.mean(dim=0,keepdim=True)

            # TODO: 5. Generate a spectrogram of the waveform. Expected shape: (1, H, W)
            spec = self.spectrogram_transform(waveform)

            # TODO: 6. Copy the spectrogram into 3 channels to get shape: (3, H, W)
            spec = spec.repeat(3,1,1)

            # TODO: 7. Apply transformations
            if self.transforms:
                spec = self.transforms(spec)

            if self.return_targets:
                return spec, target
            else:
                return spec
        except Exception as e:
            print(f'Error processing {filepath}: {e}')
            if self.return_targets:
                return None, None
            else:
                return None

Problem 5b

把Dataset实例化后装到DataLoader里面去

from torch.utils.data import Subset

# TODO: Define a transform to resize images to (224, 224)
resize_transform = torchvision.transforms.Resize((224,224))

# TODO: Instantiate a SpectrogramDataset using the folder fold1
spectrogram_dataset = SpectrogramDataset(audio_dir='data/fold1_train',transforms=resize_transform,return_targets=True)

# TODO: Create a 0.8 training and 0.2 test split
indices=list(range(len(spectrogram_dataset)))
train_indices,test_indices=train_test_split(
    indices,
    test_size=0.2,
    random_state=SEED,
    shuffle=True
)
spectrogram_train_dataset, spectrogram_test_dataset = Subset(spectrogram_dataset,train_indices),Subset(spectrogram_dataset,test_indices)

# TODO: Create training and testing dataloaders
spectrogram_train_dataloader = DataLoader(spectrogram_train_dataset,batch_size=32,shuffle=True)
spectrogram_test_dataloader = DataLoader(spectrogram_test_dataset,batch_size=32,shuffle=False)

# Printing out the shapes of data and targets in the first batch!
batch = next(iter(spectrogram_train_dataloader))

data, targets = batch
print(f"Shape of 1 batch of data: {data.shape}")
print(f"Shape of 1 batch of targets: {targets.shape}")

公式化代码, 没什么好说的

Shape of 1 batch of data: torch.Size([32, 3, 224, 224])
Shape of 1 batch of targets: torch.Size([32])

Problem 5c

实现训练循环

def train_image_classifier(model, optimizer, criterion, device, num_epochs, train_dataloader, val_dataloader):
    """
    Args:
        model: the model to train
        optimizer: the optimizer to use
        criterion: the loss function to use
        num_epochs: the number of epochs to train for
        train_dataloader: the dataloader for the training set
        val_dataloader: the dataloader for the validation set
    Returns:
        train_losses: a list of training losses for each epoch
        val_losses: a list of validation losses for each epoch
        train_accuracies: a list of training accuracies for each epoch
        val_accuracies: a list of validation accuracies for each epoch
    """
    # === SETUP ===
    model.to(device)

    # Lists to store metrics across epochs
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []

    # === EPOCH LOOP ===
    for epoch in range(num_epochs):
        # === TRAINING PHASE ===
        model.train() # Set model to training mode

        # Initialize metrics for this epoch
        train_loss = 0.0
        train_correct = 0

        # === INNER LOOP (iterate over training batches) ===
        for batch in train_dataloader:
            # TODO 1: Move the data and targets to device and cast them to the appropriate dtypes
            x, y = batch
            x, y = x.to(device),y.to(device)

            # TODO 2: Reset gradients
            optimizer.zero_grad()

            # TODO 3: Forward pass: pass inputs to model
            y_hat = model(x)

            # TODO 4: Compute loss
            loss = criterion(y_hat,y)

            # TODO 5: Backward pass/compute gradients
            loss.backward()

            # TODO 6: Update parameters
            optimizer.step()

            # TODO 7: Track training metrics for this epoch
            train_loss+=loss.item()
            _, preds = torch.max(y_hat,dim=1)
            train_correct+=torch.sum(preds==y).item()

        # === END OF INNER LOOP ===

        # Compute average training metrics for the epoch
        train_loss /= len(train_dataloader)
        train_acc = train_correct / len(train_dataloader.dataset)

        # Append this epoch's training metrics to history
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)

        print(f"Epoch {epoch + 1}: Training loss = {train_loss}\tTrain accuracy = {train_acc}")

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        model.eval() # set model to evaluation mode
        val_loss = 0.0
        val_correct = 0

        with torch.no_grad():
            # === INNER LOOP (iterate over validation batches) ===
            for batch in val_dataloader:
                # TODO 8: Move the data and targets to device and cast them to the appropriate dtypes
                x, y = batch
                x, y = x.to(device),y.to(device)

                # TODO 9: Forward pass only
                y_hat = model(x)

                # TODO 10: Compute loss
                loss = criterion(y_hat,y)

                # TODO 11: Track validation metrics
                val_loss+=loss.item()
                _, preds = torch.max(y_hat,dim=1)
                val_correct+=torch.sum(preds==y).item()

            # === END OF INNER LOOP ===

            val_loss /= len(val_dataloader)
            val_acc = val_correct / len(val_dataloader.dataset)

        val_losses.append(val_loss)
        val_accuracies.append(val_acc)
        print(f"Epoch {epoch + 1}: Validation loss = {val_loss}\tValidation accuracy = {val_acc}")

        # === END OF VALIDATION PHASE ===

    # === END OF EPOCH LOOP ===

    print(f"=" * 20 + " Final Metrics " + "=" * 20)
    print(f"Final training loss: {train_losses[-1]:.5f}\tFinal training accuracy = {train_accuracies[-1]:.5f}")
    print(f"Final validation loss: {val_losses[-1]:.5f}\tFinal validation accuracy = {val_accuracies[-1]:.5f}")
    print(f"=" * 55)

    # Return history
    return train_losses, val_losses, train_accuracies, val_accuracies

Problem 5d

这里用的模型是在ImageNet上预训练好的ConvNeXt模型, 他的最后一层的全连层的输出是1000维的, 我们要改成10维, 因为这是个10-分类问题

def replace_final_convnext_linear_layer(model, num_classes=10):
    # TODO: Access the classifier module of the model
    classifier = model.classifier[-1]

    # TODO: Get the input dimensions of the classifier's last linear layer
    in_features = classifier.in_features

    # TODO: replace the model's last linear layer with a new linear layer
    model.classifier[-1] = nn.Linear(in_features,num_classes)

    return model

预训练与微调简介

所谓预训练就是直接用别人调好了参数的模型, 比如在ImageNet上预训练好了的resnet50

model = resnet50(weights = ResNet50_Weights.IMAGENET1K_V2)

当然也可以自己初始化一版模型参数

model = resnet50(weights = None)

微调有好几种方式:

从Scratch训练, 用随机的初始化参数
采用预训练好的参数, 并且冻结除了最后一个分类头之外的所有参数, 只训练最后那个Linear Layer
全量微调, 用预训练好的参数但是所有参数都可以再次进行训练

实验也给出了冻结某一层的参数的方式:

for name, param in frozen_backbone.named_parameters():
    if "classifier" not in name: # Freeze any non-classifier layers
        param.requires_grad = False

所谓"冻结"也就是是否允许参数被更改, 当然也就是设置这个梯度的bool

Problem 5e

用全部的数据来从头训练convnext模型

# TODO: Load a ConvNeXt model with uninitialized weights
uninitialized_convnext = convnext_base(weights=None)

# TODO: Replace the final classification layer of the ConvNeXt
uninitialized_convnext = replace_final_convnext_linear_layer(model=uninitialized_convnext, num_classes=10)

# TODO: Initialize the optimizer
optimizer = AdamW(params=uninitialized_convnext.parameters(),lr=1e-4)

# TODO: Define the loss function
criterion = nn.CrossEntropyLoss()

# TODO: Train your model for 5 epochs
unintialized_train_losses, unintialized_val_losses, unintialized_train_accuracies, unintialized_val_accuracies = train_image_classifier(
    model=uninitialized_convnext, 
    train_dataloader=spectrogram_train_dataloader, 
    val_dataloader=spectrogram_test_dataloader, 
    optimizer=optimizer, 
    num_epochs=5,
    criterion=criterion, 
    device=device)

# TODO: Plot the metrics
plot_metrics(unintialized_train_losses, unintialized_val_losses, unintialized_train_accuracies, unintialized_val_accuracies)

Epoch 1: Training loss = 2.3848066329956055	Train accuracy = 0.1774193548387097
Epoch 1: Validation loss = 2.237703227996826	Validation accuracy = 0.20714285714285716
Epoch 2: Training loss = 2.1374161640803018	Train accuracy = 0.21505376344086022
Epoch 2: Validation loss = 2.2762171745300295	Validation accuracy = 0.2357142857142857
Epoch 3: Training loss = 2.1253870791859097	Train accuracy = 0.23297491039426524
Epoch 3: Validation loss = 2.116562104225159	Validation accuracy = 0.2
Epoch 4: Training loss = 2.0567029780811734	Train accuracy = 0.2078853046594982
Epoch 4: Validation loss = 2.0853104829788207	Validation accuracy = 0.19285714285714287
Epoch 5: Training loss = 2.0093974073727927	Train accuracy = 0.22580645161290322
Epoch 5: Validation loss = 2.036539649963379	Validation accuracy = 0.2642857142857143
==================== Final Metrics ====================
Final training loss: 2.00940	Final training accuracy = 0.22581
Final validation loss: 2.03654	Final validation accuracy = 0.26429
=======================================================

注意这里的学习率是超参数, 可以调的

Problem 5f

冻结除了分类头以外的层, 只训练分类头

# TODO: Load a ConvNeXt model with pretrained weights
frozen_backbone = convnext_base(weights=ConvNeXt_Base_Weights.IMAGENET1K_V1)

# TODO: Replace the final classification layer of the ConvNeXt
frozen_backbone = replace_final_convnext_linear_layer(model=frozen_backbone, num_classes=10)

# TODO: Freeze the ConvNeXt's backbone
for name,param in frozen_backbone.named_parameters():
    if "classifier" not in name:
        param.requires_grad = False
    
    else:
        param.requires_grad=True

# TODO: Initialize the optimizer
optimizer = AdamW(params=frozen_backbone.parameters(),lr=1e-4)

# TODO: Define the loss function
criterion = nn.CrossEntropyLoss()

# TODO: Train your model for 5 epochs
frozen_bb_train_losses, frozen_bb_val_losses, frozen_bb_train_accuracies, frozen_bb_val_accuracies = train_image_classifier(
    model=frozen_backbone, 
    train_dataloader=spectrogram_train_dataloader, 
    val_dataloader=spectrogram_test_dataloader, 
    optimizer=optimizer, 
    criterion=criterion, 
    num_epochs=5,
    device=device)

# TODO: Plot the metrics
plot_metrics(frozen_bb_train_losses, frozen_bb_val_losses, frozen_bb_train_accuracies, frozen_bb_val_accuracies)

Epoch 1: Training loss = 2.303494784567091	Train accuracy = 0.08064516129032258
Epoch 1: Validation loss = 2.2263582229614256	Validation accuracy = 0.17142857142857143
Epoch 2: Training loss = 2.220644950866699	Train accuracy = 0.18100358422939067
Epoch 2: Validation loss = 2.1564966201782227	Validation accuracy = 0.17857142857142858
Epoch 3: Training loss = 2.168804738256666	Train accuracy = 0.21863799283154123
Epoch 3: Validation loss = 2.100537633895874	Validation accuracy = 0.22857142857142856
Epoch 4: Training loss = 2.118125465181139	Train accuracy = 0.25806451612903225
Epoch 4: Validation loss = 2.0520622968673705	Validation accuracy = 0.3
Epoch 5: Training loss = 2.068689114517636	Train accuracy = 0.2974910394265233
Epoch 5: Validation loss = 2.0029944658279417	Validation accuracy = 0.32142857142857145
==================== Final Metrics ====================
Final training loss: 2.06869	Final training accuracy = 0.29749
Final validation loss: 2.00299	Final validation accuracy = 0.32143
=======================================================

Problem 5g

全量微调

# TODO: Load a ConvNeXt model with pretrained weights
unfrozen_convnext = convnext_base(weights=ConvNeXt_Base_Weights.IMAGENET1K_V1)

# TODO: Replace the final classification layer of the ConvNeXt
unfrozen_convnext = replace_final_convnext_linear_layer(model=unfrozen_convnext, num_classes=10)

# TODO: Make sure all the parameters are trainable
for param in unfrozen_convnext.parameters():
    param.requires_grad = True

# TODO: Initialize the optimizer
optimizer = AdamW(params=unfrozen_convnext.parameters(),lr=1e-4)

# TODO: Define the loss function
criterion = nn.CrossEntropyLoss()

# TODO: Train your model for 5 epochs
unfrozen_train_losses, unfrozen_val_losses, unfrozen_train_accuracies, unfrozen_val_accuracies = train_image_classifier(
    model=unfrozen_convnext, 
    train_dataloader=spectrogram_train_dataloader, 
    val_dataloader=spectrogram_test_dataloader, 
    optimizer=optimizer, 
    num_epochs=5,
    criterion=criterion, 
    device=device)

# TODO: Plot the metrics
plot_metrics(unfrozen_train_losses, unfrozen_val_losses, unfrozen_train_accuracies, unfrozen_val_accuracies)

Epoch 1: Training loss = 1.8407021694713168	Train accuracy = 0.3888888888888889
Epoch 1: Validation loss = 1.1633551120758057	Validation accuracy = 0.7
Epoch 2: Training loss = 0.9334362381034427	Train accuracy = 0.7419354838709677
Epoch 2: Validation loss = 0.5684262037277221	Validation accuracy = 0.85
Epoch 3: Training loss = 0.4550878521468904	Train accuracy = 0.8888888888888888
Epoch 3: Validation loss = 0.4825403094291687	Validation accuracy = 0.8357142857142857
Epoch 4: Training loss = 0.3048595061732663	Train accuracy = 0.9068100358422939
Epoch 4: Validation loss = 0.3036103412508965	Validation accuracy = 0.9
Epoch 5: Training loss = 0.19161400902602407	Train accuracy = 0.946236559139785
Epoch 5: Validation loss = 0.3878093510866165	Validation accuracy = 0.8785714285714286
==================== Final Metrics ====================
Final training loss: 0.19161	Final training accuracy = 0.94624
Final validation loss: 0.38781	Final validation accuracy = 0.87857
=======================================================

Problem 5h

分析一下上面三种方式(从Scratch训练, 冻结除了分类头以外的层, 全量微调)的绩效

Pretrained + Full Fine-tuning is the most effective approach because it allows the model to adapt to the new dataset and achieve the highest accuracy.

Disadvantages:
* It requires more data and training time.
* It may overfit to the new dataset if the amount of data is limited.

Problem 5i

推理一次, 看看预测的效果

# TODO 1: Pick an audio file to listen to and save it to the `audio_file` variable
# audio_file = '/content/drive/MyDrive/cs189/hw/hw4/data/fold1_train/101415-3-0-2.wav'
audio_file = 'data/fold1_train/101415-3-0-2.wav'
if IS_COLAB:
    from IPython.display import Audio, display
    display(Audio(audio_file, autoplay=False))
else:
    import pygame
    pygame.mixer.init()
    pygame.mixer.music.load(audio_file)
    pygame.mixer.music.play()
    while pygame.mixer.music.get_busy():
        pygame.time.Clock().tick(10)

try:
    # TODO 2: Load the .wav audio file
    waveform, sample_rate = torchaudio.load(audio_file)

    # TODO 3: Average the channels to create a single mono channel
    waveform = waveform.mean(dim=0, keepdim=True)

    # TODO 4: Generate a spectrogram
    spec = spectrogram_transform(waveform)

    # TODO 5: Resize the spectrogram into the shape (1 x 224 x 224)
    spec = resize_transform(spec)

    # TODO 6: Repeat the spectrogram 3 times to have 3 channels
    spec = spec.repeat(3,1,1)

    # TODO 7: Add a batch dimension
    spec = spec.unsqueeze(0)

    # TODO 8: Move the input to the right device and cast it to the right dtype
    spec = spec.to(device)
    spec = spec.float()

    # TODO 9: Get the model's outputs
    y_hat = unfrozen_convnext(spec)

    # TODO 10: find the class with the highest output score
    _, pred = torch.max(y_hat,dim=1)

    # TODO 11: Look up the class label of the model's prediction
    predicted_class = class_id_to_sound[pred.item()]

    # TODO 12: Parse the audio file's name to get the true class ID and the true class label
    true_class_id = os.path.basename(audio_file).split('-')[1]
    true_class = class_id_to_sound[int(true_class_id)]

    print(f"Predicted class: {predicted_class}")
    print(f"True class: {true_class}")
except Exception as e:
    print(f'Error processing {audio_file}: {e}')

Predicted class: dog_bark
True class: dog_bark

UC Berkeley CS189 Assignment 4 (Part 1)

Sun, 08 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 4

项目描述

从这个项目起基本上告别小打小闹了, 我们要用PyTorch来搭建一些较为复杂的模型, 比如说CNN, Transformer, Bert. 关于PyTorch的API用法我不会做太多解释, 和CS336笔记中一样, 我会着重说明这些线性变换的尺寸和方式, 我觉得这才是真正理解模型的重要点

学习目标如下:

-- 用PyTorch搭建自己的神经网络 -- 理解并实现自定义的Datasets, DataLoaders和训练循环 -- 学习如何复现论文中的模型架构 -- 理解ResNet架构 -- 用PyTorch实现Transformer -- 理解Transformer的原理

Problem 1a

用PyTorch写一个CNN类, 非常简单, 只要照着他给出的卷积核的尺寸往里填就行了

class CNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # TODO: Instantiate the layers your CNN will have!
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=7, stride=2, padding=1)
        self.linear = nn.Linear(in_features=16*54*54, out_features=num_classes)
        # Optional: You can also define your ReLU layers here if you'd like
        # Otherwise, if you decide to use F.relu you can just remove these lines
        self.relu1 = nn.ReLU()
        self.relu2 = nn.ReLU()

    def forward(self, x):
        # TODO: 1. Pass the input through the first conv layer
        conv1_out = self.conv1(x)

        # TODO: 2. Apply ReLU
        relu1_out = self.relu1(conv1_out)

        # TODO: 3. Pass to the 2nd conv layer
        conv2_out = self.conv2(relu1_out)

        # TODO: 4. Apply ReLU
        relu2_out = self.relu2(conv2_out)

        # TODO: 5. Flatten before linear layer
        flat = relu2_out.view(relu2_out.size(0), -1)

        # TODO: 6. Pass to linear layer
        out = self.linear(flat)

        return out

假设我们传入的尺寸为(1, 224 ,224 ,3), 具体含义是一个batch有一张图片, 一个图片有三个图层(channel), 图片的长和宽都是224, 不过我觉得没有必要对高维张量去进行具体的想象什么维度代表什么, 因为超过三维的东西就很难想得明白, 只要记住这个二维卷积核肯定是对倒数第一和第二的维度做变换就行了

conv1上的尺寸变换

conv1的尺寸为:

self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)

先考虑对最后两维的变化, 卷积核的尺寸为3, 步长为2, 填充为1, 所以输出尺寸为:

Gauss((224 + 1 * 2 - 3) / 2 + 1) = 112

其中Gauss()是向下取整函数, 又因为卷积核指定了out_channels = 16, 所以输出尺寸为(1, 16, 112, 112)

再解释一下这个输出通道的变化, 如图所示:

输入图片(3个图层)           每个输出通道的卷积核
┌─────────────┐
│   224 * 224   │ ──▶ 3 * 3 矩阵(处理R层)──┐
│   (R通道)   │                         │
├─────────────┤                         │
│   224 * 224   │ ──▶ 3 * 3 矩阵(处理G)）──┼──▶ 求和 ──▶ 输出1个通道
│   (G通道)   │                         │
├─────────────┤                         │
│   224 * 224   │ ──▶ 3 * 3 矩阵(处理B层)──┘
│   (B通道)   │
└─────────────┘

对于每个输出通道(1/16), 都有3个3*3的卷积核, 每个卷积核去处理输入当中的一个通道, 所以总共有16 * 3个卷积核, 每个的尺寸是3 * 3

也就是说, 输出通道变多的意思是我们有很多的卷积核分别在原图的不同层次上学习特征, 并且每次只输出一个通道

conv2上的尺寸变换

conv2的尺寸为:

self.conv2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=7, stride=2, padding=1)

先考虑对最后两维的变化, 卷积核的尺寸为7, 步长为2, 填充为1, 所以输出尺寸为:

Gauss((112 + 1 * 2 - 7) / 2 + 1) = 54

所以输出尺寸为(1, 16, 54, 54)

Linear上的尺寸变换

linear的尺寸为:

self.linear = nn.Linear(in_features=16*54*54, out_features=num_classes)

输入尺寸为(1, 16, 54, 54), 所以输出尺寸为(1, num_classes)

这里其实有两个步骤, 首先去做一个flatten操作, 把(1, 16, 54, 54)变成(1, 16 * 54 * 54), 然后再去做一个线性变换, 把(1, 16 * 54 * 54)变成(1, num_classes), 注意这一步总是可行的, 因为一个高维的东西总是可以把他的每个元素拿出来然后flatten到一维

检查输出尺寸

运行他的检查代码, 确保我们的实现是正确的

x = torch.rand((1, 3, 224, 224))  # 1 image, 3 RGB channels, 224 x 224 pixels
print(f"Shape of input: {x.shape}")

res = CNN(num_classes=10)
out = res(x)
print(f"Shape of output: {out.shape}")
assert out.shape == torch.Size([1, 10]), f"Expected output shape (1, 10), but got {out.shape}"

Shape of input: torch.Size([1, 3, 224, 224])
Shape of output: torch.Size([1, 10])

理解HF Dataset对象

很多时候要用HuggingFace的load_dataset来加载数据集, 所以有必要看一下加载的数据集对象的数据结构, 运行示例代码即可

All classes in the ImageNet dataset: {0: 'house_finch', 1: 'robin', 2: 'triceratops', 3: 'green_mamba', 4: 'harvestman', 5: 'toucan', 6: 'goose', 7: 'jellyfish', 8: 'nematode', 9: 'king_crab', 10: 'dugong', 11: 'Walker_hound', 12: 'Ibizan_hound', 13: 'Saluki', 14: 'golden_retriever', 15: 'Gordon_setter', 16: 'komondor', 17: 'boxer', 18: 'Tibetan_mastiff', 19: 'French_bulldog', 20: 'malamute', 21: 'dalmatian', 22: 'Newfoundland', 23: 'miniature_poodle', 24: 'white_wolf', 25: 'African_hunting_dog', 26: 'Arctic_fox', 27: 'lion', 28: 'meerkat', 29: 'ladybug', 30: 'rhinoceros_beetle', 31: 'ant', 32: 'black-footed_ferret', 33: 'three-toed_sloth', 34: 'rock_beauty', 35: 'aircraft_carrier', 36: 'ashcan', 37: 'barrel', 38: 'beer_bottle', 39: 'bookshop', 40: 'cannon', 41: 'carousel', 42: 'carton', 43: 'catamaran', 44: 'chime', 45: 'clog', 46: 'cocktail_shaker', 47: 'combination_lock', 48: 'crate', 49: 'cuirass', 50: 'dishrag', 51: 'dome', 52: 'electric_guitar', 53: 'file', 54: 'fire_screen', 55: 'frying_pan', 56: 'garbage_truck', 57: 'hair_slide', 58: 'holster', 59: 'horizontal_bar', 60: 'hourglass', 61: 'iPod', 62: 'lipstick', 63: 'miniskirt', 64: 'missile', 65: 'mixing_bowl', 66: 'oboe', 67: 'organ', 68: 'parallel_bars', 69: 'pencil_box', 70: 'photocopier', 71: 'poncho', 72: 'prayer_rug', 73: 'reel', 74: 'school_bus', 75: 'scoreboard', 76: 'slot', 77: 'snorkel', 78: 'solar_dish', 79: 'spider_web', 80: 'stage', 81: 'tank', 82: 'theater_curtain', 83: 'tile_roof', 84: 'tobacco_shop', 85: 'unicycle', 86: 'upright', 87: 'vase', 88: 'wok', 89: 'worm_fence', 90: 'yawl', 91: 'street_sign', 92: 'consomme', 93: 'trifle', 94: 'hotdog', 95: 'orange', 96: 'cliff', 97: 'coral_reef', 98: 'bolete', 99: 'ear'}
Dataset({
    features: ['image', 'label'],
    num_rows: 1000
})
Dataset({
    features: ['image', 'label'],
    num_rows: 200
})

注意到他是个嵌套的数据结构

Problem 1b

对于得到的数据集, 我们总是希望把它写成一个PyTorch的Dataset类的子类, 好处在于可以传递给DataLoader类, 各种操作都很方便, 而且:

-- 可以做批次, 打乱顺序以及并行加载 -- 将数据的预处理写成标准代码, 可读性好

总之就是比自己写一个类好得多

要完成这一点我们至少要继承torch.utils.data.Dataset类并且实现以下三个方法:

class myDataset(Dataset):

    def __init__(self, ...):
        raise NotImplementedError
        # save data
        # initialize preprocessing
        # set up configurations

    def __len__(self):
        raise NotImplementedError
        # return total number(rows) of dataset

    def __getitem__(self, idx):
        raise NotImplementedError
        # return a single data and its label from dataset
        # preprocessing will also be done here

__init__和__len__比较容易实现, 讲一下__getitem__, 先分析一下数据结构和他的Hints:

之前拿到的Dataset嵌套结构是:

Dataset({
    features: ['image', 'label'],
    num_rows: 1000
})
Dataset({
    features: ['image', 'label'],
    num_rows: 200
})

如果想获得一个image和label, 应该要:

# 在__getitem__里面
# 假设self.dataset已经初始化完毕
sample = self.dataset[idx]
img = sample['image']
label = sample['label']

接下来要分出RGB channel, 按照Hint里面要对image用convert方法

img = img.convert("RGB")

查了一下API, 这个image是PIL Image对象, 有一个convert方法转化出指定的颜色模式

接下来还需要实现一个show_images方法来画图, 这个就没什么多说的, 也就是通过类别找到图像, 切割出一些index去plot

完整代码如下:

class MiniImageNetDataset(Dataset):
    def __init__(
        self,
        dataset,
        class_id_to_name: dict[int, str],
        split: str = "train",
        transform: transforms.Compose = None,
    ):
        """
        Args:
            dataset (datasets.Dataset): HuggingFace dataset object
                Each sample in the dataset has:
                - An "image" field that contains the image data
                - A "label" field that contains the numeric class ID
            class_id_to_name (dict): dictionary mapping numeric class IDs to class names
            split (str): "train" or "val" or "test"
        """
        # TODO: Implement the __init__ method of the MiniImageNetDataset
        self.split = split
        self.dataset = dataset
        self.transform = transform
        self.class_id_to_name = class_id_to_name

    def __len__(self):
        # TODO: Implement the __len__ method of the MiniImageNetDataset
        return len(self.dataset)

    def __getitem__(self, index):
        """
        Returns a single sample from the dataset at the given index.
        Args:
            index (int): index of the item to get
        Returns:
            tuple: (image, label) where image is a tensor of shape C x H x W and label is a tensor of shape 1
        """

        # TODO: Implement the __getitem__ method of the MiniImageNetDataset
        # Don't forget to...
        # - Cast the item and label to torch.Tensor with the correct dtype
        # - Permute the image into the shape (C, H, W)
        # - Apply the transform if it's not None
        sample = self.dataset[index]
        img = sample['image']
        label = sample['label']

        img = img.convert("RGB")
        
        if self.transform is not None:
            img=self.transform(img)

        label = torch.tensor(label,dtype=torch.long)
        return img, label

    def show_images(self, num_images: int = 10):
        """
        Visualize images in a grid layout.

        Args:
            num_images (int): Number of images to display per class

        Returns:
            tuple: (fig, ax) - Matplotlib figure and axes objects
        """
        # TODO: Implement the show_images method
        # Make sure to call fig.show() or plt.show() to display your grid of images at the end!
        
        # 所有类别 ID，按从小到大排序，方便排版
        class_ids = sorted(self.class_id_to_name.keys())
        n_classes = len(class_ids)

        # 创建子图：每一行一个类别，每行 num_images 张图
        fig, axes = plt.subplots(n_classes, num_images, figsize=(num_images * 2, n_classes * 2))

        # 当只有 1 行或 1 列时，axes 不是二维数组，做个统一处理
        if n_classes == 1:
            axes = [axes]
        if num_images == 1:
            axes = [[ax] for ax in axes]

        # 对每个类别画图
        for row, cls in enumerate(class_ids):
            # 找到属于该类别的样本下标
            idxs = [i for i, y in enumerate(self.dataset["label"]) if y == cls]
            # 只取前 num_images 个
            idxs = idxs[:num_images]

            for col, idx in enumerate(idxs):
                img = self.dataset[idx]["image"]  # 通常是 PIL.Image
                ax = axes[row][col]
                ax.imshow(img)
                ax.axis("off")
                # 每一行第一个子图上写类别名字
                if col == 0:
                    ax.set_title(self.class_id_to_name[cls])

        plt.tight_layout()
        plt.show()  # 或 fig.show()
        return fig, axes

Problem 1c

先来看看如何构造数据的变换, 这里利用的是:

import torchvision.transforms as transforms

从这里(https://docs.pytorch.org/vision/0.8/transforms.html) 可以查到各种变换的用法, 题目里直接用他给的就行了, 为了把好几个变换做成变换序列, 需要用transforms.Compose, 或者(按文档中的说法)用torch.nn.Sequential也可以

train_transform = transforms.Compose([
    transforms.Resize(256),                  # Resize the shorter side to 256
    transforms.RandomCrop(224),              # Random crop into shape 224×224
    transforms.RandomHorizontalFlip(),       # Random horizontal flip
    transforms.ToTensor(),                   # Convert to tensor, scale to [0,1], and permute into shape (C, H, W)
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]) # Normalize
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

train_dataset = MiniImageNetDataset(mini_imagenet_train, class_id_to_name, split="train", transform=train_transform)
val_dataset = MiniImageNetDataset(mini_imagenet_val, class_id_to_name, split="val", transform=val_transform)

从最后两行可以看出, 用这种继承Dataset的方式构造数据集很方便, 构造完毕后写成一个DataLoader也仅需一行代码

train_dataloader = DataLoader(train_dataset,batch_size=32,shuffle=True)
val_dataloader = DataLoader(val_dataset,batch_size=32,shuffle=True)


# Print the first batch of data and labels
for batch in train_dataloader: # each batch is a tuple of (data, labels)
    data = batch[0] # data is a tensor of shape (B, 3, 224, 224)
    labels = batch[1] # labels is a tensor of shape (B,)
    print(f"Shape of data: {data.shape}")
    print(f"Shape of labels: {labels.shape}")
    break

现在我们就可以通过循环去访问train_dataloader里面的batch, 再通过下标去访问数据和标签了

Problem 1d

实现CNN的训练

训练循环的抽象

理论上所有训练循环的抽象代码都差不多, 示例如下:

Step1: 把模型转移到device上(cpu/gpu)

model.to(device)

Step2: 决定num_epochs并启动外层循环, 决定要遍历数据集多少次

for epoch in range(num_epochs):
    # training logic

Step3: 外层循环内每次都要把model设置为train模式

model.train()

Step4: 启动内层循环, 遍历数据集的x和y, 这个内层循环一次处理一个batch

for x, y in train_dataloader:
    x = x.to(device, dtype = torch.float)
    y = y.to(device, dtype = torch.long)

Step5: 清零梯度

optimizer.zero_grad()

Step6: 前向传播

y_hat = model(x)

Step7: 计算损失

loss = criterion(y_hat, y)

这里的criterion是预设的损失函数, 比如说nn.CrossEntropyLoss

Step8: 反向传播

loss.backward()

Step9: 更新参数

optimizer.step()

Step10: 记录损失

train_loss = loss.item()
preds = torch.argmax(y_hat, dim=1)
train_correct += (preds == y).sum().item()

Step11: 每个epoch结束之后计算`val_dataset`上的准确率, 这段代码是在内层循环外面的

```python
model.eval()
with torch.no_grad():
    for x,y in val_dataloader:
        x = x.to(device, dtype = torch.float)
        y = y.to(device, dtype = torch.long)

        y_hat = model(x)
        loss = criterion(y_hat, y)
        val_loss += loss.item()
        preds = torch.argmax(y_hat, dim=1)
        val_correct += (preds == y).sum().item()

val_loss /= len(val_dataloader)

完整的抽象代码如下:

# === SETUP ===
model.to(device)

# === EPOCH LOOP ===
for epoch in range(num_epochs):
    train_loss = 0.0
    train_correct = 0.0
    
    # === TRAINING PHASE ===
    model.train()  # Enable training mode
    
    for x, y in train_dataloader:
        # 1. Move data to device
        x, y = x.to(device), y.to(device)
        
        # 2. Reset gradients
        optimizer.zero_grad()
        
        # 3. Forward pass
        y_hat = model(x)
        
        # 4. Compute loss
        loss = criterion(y_hat, y)
        
        # 5. Backward pass (compute gradients)
        loss.backward()
        
        # 6. Update weights
        optimizer.step()
        
        # 7. Track metrics (optional)
        # ...
    
    # === VALIDATION PHASE ===
    model.eval()  # Disable training-specific layers
    with torch.no_grad():  # Don't compute gradients
        for x, y in val_dataloader:
            # Forward pass only, no backward pass!
            y_hat = model(x)
            loss = criterion(y_hat, y)
            # Track validation metrics

在一个epoch内, 可能会有很多batch, 每个batch我们会计算一次损失并且加入到train_loss里面, 直到这个epoch结束, 用train_loss/len(train_dataloader)就是平均损失, val_loss只在一个epoch结束之后算一次, 举例如下:

假设训练集有 1000 个样本, batch_size = 100, 那么每个 epoch 有 10 个 batch。在训练过程中, 每处理一个 batch, 我们计算一次该 batch 的 loss(比如分别是 0.85, 0.72, 0.68, ..., 0.55), 并累加到 train_loss 中。Epoch 结束后，用 train_loss / len(train_dataloader) 得到平均训练损失(比如 6.5 / 10 = 0.65). 而 val_loss 只有在整个 epoch 的训练结束后才会计算, 它反映模型在验证集上的表现。

完整代码如下, 依葫芦画瓢即可:

# Common practice to check if GPU is available
if torch.accelerator.is_available():
    device = torch.accelerator.current_accelerator()
else:
    device = "cpu"
print(f"Using device: {device}")
assert str(device) in ["cuda", "mps"], "Please make sure you are using an accelerator (e.g. GPU or MPS)"


def train(
    model, optimizer, criterion, num_epochs, train_dataloader, val_dataloader=None
):
    """
    Args:
        model: the model to train
        optimizer: the optimizer to use
        criterion: the loss function to use
        num_epochs: the number of epochs to train for
        train_dataloader: the dataloader for the training set
        val_dataloader: the dataloader for the validation set
    """
    # === SETUP ===
    model.to(device)

    # == (Optional) Perform a validation loop before training as a baseline ==
    if val_dataloader:
        model.eval()  # Set model to evaluation mode
        num_correct = 0.0
        val_loss = 0.0
        with torch.no_grad():  # Don't compute gradients
            for x, y in val_dataloader:
                x = x.to(device, dtype=torch.float)
                y = y.to(device, dtype=torch.long)
                y_hat = model(x)

                loss = criterion(y_hat, y)  # pred, target
                val_loss += loss.item()

                preds = torch.argmax(y_hat, dim=1)
                num_correct += (preds == y).sum().item()

        val_acc = num_correct / len(val_dataloader.dataset)
        tqdm.write(f"Initial validation loss: {val_loss:.4f}")
        tqdm.write(f"Initial validation accuracy: {val_acc:.4f}")

    # Lists to store metrics across epochs
    train_accuracies = []
    train_losses = []
    val_accuracies = []
    val_losses = []

    epoch_pbar = tqdm(range(num_epochs), desc="Training Progress", position=0)

    # === EPOCH LOOP ===
    for epoch in epoch_pbar:
        # === TRAINING PHASE ===
        model.train()  # Set model to training mode

        # Initialize metrics for this epoch
        train_loss = 0.0
        train_correct = 0.0

        # === INNER LOOP (iterate over training batches) ===
        for x, y in train_dataloader:
            # 1. Move data to device
            x = x.to(device, dtype=torch.float)
            y = y.to(device, dtype=torch.long)

            # 2. Reset gradients
            optimizer.zero_grad()

            # 3. Forward pass
            y_hat = model(x)

            # 4. Compute loss
            loss = criterion(y_hat, y)

            # 5. Backward pass (compute gradients)
            loss.backward()

            # 6. Update weights
            optimizer.step()

            # 7. Track metrics
            train_loss += loss.item()  # Accumulate loss (convert to Python float)
            preds = torch.argmax(
                y_hat, dim=1
            )  # Get predicted class for each item in the batch (highest prob = highest confidence)
            train_correct += (preds == y).sum().item()  # Count correct predictions

        # === END OF INNER LOOP ===

        # Compute average metrics for the epoch
        avg_train_loss = train_loss / len(
            train_dataloader
        )  # Average loss = total loss / total number of batches
        train_acc = (
            train_correct / len(train_dataloader.dataset)
        )  # Average accuracy = accuracy / total number of data points seen across the whole epoch (i.e. in the entire dataset)
        train_losses.append(avg_train_loss)
        train_accuracies.append(train_acc)

        print(f"\nEpoch {epoch + 1} - Training accuracy: {train_acc:.3f}")
        print(f"Epoch {epoch + 1} - Training loss: {train_loss:.3f}")

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        if val_dataloader:
            model.eval()  # Set model to evaluation mode
            val_loss = 0.0
            val_correct = 0.0

            with torch.no_grad():  # Don't compute gradients
                # === INNER LOOP (iterate over validation batches) ===
                for x, y in val_dataloader:
                    # 1. Move data to device
                    x = x.to(device, dtype=torch.float)
                    y = y.to(device, dtype=torch.long)

                    # 2. Forward pass only, no backward pass!
                    y_hat = model(x)

                    # 3. Compute loss
                    loss = criterion(y_hat, y)

                    # 4. Track validation metrics
                    val_loss += loss.item()
                    preds = torch.argmax(y_hat, dim=1)
                    val_correct += (preds == y).sum().item()

                # === END OF INNER LOOP ===

                # Compute average metrics for the epoch
                avg_val_loss = val_loss / len(val_dataloader)
                val_acc = val_correct / len(val_dataloader.dataset)

                val_losses.append(avg_val_loss)
                val_accuracies.append(val_acc)

            # === END OF VALIDATION PHASE ===

        # === END OF EPOCH LOOP ===

        # Print metrics for the epoch
        print(f"Epoch {epoch + 1} - Validation accuracy: {val_acc:.3f}")
        print(f"Epoch {epoch + 1} - Validation loss: {val_loss:.3f}")

        epoch_pbar.set_postfix(
            {
                "train_loss": f"{train_loss:.3f}",
                "train_acc": f"{train_acc:.3f}",
                "val_loss": f"{val_loss:.3f}",
                "val_acc": f"{val_acc:.3f}",
            }
        )

    return train_accuracies, train_losses, val_accuracies, val_losses

调用自定义的训练循环

现在需要实现train函数里面的参数

实现模型实例:

cnn = CNN(num_classes=10)

定义循环轮次:

num_epochs = 20

定义优化器:

optimizer = AdamW(cnn.parameters(),lr=0.0005)

定义损失函数:

criterion = torch.nn.CrossEntropyLoss()

开始训练

cnn_train_accuracies, cnn_train_losses, cnn_val_accuracies, cnn_val_losses = train(
    model=cnn,
    optimizer=optimizer,
    criterion=criterion,
    num_epochs=num_epochs,
    train_dataloader=train_dataloader,
    val_dataloader=val_dataloader
)

Initial validation loss: 16.1283
Initial validation accuracy: 0.1000
Training Progress:   0%|          | 0/20 [00:00<?, ?it/s]
Epoch 1 - Training accuracy: 0.238
Epoch 1 - Training loss: 68.580
Training Progress:   5%|▌         | 1/20 [00:09<03:01,  9.53s/it, train_loss=68.580, train_acc=0.238, val_loss=14.099, val_acc=0.295]Epoch 1 - Validation accuracy: 0.295
Epoch 1 - Validation loss: 14.099

Epoch 20 - Training accuracy: 0.599
Epoch 20 - Training loss: 36.430
Training Progress: 100%|██████████| 20/20 [03:13<00:00,  9.67s/it, train_loss=36.430, train_acc=0.599, val_loss=11.613, val_acc=0.525]Epoch 20 - Validation accuracy: 0.525
Epoch 20 - Validation loss: 11.613

Problem 1e

写一个画图函数来画损失曲线和准确率曲线

def plot_metrics(train_losses, val_losses, train_accuracies=None, val_accuracies=None, num_epochs=None, title=""):
    """
    Plots the training loss, training accuracy, validation loss, and validation accuracy.
    Args:
        train_losses: list of training losses
        val_losses: list of validation losses
        train_accuracies: list of training accuracies
        val_accuracies: list of validation accuracies
        num_epochs: number of epochs
        title: title of the plot
    """
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # 训练损失
    axes[0, 0].plot(train_losses, label='Training Loss')
    axes[0, 0].set_title('Training Loss')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')

    axes[0, 1].plot(train_accuracies, label='Training Accuracy')
    axes[0, 1].set_title('Training Accuracy')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy')

    axes[1, 0].plot(val_losses, label='Validation Loss')
    axes[1, 0].set_title('Validation Loss')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('Loss')

    axes[1, 1].plot(val_accuracies, label='Validation Accuracy')
    axes[1, 1].set_title('Validation Accuracy')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('Accuracy')

    plt.tight_layout()
    plt.suptitle(title)
    plt.show()
    
    
    
plot_metrics(cnn_train_losses, cnn_val_losses, cnn_train_accuracies, cnn_val_accuracies, num_epochs=20, title="CNN Training Metrics")

残差块

如果我们的某一层神经网络需要学习的函数是f(x), 我们可以让他学习g(x) = f(x) - x, 然后在输出端再加上x即可(即输出f(x) = g(x) + x)

这有个显然的好处, 对于第一层的梯度:

$$ \frac{\partial L}{\partial x_1} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x_1} = \frac{\partial L}{\partial y} \cdot \left( \frac{\partial F_n}{\partial x_1} + 1 \right) = \frac{\partial L}{\partial y} \cdot \frac{\partial F_n}{\partial x_n} \cdot \frac{\partial F_{n-1}}{\partial x_{n-1}} \cdot \cdots \cdot \frac{\partial F_2}{\partial x_2} \cdot \left( \frac{\partial F_1}{\partial x_1} + 1 \right) $$

最后一项的+1有效防止了梯度消失, 因为梯度为-1是不稳定的, 并不会一直出现, 所以可以保证梯度不会指数级别的衰减到0, 不过也许还是会震荡

Problem 2a

根据给出的架构自己实现ResNet

                    输入 x (in_channels)
                        │
          ┌─────────────┴─────────────┐
          │                           │
          ▼                           │
    ┌─────────┐                       │
    │ Conv1   │ ──→ BatchNorm1 ──→ ReLU
    └────┬────┘                       │
         │                            │
    ┌────┴────┐                       │
    │ Conv2   │ ──→ BatchNorm2        │
    └────┬────┘                       │
         │                            │
    ┌────┴────────────────────────────┐
    │              +                  │  ← 残差相加
    └────┬────────────────────────────┘
         │
      ReLU (out)

卷积层在CNN的实现当中已经解释过了, 说一说归一化层nn.BatchNorm2d, 实际上对于输入维度[B, C, H, W], 他会对每个batch的每个channel(的所有数据点)去计算均值和方差做归一化

# 输入: conv_output shape = [B=4, C=3, H=2, W=2]
# B=batch size, C=channels, H=height, W=width

x = torch.tensor([
    # batch 0
    [[[1, 2],      # channel 0
      [3, 4]],
     [[5, 6],      # channel 1  
      [7, 8]],
     [[9, 10],     # channel 2
      [11, 12]]],
    # batch 1
    [[[2, 3],
      [4, 5]],
     [[6, 7],
      [8, 9]],
     [[10, 11],
      [12, 13]]],
    # batch 2...
    # batch 3...
])

┌─────────────────────────────────────────────────────────────┐
│  BatchNorm2d 对每个通道独立计算：                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 0 (所有batch):                                     │
│  ┌─────────┬─────────┬─────────┬─────────┐                 │
│  │ batch0  │ batch1  │ batch2  │ batch3  │                 │
│  │ [[1,2], │ [[2,3], │ ...     │ ...     │                 │
│  │  [3,4]] │  [4,5]] │         │         │                 │
│  └─────────┴─────────┴─────────┴─────────┘                 │
│        │           │                               × 4    │
│        └───────────┴───────────────────────→               │
│                 展平成: [1,2,3,4, 2,3,4,5, ...]            │
│                                                             │
│  μ₀ = mean([1,2,3,4,2,3,4,5, ...])                         │
│  σ₀² = var([1,2,3,4,2,3,4,5, ...])                         │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 1:                                                 │
│  ┌─────────┬─────────┬─────────┬─────────┐                 │
│  │ [[5,6], │ [[6,7], │ ...     │ ...     │                 │
│  │  [7,8]] │  [8,9]] │         │         │                 │
│  └─────────┴─────────┴─────────┴─────────┘                 │
│        │           │                               × 4    │
│        └───────────┴───────────────────────→               │
│                 展平成: [5,6,7,8,6,7,8,9, ...]             │
│                                                             │
│  μ₁ = mean([5,6,7,8,6,7,8,9, ...])                         │
│  σ₁² = var([5,6,7,8,6,7,8,9, ...])                         │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 2:  同理计算 μ₂, σ₂²                               │
│                                                             │
└─────────────────────────────────────────────────────────────┘

结果：3个通道 → 3个均值(μ₀, μ₁, μ₂) + 3个方差(σ₀², σ₁², σ₂²)

再说一下这里的残差连接, 当x经过主路径几个卷积层的变换后, 此时肯定是不能和原始输入x相加了, 所以需要一个1x1的卷积层来把x的通道数变成和主路径的输出通道数一样, 然后再相加

class ResidualBlock(nn.Module):
    def __init__(
        self,
        in_channels=3,
        out_channels=64,
        kernel_size=3,
        initial_downsample=False, # whether to downsample the input
        verbose=False, # whether to print debug statements
    ):
        super().__init__()
        self.verbose = verbose


        # TODO: Define the first conv layer
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
        kernel_size=kernel_size,padding=1,stride=1 if not initial_downsample else 2)
        # TODO: Define the first batchnorm layer
        self.batchnorm1 = nn.BatchNorm2d(num_features=out_channels)
        # TODO: Define the second conv layer
        self.conv2 = nn.Conv2d(in_channels=out_channels,out_channels=out_channels,kernel_size=kernel_size,
        padding=1,stride=1)
        # TODO: Define the second batchnorm layer
        self.batchnorm2 = nn.BatchNorm2d(num_features=out_channels)
        # TODO: Define the residual connection
        self.residual_connection = nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=1,
        stride=1 if not initial_downsample else 2)

    def forward(self, x):
        # TODO: Implement the forward pass
        conv1 = self.conv1
        if self.verbose:
            print(f"conv1 shape: {conv1.shape}")
        bn1 = self.batchnorm1
        if self.verbose:
            print(f"bn1 shape: {bn1.shape}")

        relu=nn.ReLU()
        z1 = relu(bn1(conv1(x)))

        conv2 = self.conv2
        if self.verbose:
            print(f"conv2 shape: {conv2.shape}")
        bn2 = self.batchnorm2
        if self.verbose:
            print(f"bn2 shape: {bn2.shape}")

        residual = self.residual_connection
        if self.verbose:
            print(f"residual shape: {residual.shape}")

        out = bn2(conv2(z1))+residual(x)
        # Don't forget to apply ReLU activation after adding the residual
        out = relu(out)
        return out

运行维度检查的代码:

x = torch.rand((1, 3, 224, 224))  # 1 image, 3 RGB channels, 224 x 224 pixels
print(f"Shape of input: {x.shape}")

res = ResidualBlock(in_channels=3, out_channels=64, kernel_size=3, initial_downsample=True, verbose=False)
resblock_out = res(x)
print(f"Shape of output: {resblock_out.shape}")
assert resblock_out.shape == torch.Size([1, 64, 112, 112]), f"Expected output shape (1, 64, 112, 112), but got {resblock_out.shape}"

Shape of input: torch.Size([1, 3, 224, 224])
Shape of output: torch.Size([1, 64, 112, 112])

# Verify the shapes of the weights
conv1_shape = res.conv1.weight.shape
print(f"Shape of conv1 weight: {conv1_shape}")
assert conv1_shape == torch.Size([64, 3, 3, 3]), f"Expected Conv1 weights to have shape (64, 3, 3, 3), but got {conv1_shape}"

conv2_shape = res.conv2.weight.shape
print(f"Shape of conv2 weight: {conv2_shape}")
assert conv2_shape == torch.Size([64, 64, 3, 3]), f"Expected Conv2 weights to have shape (64, 64, 3, 3), but got {conv2_shape}"

residual_shape = res.residual_connection.weight.shape
print(f"Shape of residual weight: {residual_shape}")
assert residual_shape == torch.Size([64, 3, 1, 1]), f"Expected residual connection to have shape (64, 3, 1, 1), but got {residual_shape}"

Shape of conv1 weight: torch.Size([64, 3, 3, 3])
Shape of conv2 weight: torch.Size([64, 64, 3, 3])
Shape of residual weight: torch.Size([64, 3, 1, 1])

Problem 2b

刚刚我们已经实现了残差块, 得益于PyTorch的继承和封装机制, 我们可以很方便的再实现ResNet-18

class ResNet18(nn.Module):
    def __init__(self, num_classes, verbose=False):

        # TODO: Instantiate all the layers/stages of ResNet-18
        # Hint: don't forget to call super().__init__() first!
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=2,padding=3)
        self.batchnorm1=nn.BatchNorm2d(num_features=64)
        self.relu=nn.ReLU()
        self.maxpool1=nn.MaxPool2d(kernel_size=3,stride=2,padding=1)

        self.residualblock1=ResidualBlock(in_channels=64,out_channels=64,kernel_size=3)

        self.residualblock2=ResidualBlock(in_channels=64,out_channels=64,kernel_size=3)

        self.residualblock3=ResidualBlock(in_channels=64,out_channels=128,kernel_size=3,initial_downsample=True)

        self.residualblock4=ResidualBlock(in_channels=128,out_channels=128,kernel_size=3)

        self.residualblock5=ResidualBlock(in_channels=128,out_channels=256,kernel_size=3,initial_downsample=True)

        self.residualblock6=ResidualBlock(in_channels=256,out_channels=256,kernel_size=3)

        self.residualblock7=ResidualBlock(in_channels=256,out_channels=512,kernel_size=3,initial_downsample=True)

        self.residualblock8=ResidualBlock(in_channels=512,out_channels=512,kernel_size=3)

        self.avgpool=nn.AdaptiveAvgPool2d(output_size=(1,1))

        self.flatten=nn.Flatten()

        self.fc=nn.Linear(in_features=512,out_features=num_classes)
        
        self.model=nn.Sequential(self.conv1,self.batchnorm1,self.relu,self.maxpool1,self.residualblock1,self.residualblock2,
        self.residualblock3,self.residualblock4,self.residualblock5,self.residualblock6,
        self.residualblock7,self.residualblock8,self.avgpool,self.flatten,self.fc)
        
    def forward(self, x):

        # TODO: Implement the forward pass of ResNet-18
        out=self.model(x)

        return out

看起来很吓人, 不过是搭积木罢了

x = torch.rand((1, 3, 224, 224))  # 1 image, 3 RGB channels, 224 x 224 pixels
print(f"Shape of input: {x.shape}")

resnet18 = ResNet18(num_classes=10, verbose=True)
resnet18_out = resnet18(x)
print(f"Shape of output: {resnet18_out.shape}")
assert resnet18_out.shape == torch.Size([1, 10]), f"Expected output shape (1, 10), but got {resnet18_out.shape}"

Shape of input: torch.Size([1, 3, 224, 224])
Shape of output: torch.Size([1, 10])

Problem 2c

直接调用之前的train函数开始训练即可

# TODO: Instantiate a ResNet18 model
resnet = ResNet18(num_classes=10)

# TODO: Set the number of epochs
num_epochs = 50

# TODO: Instantiate an optimizer
optimizer = AdamW(resnet.parameters(),lr=0.001)

# TODO: Instantiate a loss criterion
criterion = nn.CrossEntropyLoss()

# TODO: Train the model

resnet_train_accuracies, resnet_train_losses, resnet_val_accuracies, resnet_val_losses = train(
    model=resnet,
    optimizer=optimizer,
    criterion=criterion,
    num_epochs=num_epochs,
    train_dataloader=train_dataloader,
    val_dataloader=val_dataloader
)

Problem 2d

同样的, 调用plot_metrics函数画出损失曲线

plot_metrics(resnet_train_losses, resnet_val_losses, resnet_train_accuracies, resnet_val_accuracies, num_epochs=50, title="ResNet Training Metrics")

Transformer

这里要求我们实现Transformer模型, 虽然说这个完整的架构比CS336里面那个要实现的东西多一点, 但实际上简单很多, 因为这里全程都用PyTorch实现, 而那边基本上都是手搓

和之前做ResNet一样, 先从子块开始实现, 然后再搭积木拼装回去

这一部分的更详细实现和原理参看我CS336的Assignment1 Part2的文章

Problem 3a

要求实现softmax, 直接用torch.exp和torch.sum即可

def softmax(x: torch.Tensor):
    """
    Compute the softmax of each element in x.
    """
    x_exp=torch.exp(x)
    x_sum=torch.sum(x_exp,dim=-1,keepdim=True)
    return x_exp/x_sum

Problem 3b

要求实现点积注意力机制, 本质上就是几个矩阵(Tensor)相乘而已

import math
def scaled_dot_product_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask=None, verbose=False):
    """
    Q: matrix of shape (B, N, d_k)
    K: matrix of shape (B, N, d_k)
    V: matrix of shape (B, N, d_v)
    mask: boolean matrix of shape (B, N, N). Values where mask is True will be masked out

    B is the batch size
    N is the number of query/key/value vectors
    d_k is the dimensions of the query/key vectors
    d_v is the dimensions of the value vectors
    """
    assert Q.shape[-1] == K.shape[-1], "Q and K must have same d_k dimension"
    assert K.shape[-2] == V.shape[-2], "K and V must have same sequence length"


    # TODO: Implement the attention function
    return torch.matmul((softmax(torch.matmul(Q,K.transpose(-2,-1))/math.sqrt(Q.shape[-1]))),V)

Problem 3c

要求实现注意力头, 注意这里的输入有两种, 一种是输入x, 然后通过三个变换矩阵分别变换到$$Q,K,V$$, 另一种是在Outputs模块当中, 左边的两个输入要读取Inputs模块的输出, 右边的一个输入为x, 通过一个verbose参数去控制一下就好了

class AttentionHead(nn.Module):
    def __init__(self, d_model=512, d_k=64, verbose=False):

        # TODO: Implement the __init__ method
        # Don't forget to call super().__init__()!
        super().__init__()
        self.W_q=nn.Linear(d_model,d_k)
        self.W_k=nn.Linear(d_model,d_k)
        self.W_v=nn.Linear(d_model,d_k)
        self.verbose=verbose


    def forward(self, x, encoder_output=None, mask=None):
        """
        x is the input to use for the queries, keys, and values
        encoder_output is the output from the encoder (used for cross-attention)
        mask is the mask to use for the attention
        """


        # TODO: Implement the forward method
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        if encoder_output is not None:
            Q = self.W_q(x)
            K = self.W_k(encoder_output)
            V = self.W_v(encoder_output)

        out = scaled_dot_product_attention(Q,K,V,mask=mask,verbose=self.verbose)

        return out

所谓的"三个线性变换到QKV空间"其实就是三个线性层

Problem 3d

要求实现多头机制, 其实只不过是初始化若干个AttentionHead类, 然后把x传入每个类, 最后把每个类的输出拼接起来, 这里只要注意一下每个头的维度均为d_k = d_model / num_heads即可, 还有就是计算完的注意力要乘以一个W_o矩阵

class MultiHeadAttention(nn.Module):
    def __init__(
        self,
        num_heads=8,
        d_model=512,  # dimension of the embeddings
        verbose=False,
    ):
        """
        Args:
            num_heads: number of attention heads
            d_model: dimension of the embeddings
            verbose: whether to print debug information
        """

        # TODO: Implement the __init__ method
        super().__init__()
        self.d_k = d_model//num_heads
        self.heads = nn.ModuleList([AttentionHead(d_model=d_model,d_k=self.d_k,verbose=verbose) for _ in range(num_heads)])
        self.W_out = nn.Linear(d_model,d_model)

    def forward(self, x, encoder_output=None, mask=None):

        # TODO: Implement the forward method
        out = torch.cat([head(x,encoder_output,mask)for head in self.heads],dim=-1)
        out = self.W_out(out)

        return out

Problem 3e

实现架构图的左半部分, 非常纯粹的看图说话, 而且用torch实现这个ffn也非常简单, 不过是一个线性层加上一个ReLU而已

如果想看LayerNorm, Linear, ReLU的实现, 参看我CS336的Assignment1 Part2的文章

class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.multiheadlayer=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.layer_norm1=nn.LayerNorm(d_model)
        self.ffn=nn.Sequential(
            nn.Linear(d_model,d_model),
            nn.ReLU(),
            
        )

    def forward(self, x, mask=None):

        # TODO: Implement the forward method
        # 1. Self-attention
        out = self.multiheadlayer(x,mask=mask)
        # 2. Residual connection
        out = x + out
        # 3. LayerNorm
        out = self.layer_norm1(out)
        norm1_out = out
        # 4. Feedforward network
        out=self.ffn(out)
        # 5. Residual connection
        out=out+norm1_out
        # 6. LayerNorm
        out=self.layer_norm1(out)
        # 7. Return the output
        return out

Problem 3f

实现架构图的右半部分Decoder, 和上面差不多, 不再赘述了, 这里我没有写mask功能, 实际上mask功能就是构造一个上三角为True, 下三角为False的矩阵, 然后让注意力机制在计算注意力的时候, 把下三角的注意力权重设为负无穷, 这样在计算注意力权重的时候, 下三角的注意力权重就会变成0, 从而实现mask功能

具体可以看我CS336的Assignment1 Part2的文章, 那里非常详细的说了这个掩码功能

class DecoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.masked_multiheadlayer1=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.layernorm1=nn.LayerNorm(d_model)
        self.masked_multiheadlayer2=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.ffn=nn.Sequential(
            nn.Linear(d_model,4*d_model),
            nn.ReLU(),
            nn.Linear(4*d_model,d_model)
        )
        self.layernorm2=nn.LayerNorm(d_model)

    def forward(self, x, encoder_output):

        # TODO: Implement the forward method
        mask = torch.triu(torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device), diagonal=1)
        # Create a look-ahead mask

        # 1. Masked self-attention
        out = self.masked_multiheadlayer1(x,mask=mask)
        # 2. Residual connection
        out = x + out
        # 3. LayerNorm
        out = self.layernorm1(out)
        norm1 = out
        # 4. Cross-attention
        out = self.masked_multiheadlayer2(out,encoder_output=encoder_output)
        # 5. Residual connection
        out = out + norm1
        # 6. LayerNorm
        out = self.layernorm2(out)
        norm2=out
        # 7. Position-wise feed forward network
        out = self.ffn(out)
        # 8. Residual connection
        out = out + norm2
        # 9. LayerNorm
        out = self.layernorm2(out)
        # 10. Return the final normalized output
        return out
        # 10. Return the final normalized output

Problem 3g

要求实现RoPE, 这里解释起来比较复杂, 建议看336的notes, 这里我直接给出代码

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=1024):
        """
        max_len = max number of tokens in the sequence
        """
        super().__init__()
        self.pos = torch.zeros((max_len, d_model))
        numerator = torch.arange(end=max_len, dtype=torch.float32).reshape(-1, 1) # shape (seq_len, 1)
        denominator = torch.pow(10000, torch.arange(start=0, end=d_model, step=2, dtype=torch.float32) / d_model).reshape(1, -1) # shape (1, d_model / 2)

        quotient = numerator / denominator # shape (seq_len, d_model). quotient[i, j] = numerator[i, 0] / denominator[0, j]
        self.pos[:, 0::2] = torch.sin(quotient) # assign result of trig to all even indices for every token positions
        self.pos[:, 1::2] = torch.cos(quotient) # assign result of trig to all odd indices for every token position

    def forward(self, x):
        """
        Input x is of shape (B, seq_len, d_embedding)
        """
        out = x + self.pos[:x.shape[1], :].to(device=x.device)
        return out

Problem 3h

拓展一下encode部分, 本质上encode就三件事情: 词嵌入, RoPE, 过含有多头注意力, layernorm, ffn的EncodeLayer

class TransformerEncoder(nn.Module):
    def __init__(
        self, vocab_size=1024, d_model=512, num_layers=6, num_heads=8, verbose=False
    ):

        # TODO: Implement the __init__ method
        super().__init__()
        self.embedding=nn.Embedding(num_embeddings=vocab_size,embedding_dim=d_model)
        self.positional_encoding=PositionalEncoding(d_model=d_model)
        self.encoder_layers=nn.ModuleList([EncoderLayer(d_model=d_model,num_heads=num_heads,verbose=verbose) for _ in range(num_layers)])
        self.num_layers=num_layers
        self.num_heads=num_heads

    def forward(self, x):

        # 1. Embed the input sequence of tokens
        x = self.embedding(x)
        # 2. Add positional encodings
        x = self.positional_encoding(x)
        # 3. Sequentially pass the result through each EncoderLayer
        for layer in self.encoder_layers:
            x = layer(x)
        # 4. Return the final tensor
        return x

Problem 3i

拓展decode部分, 和上面一样

class TransformerDecoder(nn.Module):
    def __init__(self, vocab_size=1024, d_model=512, num_layers=6, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.embedding=nn.Embedding(num_embeddings=vocab_size,embedding_dim=d_model)
        self.positional_encoding=PositionalEncoding(d_model=d_model)
        self.num_layers=num_layers
        self.decoder_layers=nn.ModuleList([DecoderLayer(d_model=d_model,num_heads=num_heads,verbose=verbose) for _ in range(num_layers)])
        self.linear=nn.Linear(d_model,vocab_size)

    def forward(self, x, encoder_output):

        # 1. Embed the input sequence of tokens
        x = self.embedding(x)

        # 2. Add positional encodings
        x = self.positional_encoding(x)

        # 3. Sequentially pass the result through each DecoderLayer
        for layer in self.decoder_layers:
            x = layer(x,encoder_output)

        # 4. Apply the linear layer to predict the next token in our vocabulary
        x = self.linear(x)

        # 5. Return the final tensor
        return x

Problem 3j

全部组装成Transformer类即可

class Transformer(nn.Module):
    def __init__(
        self,
        vocab_size=1024,
        d_model=512,
        num_encoder_layers=6,
        num_decoder_layers=6,
        num_heads=8,
        verbose=False,
    ):

        # TODO: Implement the __init__ method
        super().__init__()
        self.encoder=TransformerEncoder(vocab_size=vocab_size,d_model=d_model,num_layers=num_encoder_layers,num_heads=num_heads,verbose=verbose)
        self.decoder=TransformerDecoder(vocab_size=vocab_size,d_model=d_model,num_layers=num_decoder_layers,num_heads=num_heads,verbose=verbose)

    def forward(self, x):

        # TODO: Implement the forward method
        # 1. Pass `x` through the encoder
        encoder_output = self.encoder(x)
        # 2. Pass `x` and the encoder's output through the decoder
        x = self.decoder(x,encoder_output)
        # 3. Return the output of the decoder
        return x

接下来是"鉴赏"部分, 课程组帮我们准备好了TinyStories数据集并且做好了分词, 只不过这是个word级别的分词, 而我们在336里面实现的是byte级的

Problem 3k

题目让我们处理一下段落, 其实也就是完成以下任务:

-- 按换行把段落分开 -- 每一段用之前给出的extract_full_words得到一个词列表 -- 在段落的开头和结尾加上<start>和<end> -- 对于段落中的每个词, 用word_to_id得到一个ID(毕竟输入模型的不能是词语) -- 跳过每个比max_seq_length还短的段落 (这里不知道是不是课程组笔误了, 应该记作min_seq_length才对) -- 用滑动窗口得到inputs和targets, 保证targets永远比inputs长1(预测下一个token)

def create_sequences(stories, word_to_id, max_seq_length=20, pad_token=PAD_TOKEN, start_token=START_TOKEN, end_token=END_TOKEN):
    """
    Create input/target sequences from paragraphs

    Args:
        stories (list): List of stories to create sequences from
        word_to_id (dict): Dictionary mapping words to their token IDs
        max_seq_length (int): Maximum sequence length to create
        pad_token (str): Token for the pad token
        start_token (str): Token for the start of a sequence
        end_token (str): Token for the end of a sequence

    Returns:
        inputs (tensor): Tensor of input sequences
        targets (tensor): Tensor of target sequences
    """
    inputs = []
    targets = []

    # List of paragraphs
    paragraphs = []

    # TODO: Split the stories into paragraphs and append them to the `paragraphs` list
    for story in stories:
        story_paragraphs = story.split("\n\n")
        paragraphs.extend(story_paragraphs)


    # TODO: Iterate over each paragraph, extract full words, add `<start>` and `<end>` tokens, convert to IDs, and build sequences
    # Make sure to skip paragraphs that are too short to create any sequences
    for paragraph in paragraphs:
        full_words = extract_full_words(paragraph.lower())
        full_words.insert(0,start_token)
        full_words.append(end_token)
        if len(full_words) < max_seq_length:
            continue

        for i in range(len(full_words)-max_seq_length):
            input_seq = full_words[i:i+max_seq_length]
            target_seq = full_words[i+1:i+max_seq_length+1]

            input_seq = [word_to_id.get(word,0) for word in input_seq]
            target_seq = [word_to_id.get(word,0) for word in target_seq]
            inputs.append(input_seq)
            targets.append(target_seq)
            
            

    # TODO: Convert the inputs and targets to tensors of type `torch.long`
    inputs = torch.tensor(inputs,dtype=torch.long)
    targets = torch.tensor(targets,dtype=torch.long)
    return inputs, targets

# Generate sequences
max_seq_length = 20
inputs, targets = create_sequences(tiny_stories_text, word_to_id, max_seq_length)

print(f"Total sequences created: {len(inputs)}")
print(f"Each input sequence length: {len(inputs[0])}")
print(f"Each target sequence length: {len(targets[0])}")

Problem 3l

用刚刚实现的数据集去构造DataLoader即可

# TODO: Split into training and validation datasets (80/20 split)
train_inputs = inputs[:2000]
train_targets = targets[:2000]
val_inputs = inputs[2000:2200]
val_targets = targets[2000:2200]

print(f"Number of training sequences: {len(train_inputs)}")
print(f"Number of validation sequences: {len(val_inputs)}")

# TODO: Create TensorDatasets
train_dataset = TensorDataset(train_inputs,train_targets)
val_dataset = TensorDataset(val_inputs,val_targets)

# TODO: Create DataLoaders
batch_size = 32
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
val_loader = DataLoader(val_dataset,batch_size=batch_size,shuffle=True)

print(f"Training batches: {len(train_loader)}")
print(f"Validation batches: {len(val_loader)}")

batch = next(iter(train_loader)) # Get the first batch from the training DataLoader

train_input, train_target = batch # Unpack the batch into input and target tensors
print(f"First 10 tokens of training input in first batch: {train_input[0][:10]}")
print(f"First 10 tokens of training target in first batch: {train_target[0][:10]}")

Number of training sequences: 2000
Number of validation sequences: 200
Training batches: 63
Validation batches: 7
First 10 tokens of training input in first batch: tensor([4546, 1945, 2585, 1379, 3266, 4146,  629, 2357, 2319, 3820])
First 10 tokens of training target in first batch: tensor([1945, 2585, 1379, 3266, 4146,  629, 2357, 2319, 3820, 3896])

Problem 3m

开始训练

# Transformer training loop
def train_transformer(
    model,
    device,
    optimizer,
    criterion,
    num_epochs,
    train_dataloader,
    val_dataloader,
    vocab,
):
    # === SETUP ===
    # Move model to device
    model.to(device)

    # Lists to store metrics across epochs
    train_losses = []
    val_losses = []

    # === EPOCH LOOP ===
    for epoch in tqdm(range(num_epochs)):
        # === TRAINING PHASE ===
        model.train()  # Set model to training mode

        # Initialize metrics for this epoch
        total_train_loss = 0.0

        # === INNER LOOP (iterate over training batches) ===
        for batch_inputs, batch_targets in tqdm(train_dataloader, desc="Training"):
            # Move inputs and targets to device
            batch_inputs = batch_inputs.to(device)
            batch_targets = batch_targets.to(device)

            # Reset the gradients
            optimizer.zero_grad()

            # Forward pass
            logits = model(batch_inputs)  # [batch_size, seq_len, vocab_size]

            # Reshape for loss calculation
            loss = criterion(logits.reshape(-1, len(vocab)), batch_targets.reshape(-1))
            # Compute loss
            loss.backward()

            # Optional: Add gradient clipping to prevent exploding gradients
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            # Update the model parameters
            optimizer.step()

            # Accumulate the loss
            total_train_loss += loss.item()

        # === END OF INNER LOOP ===

        # Compute the average loss for this epoch
        avg_train_loss = total_train_loss / len(train_dataloader)
        train_losses.append(avg_train_loss)

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        model.eval()  # Set model to evaluation mode
        total_val_loss = 0.0

        with torch.no_grad():
            # === INNER LOOP (iterate over validation batches) ===
            for batch_inputs, batch_targets in tqdm(val_dataloader, desc="Validation"):
                # Move inputs and targets to device
                batch_inputs = batch_inputs.to(device)
                batch_targets = batch_targets.to(device)

                # Forward pass
                logits = model(batch_inputs)

                # Compute loss
                loss = criterion(
                    logits.reshape(-1, len(vocab)), batch_targets.reshape(-1)
                )

                # Accumulate the loss
                total_val_loss += loss.item()
        # === END OF INNER LOOP ===

        # Compute the average loss for this epoch
        avg_val_loss = total_val_loss / len(val_dataloader)
        val_losses.append(avg_val_loss)

        # === END OF VALIDATION PHASE ===

        # Print epoch results
        print(f"\nEpoch {epoch + 1} Results:")
        print(f"  Training Loss: {avg_train_loss:.4f}")
        print(f"  Validation Loss: {avg_val_loss:.4f}")

    # === END OF EPOCH LOOP ===

    return train_losses, val_losses

# TODO: Instantiate a Transformer model

vocab_size=len(vocab)

transformer = Transformer(vocab_size=vocab_size,
                                    d_model=512,
                                    num_encoder_layers=6,
                                    num_decoder_layers=6,
                                    num_heads=8,
                                    verbose=False).to(device)

# TODO: Set the optimizer
optimizer = AdamW(transformer.parameters(),lr=0.0001)

# TODO: Set the loss criterion
criterion = nn.CrossEntropyLoss()

# TODO: Set the number of epochs
num_epochs= 50

# TODO: Train the model
transformer_train_losses, transformer_val_losses = train_transformer(
    model=transformer,
    optimizer=optimizer,
    criterion=criterion,
    num_epochs=num_epochs,
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    device=device,
    vocab=vocab
)

Problem 3n

和之前一样画出损失曲线

# TODO: Plot the training and validation loss across the epochs

def plot_metrics_transformer(train_losses, val_losses, num_epochs=None, title=""):
    """
    Plots the training loss, training accuracy, validation loss, and validation accuracy.
    Args:
        train_losses: list of training losses
        val_losses: list of validation losses
        train_accuracies: list of training accuracies
        val_accuracies: list of validation accuracies
        num_epochs: number of epochs
        title: title of the plot
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 8))

    # 训练损失
    axes[0].plot(train_losses, label='Training Loss')
    axes[0].set_title('Training Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')


    axes[1].plot(val_losses, label='Validation Loss')
    axes[1].set_title('Validation Loss')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Loss')



    plt.tight_layout()
    plt.suptitle(title)
    plt.show()
    
plot_metrics_transformer(train_losses=transformer_train_losses,val_losses=transformer_val_losses,
num_epochs=50,title="Transformer Training Metrics")

print(f"Final Training Loss: {transformer_train_losses[-1]:.4f}")
print(f"Final Validation Loss: {transformer_val_losses[-1]:.4f}")

Problem 3o

接下来就可以让我们的模型帮我们生成点东西了, 自由发挥即可

# TODO: Generate your own tiny story!
sample_text = "I love you"


sample_text = generate_sample(
    model=transformer,
    word_to_id=word_to_id,
    id_to_word=id_to_word,
    max_seq_len=20,
    prompt=sample_text,   
    max_length=50
)


print(f"Sample generation: {sample_text}")

Sample generation: shoes love you re lucky you re lucky you re lucky you re lucky you have a shield timmy looked at his shirt and saw that he was wearing his favorite superhero shirt he felt better and said thanks billy you re a good friend <end>

UC Berkeley CS189 Assignment 3

Sat, 07 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 3

项目描述

这个lab主要聚焦于实现优化器和反向传播(逐项求导), 和CS336的lab1当中的实现类似

例子: 多元微分

我们用课程组给出的例子来复习一下多元微分

中间变量是从v1到v7

$$ \frac{\partial v_7}{\partial v_7} = 1 $$

$$ v_7 = v_6 - v_5 \implies \frac{\partial v_7}{\partial v_6} = 1, \quad \frac{\partial v_7}{\partial v_5} = -1 $$

$$ v_6 = v_4 + v_3 \implies \frac{\partial v_6}{\partial v_4} = 1, \quad \frac{\partial v_6}{\partial v_3} = 1 $$

$$ \frac{\partial v_7}{\partial v_4} = \frac{\partial v_7}{\partial v_6} \cdot \frac{\partial v_6}{\partial v_4} = 1 \cdot 1 = 1 $$

$$ \frac{\partial v_7}{\partial v_3} = \frac{\partial v_7}{\partial v_6} \cdot \frac{\partial v_6}{\partial v_3} $$

$$ \frac{\partial v_6}{\partial v_3} = \frac{\partial v_6}{\partial v_3} + \frac{\partial v_6}{\partial v_4} \cdot \frac{\partial v_4}{\partial v_3} = 1 + e^{v_3} $$

$$ \frac{\partial v_7}{\partial v_3} = 1 + e^{v_3} $$

$$ \frac{\partial v_7}{\partial v_2} = \frac{\partial v_7}{\partial v_6} \cdot \frac{\partial v_6}{\partial v_2} + \frac{\partial v_7}{\partial v_5} \cdot \frac{\partial v_5}{\partial v_2} = \frac{\partial v_6}{\partial v_2} - \frac{\partial v_5}{\partial v_2} = \frac{\partial v_6}{\partial v_2} - \cos(v_2) $$

$$ \frac{\partial v_6}{\partial v_2} = \frac{\partial v_6}{\partial v_4} \cdot \frac{\partial v_4}{\partial v_2} + \frac{\partial v_6}{\partial v_3} \cdot \frac{\partial v_3}{\partial v_2} = \frac{\partial v_4}{\partial v_2} + \frac{\partial v_3}{\partial v_2} = \frac{\partial v_4}{\partial v_3} \cdot \frac{\partial v_3}{\partial v_2} + v_1 = e^{v_3} \cdot v_1 + v_1 $$

$$ \frac{\partial v_7}{\partial v_1} = \frac{\partial v_7}{\partial v_6} \cdot \frac{\partial v_6}{\partial v_1} = \frac{\partial v_6}{\partial v_1} = \frac{\partial v_6}{\partial v_4} \cdot \frac{\partial v_4}{\partial v_1} + \frac{\partial v_6}{\partial v_3} \cdot \frac{\partial v_3}{\partial v_1} = \frac{\partial v_4}{\partial v_1} + \frac{\partial v_3}{\partial v_1} = \frac{\partial v_4}{\partial v_3} \cdot \frac{\partial v_3}{\partial v_1} + \frac{\partial v_3}{\partial v_1} \ = e^{v_3} \cdot v_2 + v_2 $$

所谓反向传播就是把每一次算出来的中间结果记住, 然后用链式法则一步步的向前求导,比如说最初的输入是x,y, 输出是z, 通过计算z对一大堆中间变量的偏导数, 最终计算得到 $$ \frac{\partial z}{\partial x} ,,and ,, \frac{\partial z}{\partial y} $$

如果对上面的求偏导过程有疑问, 建议去看一看微积分(下)或者数学分析(三)

Problem 1

既然要求微分, 首先要实现Tensor之间的运算, 注意到正如上面那个图所示, 我们需要三个类:

-- class BearGrad 来给出一个可调用的函数, 根据上游梯度来计算下游梯度 -- class BearParent 来记录图的边, 即每个节点的parent节点是什么 -- class BearTensor 来提供各种基础运算, 例如加减乘除

@dataclass
class BearGrad:
    '''Stores how to compute the downstream gradient from the upstream gradient 
    - `op_str` is just a string describing the operation; you don't need to use this for this HW, but it may help in debugging if you print out what operations create your computation graph.
    - `fn` is a function that takes in the upstream loss gradient and outputs the loss gradient that should be passed downstream. This essentially applies the chain rule at the current node in the computation graph.
    '''
    fn: callable
    op_str: str | None = None

这是梯度计算器, fn函数用于计算下游梯度, op_str用来描述这个操作

例如我们有一个中间变量c = a + b

a ────┐
      ├──→ c = a + b ──→ 上游梯度 (∂L/∂c)
b ────┘

这个BearGrad类需要给出L对a和b的偏导数, 即∂L/∂a和∂L/∂b, 这种情况下fn大概是一个加法的梯度函数:

def add_grad_fn(upstream_grad):
    # ∂L/∂a = ∂L/∂c * ∂c/∂a = ∂L/∂c * 1
    # ∂L/∂b = ∂L/∂c * ∂c/∂b = ∂L/∂c * 1
    return upstream_grad  # 因为 ∂c/∂a = 1

父节点追踪器:

@dataclass
class BearParent:
    '''This class represents a parent of a node; a node must track its parents so that it knows who to propagate its gradients to.
    - `grad` is the `BearGrad` object for this parent; we can apply the `fn` from this to compute the gradient that should be passed downstream.
    - `parent` is the `BearTensor` object for the parent
    '''
    parent: BearTensor
    grad: BearGrad

    def __hash__(self):
        return id(self)

    def __eq__(self, other):
        return self is other

要记录父节点和grad, grad代表了"在已知上游梯度的情况下, 如何算出下游梯度"

顺便解释一下类装饰器@dataclass, 用了这个装饰器就会自动生成一些方法, 比如:

__init__
__repr__ (字符串表示)
__eq__ (相等比较)

等等

class BearTensor:
    '''`BearTensor`: This represents a node in our computation graph
    - `value` is the underlying data in the Tensor; this is computed during the forward pass
    - `parents` keeps track of the parents in the computation graph to whom we pass our gradient to
    - `adjoint` is the gradient (the transpose of it technically) that we compute in the backward pass
    '''
    def __init__(self, name: str, value: np.ndarray, parents: list[BearParent] | None = None):
        self.name = name
        self.value = value
        self.parents = parents if parents is not None else []
        self.adjoint: float | np.ndarray = 0.0

我们要实现各种运算和相对于那种运算的梯度公式

- Addition (`__add__`)
- Subtraction (`__sub__`)
- Multiplication (`__mul__`)
- Power (`__pow__`)
- Matrix multiplication (`__matmul__`)
- Dot product (`dot`)
- Sum (`sum`)
- Mean (`mean`)
- ReLU (`relu`)
- Sigmoid (`sigmoid`)

先来看__add__的实现, 这个运算要对self.value和other.value进行加法运算, 新BearTensor.value就是self.value + other.value

接下来思考这个问题: 返回的新的BearTensor的parents应该是什么?

return BearTensor(name=f"{self.name} + {other.name}", value=new_value, parents=parents)

新的parents应该有两个节点, 分别为self和other, 但是parent必须为BearParent类, 而构造这个类的时候需要两个参数: BearTensor和BearGrad, 显然BearTensor已经有了, BearGrad需要我们手动演算一下求导的函数

设有两个 Tensor:

$a$: BearTensor, 值为 $a$
$b$: BearTensor, 值为 $b$
$c = a + b$: BearTensor, 值为 $c = a + b$

我们需要计算 $\frac{\partial L}{\partial a}$ 和 $\frac{\partial L}{\partial b}$，其中 $L$ 是最终损失函数：

$$ \begin{align*} \frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial a} \

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial b} \

\frac{\partial c}{\partial a} = \frac{\partial (a + b)}{\partial a} = 1 \

\frac{\partial c}{\partial b} = \frac{\partial (a + b)}{\partial b} = 1\

\frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot 1 = \frac{\partial L}{\partial c}\

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot 1 = \frac{\partial L}{\partial c} \end{align*} $$

换言之, 这里的两个下游梯度都等于传过来的上游梯度, 所以两个fn函数的返回值保持不变

反向传播就是已知L对c的梯度, 求L对a和b的梯度, 即∂L/∂a和∂L/∂b

    def __add__(self, other: BearTensor) -> BearTensor:
        if self.value.shape != other.value.shape:
            raise ValueError("Shapes must match")
        
        new_value = self.value + other.value

        def grad_fn(upstream_grad):
            return upstream_grad

        def grad_fn_other(upstream_grad):
            return upstream_grad

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("add")')),
            BearParent(parent=other, grad=BearGrad(fn=grad_fn_other,op_str='text("add")'))
        ]

        return BearTensor(name=f"{self.name} + {other.name}", value=new_value, parents=parents)

减法也一样

$$ \begin{align*} \frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial a} \

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial b} \

\frac{\partial c}{\partial a} = \frac{\partial (a - b)}{\partial a} = 1 \

\frac{\partial c}{\partial b} = \frac{\partial (a - b)}{\partial b} = -1\

\frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot 1 = \frac{\partial L}{\partial c}\

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot -1 = -\frac{\partial L}{\partial c} \end{align*} $$

乘法 $$ \begin{align*} \frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial a} \

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial b} \

\frac{\partial c}{\partial a} = \frac{\partial (a \cdot b)}{\partial a} = b \

\frac{\partial c}{\partial b} = \frac{\partial (a \cdot b)}{\partial b} = a \

\frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot b \

\frac{\partial L}{\partial b} = \frac{\partial L}{\partial c} \cdot a \end{align*} $$

    def __mul__(self, other: BearTensor) -> BearTensor:
        if self.value.shape != other.value.shape:
            raise ValueError("Shapes must match")
        
        new_value = self.value * other.value

        def grad_fn(upstream_grad):
            return other.value*upstream_grad
        
        def grad_fn_other(upstream_grad):
            return self.value*upstream_grad

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("mul")')),
            BearParent(parent=other, grad=BearGrad(fn=grad_fn_other,op_str='text("mul")'))
        ]

        return BearTensor(name=f"{self.name} * {other.name}", value=new_value, parents=parents)

幂运算, 此时另外一个输入变成固定的数power, 记作n

$$ \begin{align*} \frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot \frac{\partial c}{\partial a} \

\frac{\partial c}{\partial a} = \frac{\partial (a^n)}{\partial a} = n \cdot a^{n-1} \

\frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot n \cdot a^{n-1} \

\end{align*} $$

    def __pow__(self, power: float) -> BearTensor:
        
        new_value=self.value**power

        def grad_fn(upstream_grad):
            return upstream_grad*power*self.value**(power-1)

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("pow")'))
        ]

        return BearTensor(name=f"{self.name}**{power}", value=new_value, parents=parents)

矩阵乘法, 注意这里要用向量微积分

$$ \begin{align*} \frac{\partial L}{\partial A} &: (m,n) \ \frac{\partial L}{\partial C} &: (m,p) \ C &= A \cdot B : (m,p) \ A &: (m,n), \quad B : (n,p) \[8pt]

\frac{\partial L}{\partial A} &= \frac{\partial L}{\partial C} \cdot B^T \ (m,n) &= (m,p) \cdot (p,n) \[8pt]

\frac{\partial L}{\partial B} &= A^T \cdot \frac{\partial L}{\partial C} \ (n,p) &= (n,m) \cdot (m,p) \end{align*} $$

    def __matmul__(self, other: BearTensor) -> BearTensor:
        # This may be helpful to avoid shape mismatch errors
        def ensure_2d(x):
            '''If x is a scalar or 1-D, convert to 2-D'''
            x = np.asarray(x)
            if x.ndim == 0:          # scalar
                return x.reshape(1, 1)
            elif x.ndim == 1:        # vector
                return x.reshape(-1, 1)  # column vector convention
            else:                     # already 2D
                return x
        
        new_value=np.matmul(ensure_2d(self.value), ensure_2d(other.value))

        def grad_fn(upstream_grad):

            return np.matmul(upstream_grad,other.value.T)

        def grad_fn_other(upstream_grad):

            return np.matmul(self.value.T,upstream_grad)

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("matmul")')),
            BearParent(parent=other, grad=BearGrad(fn=grad_fn_other,op_str='text("matmul")'))
        ]

        return BearTensor(name=f"{self.name} @ {other.name}", value=new_value, parents=parents)

记住向量求导的尺寸约定: 偏导数(矩阵)的尺寸保持为"分母"的尺寸, 比如$$\frac{\partial L}{\partial A}$$的尺寸为$$(m,n)$$, 因为$$A$$的尺寸为$$(m,n)$$

点积运算

$$ \begin{align*} c &= \mathbf{a} \cdot \mathbf{b} = \sum_{i=1}^n a_i b_i, \quad \mathbf{a}{(n)}, \mathbf{b}{(n)} \[8pt]

\frac{\partial c}{\partial \mathbf{a}} &= \mathbf{b} \[6pt]

\frac{\partial L}{\partial \mathbf{a}}{(n)} &= \frac{\partial L}{\partial c}{(1)} \cdot \mathbf{b}_{(n)} \[12pt]

\frac{\partial c}{\partial \mathbf{b}} &= \mathbf{a} \[6pt]

\frac{\partial L}{\partial \mathbf{b}}{(n)} &= \frac{\partial L}{\partial c}{(1)} \cdot \mathbf{a}_{(n)} \end{align*} $$

    def dot(self, other: BearTensor) -> BearTensor:
        if self.value.ndim != 1 or other.value.ndim != 1:
            raise ValueError("dot() only supports 1-D BearTensors (like torch.dot).")

        new_value=np.dot(self.value,other.value)

        new_value=np.array([new_value])
            
        def grad_fn(upstream_grad):

            return upstream_grad*other.value

        def grad_fn_other(upstream_grad):

            return upstream_grad*self.value

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("dot")')),
            BearParent(parent=other, grad=BearGrad(fn=grad_fn_other,op_str='text("dot")'))
        ]

        return BearTensor(name=f"{self.name} @ {other.name}", value=new_value, parents=parents)

自求和 $$ \begin{align*} s &= \sum_{i=1}^n x_i = \text{sum}(\mathbf{x}), \quad \mathbf{x}{(n)} \implies s{(1)} \[8pt]

\frac{\partial s}{\partial \mathbf{x}} &= \begin{bmatrix} 1 & 1 & \cdots & 1 \end{bmatrix}_{(n)} \[6pt]

\frac{\partial L}{\partial \mathbf{x}}{(n)} &= \frac{\partial L}{\partial s}{(1)} \cdot \mathbf{1}_{(n)} \end{align*} $$

    def sum(self) -> BearTensor:
        
        new_value=np.array([np.sum(self.value)])

        def grad_fn(upstream_grad):
            return np.ones_like(self.value)*upstream_grad

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("sum")'))
        ]

        return BearTensor(name=f"sum({self.name})", value=new_value, parents=parents)

自平均运算

$$ \begin{align*} \mu &= \text{mean}(\mathbf{x}) = \frac{1}{n} \sum_{i=1}^n x_i, \quad \mathbf{x}{(n)} \implies \mu{(1)} \[8pt]

\frac{\partial \mu}{\partial \mathbf{x}} &= \frac{1}{n} \begin{bmatrix} 1 & 1 & \cdots & 1 \end{bmatrix}_{(n)} \[6pt]

\frac{\partial L}{\partial \mathbf{x}}{(n)} &= \frac{\partial L}{\partial \mu}{(1)} \cdot \frac{1}{n} \cdot \mathbf{1}_{(n)} \end{align*} $$

    def mean(self) -> BearTensor:
        
        new_value=np.array([np.mean(self.value)])

        def grad_fn(upstream_grad):
            return np.ones_like(self.value)*upstream_grad/self.value.size

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("mean")'))
        ]

        return BearTensor(name=f"mean({self.name})", value=new_value, parents=parents)

ReLU运算

$$ \begin{align*} y &= \text{ReLU}(x) = \max(0, x) \[8pt]

\frac{\partial y}{\partial x} &= \begin{cases} 1 & \text{if } x > 0 \ 0 & \text{if } x \leq 0 \end{cases} = \mathbb{1}(x > 0) \[8pt]

\frac{\partial L}{\partial x} &= \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} = \frac{\partial L}{\partial y} \cdot \mathbb{1}(x > 0) \end{align*} $$

    def relu(self) -> BearTensor:
        
        new_value=np.maximum(0, self.value)

        def grad_fn(upstream_grad):
            return (self.value>0)*upstream_grad

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("relu")'))
        ]

        return BearTensor(name=f"relu({self.name})", value=new_value, parents=parents)

注意这个求导在数学上是有瑕疵的, ReLU函数在0处并不可导

Sigmoid函数

$$ \begin{align*} y &= \sigma(x) = \frac{1}{1 + e^{-x}} \[8pt]

\frac{dy}{dx} &= \sigma(x)(1 - \sigma(x)) = y(1 - y) \[8pt]

\frac{\partial L}{\partial x} &= \frac{\partial L}{\partial y} \cdot \frac{dy}{dx} = \frac{\partial L}{\partial y} \cdot y(1 - y) \end{align*} $$

    def sigmoid(self) -> BearTensor:
        
        new_value=1 / (1 + np.exp(-self.value))

        def grad_fn(upstream_grad):
            return upstream_grad*new_value*(1-new_value)

        parents=[
            BearParent(parent=self, grad=BearGrad(fn=grad_fn,op_str='text("sigmoid")'))
        ]

        return BearTensor(name=f"sigmoid({self.name})", value=new_value, parents=parents)

课程组给了一个链式图的demo绘图代码, 图如下

当然, 我们可以修改里面的代码, 比如说绘制一个更复杂的多元函数的链式图

# # Example code
# a = BearTensor("a", np.array([2, 3]))
# b = BearTensor("b", np.array([1, 1]))
# c = a + b
# draw_graph(c, "demo_graph.typ")

a = BearTensor("a", np.array([2, 3]))
b = BearTensor("b", np.array([1, 1]))
c = BearTensor("c", np.array([1, 4]))
d = a * b + b * c + a * c
draw_graph(d, "demo_graph.typ")

Problem 2

要实现反向传播, 首先要对所有的节点进行拓扑排序, 这样才能搞清楚传播路径

def topological_sort(node):
    '''Return a list of nodes ordered topologically'''

    visited = set()
    sorted_nodes= []
    
    def dfs(current_node):
        if current_node in visited:
            return

        visited.add(current_node)

        for parent_info in current_node.parents:
            dfs(parent_info.parent)

        sorted_nodes.append(current_node)

    dfs(node)
    return sorted_nodes

搞清楚这个结点添加的顺序, 举一个简单的例子:

dfs(loss):
  loss 不在 visited → 添加
  遍历 loss.parents → [ReLU]
    dfs(ReLU):
      ReLU 不在 visited → 添加
      遍历 ReLU.parents → [x2]
        dfs(x2):
          x2 不在 visited → 添加
          遍历 x2.parents → [x1]
            dfs(x1):
              +1 不在 visited → 添加
              遍历 +1.parents → [x]
                dfs(x):
                  x 不在 visited → 添加
                  遍历 x.parents → []
                  sorted_nodes.append(x)   ✓
              sorted_nodes.append(x1)      ✓
            sorted_nodes.append(x2)         ✓
          sorted_nodes.append(ReLU)        ✓
      sorted_nodes.append(loss)             ✓

最终: [x, x1, x2, ReLU, loss]

这里x是最初的输入, 没有任何依赖, 我们最终需要的梯度就是loss对x的梯度, 为了得到这个梯度, 需要计算一系列loss对中间变量(x1, x2, ReLU)的梯度, 然后通过链式法则得到loss对x的梯度

还需要从某个节点开始, 重置他和他上游(父)节点的所有梯度, 这也需要用递归实现

def reset_children(self):
    """Resets the gradient in the current node to zero and all nodes before it in the computation graph."""
    
    def reset_gradient(node):
        if node.adjoint is not None:
            node.adjoint = 0

        for parent in node.parents:
            reset_gradient(parent.parent)

    reset_gradient(self)

现在实现反向传播, 思考一下, 以上面那个例子当中的[x, x1, x2, ReLU, loss], 遍历的顺序应该是从loss开始, 而不是从x开始

遍历到每个节点x, 都需要再遍历x的父节点, 并且更新父节点的梯度

def backward(self):
    """
    Take a node in the computation graph, reset all gradients, and perform backpropagation 
    to compute the adjoints (gradients) for all nodes in the graph.

    Hint: After resetting the gradients, what should the gradient at the current node be?
    """
    self.reset_children()

    sorted_nodes=topological_sort(self)

    sorted_nodes.reverse()

    self.adjoint=1.0

    for node in sorted_nodes:
        
        for parent_info in node.parents:
            parent=parent_info.parent
            grad_fn=parent_info.grad.fn

            parent_grad=grad_fn(node.adjoint)

            parent.adjoint+=parent_grad

Problem 3

实现SGD, Momentum, 实现AdamW优化器

首先考虑优化器的基类, 需要有一个学习率和一堆的参数, 参数是保存为list[BearTensor], value属性保存参数值, adjoint属性保存梯度

class Optimizer:
    def __init__(self, params: list[BearTensor], lr: float):
        self.params = params
        self.lr = lr

    def zero_grad(self):
        for p in self.params:
            p.adjoint = np.zeros_like(p.value)

    def step(self):
        # raise NotImplementedError
        for p in self.params:
            p.value -= self.lr * p.adjoint

SGD直接照抄基类就行

class SGD(Optimizer):
    def __init__(self, params: list[BearTensor], lr: float):
        super().__init__(params, lr)

    def step(self):
        
        for p in self.params:
            p.value -= self.lr * p.adjoint
        # pass

注意Momentum优化器的公式:

$$ \boxed{ \begin{align*} m_t &= \beta \cdot m_{t-1} + g_t \quad \text{(动量项)} \[6pt] v_t &= -\eta \cdot m_t \quad \text{(速度/增量)} \[6pt] w_{t+1} &= w_t + v_t = w_t - \eta \cdot m_t \quad \text{(参数更新)} \end{align*} } $$

先把这些notation映射到我们类里的属性:

$$ \begin{array}{c|c|c} \text{符号} & \text{代码对应} & \text{含义} \ \hline g_t & \texttt{p.adjoint} & 当前梯度 \ m_t & \texttt{momentum_term} & 动量项 \ v_t & \texttt{self.velocities[i]} & 速度/参数增量 \ \beta & \texttt{self.beta} & 动量系数 \ \eta & \texttt{self.lr} & 学习率 \ w_t & \texttt{p.value} & 参数值 \end{array} $$

接下来没有任何难度了, 在step方法里遍历参数然后逐个更新就行

class Momentum(Optimizer):
    def __init__(self, params: list[BearTensor], lr: float, beta: float = 0.9):
        super().__init__(params, lr)
        self.beta = beta
        self.velocities = [np.zeros_like(p.value) for p in self.params]

    def step(self):
        for i, p in enumerate(self.params):
            if self.lr != 0:
                prev_momentum = self.velocities[i] / (-self.lr)
            else:
                prev_momentum = 0
            

            momentum_term = prev_momentum * self.beta + p.adjoint
            

            self.velocities[i] = -self.lr * momentum_term
            

            p.value += self.velocities[i]

对Adam优化器也做同样的理解

$$ \boxed{ \begin{align*} \text{1. 一阶矩:} \quad & m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \[6pt] \text{2. 二阶矩:} \quad & v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \[6pt] \text{3. 偏差校正:} \quad & \hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}t = \frac{v_t}{1 - \beta_2^t} \[6pt] \text{4. 参数更新:} \quad & w{t+1} = w_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \varepsilon} \end{align*} } $$

$$ \begin{array}{c|c|c} \text{符号} & \text{代码对应} & \text{含义} \ \hline g_t & \texttt{p.adjoint} & 当前梯度 \ m_t & \texttt{self.ms[i]} & 一阶矩估计（动量） \ v_t & \texttt{self.vs[i]} & 二阶矩估计（方差） \ \hat{m}_t & \texttt{m_hat} & 校正后一阶矩 \ \hat{v}_t & \texttt{v_hat} & 校正后二阶矩 \ \beta_1 & \texttt{self.beta1} & 一阶矩衰减系数 (0.9) \ \beta_2 & \texttt{self.beta2} & 二阶矩衰减系数 (0.999) \ \varepsilon & \texttt{self.eps} & 防止除零 (1e-8) \ t & \texttt{self.t} & 时间步 \ \eta & \texttt{self.lr} & 学习率 \end{array} $$

class Adam(Optimizer):
    def __init__(
        self,
        params: list[BearTensor],
        lr: float,
        beta1: float = 0.9,
        beta2: float = 0.999,
        eps: float = 1e-8,
    ):
        super().__init__(params, lr)
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.t = 0
        self.ms = [np.zeros_like(p.value) for p in self.params]
        self.vs = [np.zeros_like(p.value) for p in self.params]

    def step(self):
        self.t+=1
        for i,p in enumerate(self.params):
            self.ms[i]=self.beta1*self.ms[i]+(1-self.beta1)*p.adjoint

            self.vs[i]=self.beta2*self.vs[i]+(1-self.beta2)*((p.adjoint)**2)

            m_hat=self.ms[i]/(1-self.beta1**self.t)
            v_hat=self.vs[i]/(1-self.beta2**self.t)

            p.value-=self.lr*m_hat/(np.sqrt(v_hat)+self.eps)

运行课程组的代码查看三个优化器的训练曲线

Problem 4

这里要我们自己写一个训练循环, 唯一要注意的就是把所有的参数注册成BearTensor类, 其他没有什么难点了

from sklearn.datasets import fetch_openml
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt

# DO NOT CHANGE THE PREPROCESSING CODE
data = fetch_openml("wine-quality-red", as_frame=True)
X = data.data.to_numpy()
y = data.target.to_numpy().astype(float)

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.4, random_state=42, shuffle=True
)

X_tensor = BearTensor("X_train", X_train)               # shape: (N_train, input_dim)
y_tensor = BearTensor("y_train", y_train.reshape(-1,1)) # shape: (N_train,1)

# Add your training code here!
np.random.seed(42)

input_dim=X_train.shape[1] # 输入维度
hidden_dim=32 # 隐藏层维度
output_dim=1 # 输出维度

W1_init=np.random.randn(input_dim,hidden_dim)*np.sqrt(2.0/input_dim) # 第一个hidden_layer矩阵, input_dim * hidden_dim
W2_init=np.random.randn(hidden_dim,output_dim)*np.sqrt(2.0/input_dim) # 第二个hidden_layer矩阵, hidden_dim * output_dim

W1=BearTensor("W1",W1_init.copy()) # 把参数矩阵注册成为BearTensor类参数
W2=BearTensor("W2",W2_init.copy()) # 同上

params=[W1,W2]  # 注意我们自己实现的Adam类接受的params是一个list[BearTensor]

optimizer=Adam(params,lr=0.01) # 创建优化器, 学习率为0.01

num_epochs=100 # 训练轮数

train_losses=[] # 训练损失
test_losses=[] # 测试损失

for epoch in range(num_epochs):

    optimizer.zero_grad() # 清零梯度

    hidden=(X_tensor@W1).sigmoid() # 前向传播
    output=hidden@W2 # 同上
    
    loss=((output-y_tensor)**2).mean() # 计算损失

    loss.backward()
    optimizer.step() # 更新参数

    train_loss=loss.value.item() # 记录训练损失
    train_losses.append(train_loss) # 记录训练损失     

    # 计算测试损失（不需要梯度）
    X_test_tensor = BearTensor("X_test", X_test)
    y_test_tensor = BearTensor("y_test", y_test.reshape(-1, 1))
    
    # 测试前向传播（使用训练好的权重）
    hidden_test = (X_test_tensor @ W1).sigmoid()
    output_test = hidden_test @ W2
    test_loss = ((output_test - y_test_tensor) ** 2).mean()
    test_losses.append(test_loss.value.item())

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss.value.item():.4f}")

# 打印最终MSE
final_mse = test_losses[-1]
print(f"\nFinal Test MSE: {final_mse:.4f}")

# 绘制损失曲线
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.title('Training and Test Loss')
plt.legend()
plt.grid(True)
plt.show()

# --- Prediction function ---
def predict(x):
    """Output the scalar prediction for a single training datapoint x"""
    x_tensor=BearTensor("x_input",x.reshape(1,-1))

    hidden=(x_tensor@W1).sigmoid()
    output=hidden@W2

    return output.value.item()
    # pass

Epoch 10/100, Train Loss: 11.8084, Test Loss: 10.9215
Epoch 20/100, Train Loss: 4.0232, Test Loss: 3.5952
Epoch 30/100, Train Loss: 0.9623, Test Loss: 0.8673
Epoch 40/100, Train Loss: 0.4591, Test Loss: 0.5058
Epoch 50/100, Train Loss: 0.5181, Test Loss: 0.5702
Epoch 60/100, Train Loss: 0.4653, Test Loss: 0.4989
Epoch 70/100, Train Loss: 0.4054, Test Loss: 0.4402
Epoch 80/100, Train Loss: 0.3933, Test Loss: 0.4299
Epoch 90/100, Train Loss: 0.3896, Test Loss: 0.4261
Epoch 100/100, Train Loss: 0.3834, Test Loss: 0.4217

Final Test MSE: 0.4217

UC Berkeley CS189 Assignment 2(Part 2)

Thu, 05 Feb 2026 00:00:00 GMT

CS189 Assignment 2

项目描述

书接上回, 这个lab里我们主要关心两个问题:

如何对模型的能力建模, 即如何通过PK当中模型的输赢情况以及一些潜在的其他特征, 通过逻辑回归的方式给每个模型一个能力分数
新特征对模型能力的影响, 如果加入一些额外的特征, 模型的排名会发生什么变化

Part1代码复用

这里要用到subselect_battles函数, 直接从part1复制过来就行, 不再赘述了

Problem 4

概率建模

我们希望先搞一个LeaderBoard, 对于每一个模型m, 能给他一个"strength score" $$S_m$$, 这个score能反应:

-- score的排名反映了此模型对战另一个模型时的胜率

-- 量化的看, 对于模型A和B, A战胜B的概率应该是$$S_A - S_B$$的函数

换言之, 我们希望找到一个函数$$f$$使得: $$ P(A ,, beats ,, B) = f(S_A - S_B) $$

比如说用sigmoid函数: $$ P(A ,, beats ,, B) = \frac{1}{1 + e^{-(S_A - S_B)}} $$

本质上这就是学习一个logistic regression

数据结构设计

可以把一行数据表示为这种数据结构:

-- 一个feature vector表示两个model

-- 一个label表示胜出的模型

举个例子, 原始的行为:

row = {'model_a': 'gpt-4o-2024-05-13', 'model_b': 'claude-3-opus-20240229', 'winner': 'model_a'}`

可表述为这两行(两行是因为PK是相互的, A赢了B同时也代表B赢了A)

Feature 1:[1, -1]
Label: 1

index 0 表示gpt, index 1表示claude, 1表示gpt胜出

Feature 2:[-1, 1]
Label: 0

index 0 表示gpt, index 1表示claude, 0表示claude未胜出

通俗一点的说就是如果Label是站在feature = 1的视角来看的, 如果1赢了-1, 那么Label就是1, 如果1输了, 那么Label就是0

为什么同一行要解释成两个不同的feature vector呢? 对于含有A,B模型的一行, 我们可能要建模 $$ P(A ,, beats ,, B) = \sigma(S_A - S_B) P(B ,, beats ,, A) = \sigma(S_B - S_A) $$

显然这两个feature vector就是[S_A - S_B, S_B - S_A]的系数矩阵

$$ \begin{bmatrix} 1 & -1 \ -1 & 1 \end{bmatrix} $$

多个模型的情况

假如有多个模型, 也只不过是把所有的Feature Vector和Label(以前面那种方法得到的)组成X和y而已

比如说有5个模型

索引:    0         1          2        3         4
模型:  [GPT-4o, Claude-3, Gemini, Llama-3, PaLM-2]

如果GPT赢了Claude, 那么对应的Feature Vector和Label应该是

[+1, 0, 0, -1, 0], 1
[-1, 0, 0, +1, 0], 0

没参加PK的模型的值赋为0即可

Problem 4a

遍历每一行, 对于每一行再去遍历所有的models, 遍历到model_a就记为1, model_b就记为-1, 其他的就记为0, 再根据winner打个标就行

def turn_into_features(df, models):
    '''
    Convert pairwise battle results into feature matrix X and label vector y
    suitable for logistic regression based on the Bradley-Terry model
    '''
    # TODO:
    # 1. Iterate through each row in the DataFrame.
    # 2. For each battle, create a feature vector:
    #    - Assign +1 to the column corresponding to 'model_a'.
    #    - Assign -1 to the column corresponding to 'model_b'.
    #    - All other columns should be 0.
    # 3. Append the label:
    #    - 1 if 'model_a' is the winner.
    #    - 0 if 'model_b' is the winner.
    # 4. Return the feature matrix X and label vector y as numpy arrays.
    X = []
    y = []


    for _ ,row in df.iterrows():
        a=row['model_a']
        b=row['model_b']

        win_a=1 if row['winner']==a else 0

        x_ab=[1 if col==a else -1 if col==b else 0 for col in models]

        y_ab=win_a

        x_ba=[1 if col==b else -1 if col==a else 0  for col in models]

        y_ba=1-win_a

        X.append(x_ab)
        y.append(y_ab)
        X.append(x_ba)
        y.append(y_ba)
    return np.array(X), np.array(y)

X, y = turn_into_features(selected_battles_no_ties, selected_models)
X.shape, y.shape

注意为什么两个y是互斥的, 因为都是站在每一个feature = 1的视角去看, 如果第一个feature = 1赢了-1, 那么第二个feature = 1肯定就输了, 所以一定是互斥的

array([[ 0,  0,  0, ...,  0,  0, -1],
       [ 0,  0,  0, ...,  0,  0,  1],
       [ 0,  0,  0, ...,  0,  1,  0],
       ...,
       [-1,  0,  0, ...,  0,  0,  0],
       [ 0,  0,  0, ...,  0,  0,  0],
       [ 0,  0,  0, ...,  0,  0,  0]], shape=(51066, 20))

这个标签的设计是合理的, 假设feature vector是[1,-1]而A赢了, 那么label是1, 此时当然也希望$$\sigma(S_A - S_B)$$是接近于1的

Problem 4b

用LogisticRegression来训练即可

from sklearn.linear_model import LogisticRegression


model = LogisticRegression(fit_intercept=False)
scores = model.fit(X,y).coef_[0]

results = {"Model": selected_models, "Score": scores}
results_df = pd.DataFrame(results).sort_values("Score", ascending=False).reset_index(drop=True)
results_df

解释一下为什么对于两个模型A和B, 输出 $$ \sigma(S_A - S_B) $$ 就代表A战胜B的概率, 这和我们之前设计的特征工程有关, 比如说训练集上所有情况下A都战胜了B, 那么feature vector和label都为[1,-1]和1, 也就是说训练得到的[w_1, w_2]大概率会满足 $$ \sigma( \begin{pmatrix} 1,-1 \end{pmatrix}

\begin{pmatrix} w_1 \ w_2 \end{pmatrix}) = 1 $$

由于特征工程的设计，权重向量 $\mathbf{w}$ 中的每个元素 $w_i$ 就是模型 $i$ 的 strength $S_i$。因此 $\sigma(S_A - S_B)$ 直接给出了 A 战胜 B 的预测概率。

Problem 5a

我们要用重采样的方式来重复训练逻辑回归模型, 每个逻辑回归模型会给我们这些LLM一组分数, 在得到许多组分数之后, 我们就可以得到每个LLM的分数和置信区间

举个例子, 每次训练完成之后, 逻辑回归模型会给gpt4o这个模型一个分数, 假如我们训练了一百次逻辑回归, 那就会有一百个分数, 我们可以取这个分数的均值, 还可以取2.5%分位数和97.5%分位数, 得到置信区间

def get_bootstrapped_score(X, y, models, category_name="Overall", n_bootstrap=10):
    """
    Bootstraps logistic regression model scores to estimate confidence intervals.
    Args:
        X: Feature matrix
        y: Labels
        models: List of model names (order matches columns of X)
        n_bootstrap: Number of bootstrap samples
    Returns:
        results_df: DataFrame with Model, Average Score, Lower Bound, Upper Bound
        mean_scores: Mean of bootstrapped scores (np.array)
        confidence_intervals: 2.5 and 97.5 percentiles (np.array shape [2, n_models])
    """
    #TODO

    np.random.seed(189)  # for reproducibility
    bootstrap_scores = []
    for i in range(n_bootstrap):
        indices = np.random.choice(len(X), size=len(X), replace=True)
        X=X[indices]
        y=y[indices]
        model = LogisticRegression(fit_intercept=False)
        model.fit(X,y)
        bootstrap_scores.append(model.coef_[0])
    bootstrap_scores = np.array(bootstrap_scores)
    mean_scores = bootstrap_scores.mean(axis=0)
    confidence_intervals = np.percentile(bootstrap_scores, [2.5, 97.5], axis=0)

    results = {
        "Model": models,
        "Average Score": mean_scores,
        "Lower Bound": confidence_intervals[0],
        "Upper Bound": confidence_intervals[1],
        "Category": category_name,
    }
    results_df = pd.DataFrame(results)
    return results_df, mean_scores, confidence_intervals
results_df, mean_scores, confidence_intervals = get_bootstrapped_score(X, y, selected_models, n_bootstrap=10)

# Test that confidence intervals make sense
assert (confidence_intervals[0] <= confidence_intervals[1]).all(), "Every lower bound must be <= upper bound."
assert ((confidence_intervals[0] <= mean_scores) & (mean_scores <= confidence_intervals[1])).all(), "Each mean score should lie within its CI."

代码写起来也简单, 只要注意一下np.percentile的用法就行

可视化每个模型的置信区间:

在之前的单次实验当中, 虽然llama-3-70b-instruct的分数最高, 但是从这个图看来, 最好的应该是gemini-1.5-pro-exp-0801(我主观上也这么认为)

Problem 5b

有了置信区间, 我们可以比较保守的给出一个A模型好于B模型的定义, 即若A的Lower Bound大于B的Upper Bound, 那么我们就认为A好于B, 我们可以根据这个规则给出rank

注意这里的代码需要遍历两次, 对于每个模型, 还要遍历除了他之外的所有模型, 逐个比较强弱关系, 我们用count来记录A模型超过了多少其他模型, 注意这个条件十分严格, 一定要置信区间下界大于其他模型的上界才行

def assign_rank(row, df=results_df):
    """
    Input:
        row : pd.Series
            A row of the DataFrame (representing a model’s metrics).
        df : pd.DataFrame (default = results_df)
            DataFrame containing model performance with 'Lower Bound' and 'Upper Bound'.

    Output:
        int : The rank of the model, defined as (# of models confidently better) + 1.
    """

    count = 0
    for _,row2 in df.iterrows():
        if row['Lower Bound'] > row2['Upper Bound']:
            count+=1

    return count


results_df['Rank'] = results_df.apply(lambda r: assign_rank(r, results_df), axis=1)
results_df = results_df.sort_values(by="Rank", ascending=True)
results_df

注意结果上有很多比较好的模型都并列rank = 0了

Problem 7a

假设我们现在收到的对话语料的格式为:

conv_example = {
    "conversation": [
        {"role": "user", "content": "How do I sum a list in Python?"},
        {"role": "assistant", "content": "Use the built-in function: sum(your_list)."}
    ]
}

需要计算所有role == assistant的content的token长度, 这里题目告诉我们调用gpt2的分词器即可

import tiktoken
def calculate_response_length(conv):
    enc=tiktoken.get_encoding('gpt2')

    assistant_contents=[msg['content'] for msg in conv['conversation'] if msg['role']=='assistant']

    joined_text="\n\n".join(assistant_contents)

    tokens=enc.encode(joined_text,disallowed_special=())

    return len(tokens)

先按照条件筛选, 然后把每条句子join起来, 用分词器编码之后返回长度即可

Problem 7b

回忆一下我们的PK数据的格式, 每一行当中有这四列:

model_a model_b conversation_a conversation_b

现在希望把model_a和conversation_a分出来, 计算一下conversation_a的token长度(用7a实现的那个函数), 对conversation_b也做同样的操作, 然后竖着concat起来, 相当于以前有n行, 处理之后变成2n行了

#TODO: 
battles_a = selected_battles_no_ties[['conversation_a','model_a']].rename(columns={'conversation_a':'conversation','model_a':'model'})

battles_a['response_length']=battles_a['conversation'].apply(lambda c:calculate_response_length({'conversation':c}))

battles_b = selected_battles_no_ties[['conversation_b','model_b']].rename(columns={'conversation_b':'conversation','model_b':'model'})

battles_b['response_length']=battles_b['conversation'].apply(lambda c:calculate_response_length({'conversation':c}))

per_model_battles=pd.concat([battles_a,battles_b])
per_model_battles.head()

注意这个conversation是个比较复杂的嵌套结构, 可以理解为他是一个list, 当中的每一个值是User/Assistant说的一句话

[{'content': 'Is there any Artificial Superintelligence?', 'num_tokens': 7, 'role': 'user'}
 {'content': "As a large language model, I am not able to provide an opinion or belief on whether or not artificial superintelligence exists. The concept of artificial superintelligence is a topic of much debate and speculation.\n\nHere's what we know:\n\n* **Current AI:** Today's AI, including myself, is very good at specific tasks, like generating text, translating languages, and playing games. However, we lack the general intelligence and common sense of humans.\n* **Superintelligence Definition:** Artificial superintelligence is generally defined as an AI that surpasses human intelligence in all aspects. This includes things like problem-solving, creativity, and learning.\n* **No Consensus:** There is no scientific consensus on whether or not artificial superintelligence is possible or even desirable. Some experts believe it is inevitable, while others believe it is highly unlikely or even dangerous.\n\nThe development of artificial superintelligence is a complex and multifaceted issue with profound implications for society. It is important to continue to have open and honest discussions about the potential benefits and risks of this technology.", 'num_tokens': 208, 'role': 'assistant'}
 {'content': 'What is your knowledge cutoff date?', 'num_tokens': 7, 'role': 'user'}
 {'content': "I can access pretty up-to-date information, which means I don't really have a knowledge cut-off date. However, keep in mind:\n\n* **I am not constantly updating:** I don't have real-time information like a search engine. My knowledge is based on the massive dataset I was trained on.\n\n* **Training data has a limit:** That training dataset does have a point in time when it was last updated. While it's extensive, there will always be newer information I'm not aware of.\n\nSo, while I strive to be as current as possible, it's best to double-check anything time-sensitive with a reliable source.", 'num_tokens': 127, 'role': 'assistant'}]

按照之前得到的LLM的rank顺序, plot每个模型的token长度:

model_lineup = results_df.sort_values("Rank")['Model'].tolist()
avg_lengths = per_model_battles.groupby("model")["response_length"].mean().reset_index()
avg_lengths["model"] = pd.Categorical(avg_lengths["model"], categories=model_lineup, ordered=True)
avg_lengths = avg_lengths.sort_values("model").reset_index(drop=True)

# Add a numeric rank column for trendline fitting
avg_lengths["rank"] = avg_lengths.index + 1  # 1 = best, etc.

# Fit a linear trendline (polyfit) to the response length vs. rank
z = np.polyfit(avg_lengths["rank"], avg_lengths["response_length"], 1)
p = np.poly1d(z)
trendline = p(avg_lengths["rank"])

fig = go.Figure()

fig.add_trace(go.Scatter(
    x=avg_lengths["model"],
    y=avg_lengths["response_length"],
    mode='lines+markers',
    name='Avg Response Length'
))

fig.add_trace(go.Scatter(
    x=avg_lengths["model"],
    y=trendline,
    mode='lines',
    name='Trendline',
    line=dict(dash='dash', color='red')
))

fig.update_layout(
    title="Average Response Length of Models (with Trendline)",
    xaxis_title="Model (sorted by performance)",
    yaxis_title="Average Response Length",
    xaxis_tickangle=45,
    yaxis=dict(range=[0, max(avg_lengths["response_length"].max(), trendline.max()) * 1.05])  # y-axis starts at 0
)

fig.show()

可以看出模型好坏和对话的token长度并无直接关系

Problem 8a

从conv_metadata字段当中拿到一些额外的特征, 先给出一个正则化函数normdiff, 相当于计算出两个model的某一个特征得分之差

$$ \text{normdiff}(a, b) = \begin{cases} 0 & \text{if } a + b = 0 \[6pt] \dfrac{a - b}{a + b} & \text{otherwise} \end{cases} $$

先看一下在这个嵌套字段conv_metadata当中有什么

selected_battles_no_ties['conv_metadata'].iloc[0]

{'bold_count_a': {'**': 5, '__': 0},
 'bold_count_b': {'**': 0, '__': 0},
 'context_a_tokens': 222,
 'context_b_tokens': 230,
 'header_count_a': {'h1': 0, 'h2': 0, 'h3': 0, 'h4': 0, 'h5': 0, 'h6': 0},
 'header_count_b': {'h1': 0, 'h2': 0, 'h3': 0, 'h4': 0, 'h5': 0, 'h6': 0},
 'list_count_a': {'ordered': 0, 'unordered': 5},
 'list_count_b': {'ordered': 0, 'unordered': 0},
 'sum_assistant_a_tokens': 335,
 'sum_assistant_b_tokens': 287,
 'sum_user_tokens': 14,
 'turns': 2}

总之就是定义了一些特征, 然后分别在两个model上去计算这些特征, 需要注意的一点是, 特征的计算方式是求和计算, 比如说我想通过blod_count_a计算得到bold_a, 那在上面这个例子上应该是5+0=5

def add_style_features(df):
    """
    Adds normalized style feature difference columns to the DataFrame.
    The columns added are:
      - style_bold_count
      - style_header_count
      - style_list_count
      - style_sum_assistant_tokens
    """
    def normdiff(a, b):
        denom = a + b
        return 0 if denom == 0 else (a - b) / denom

    
    style_bold = []
    style_header = []
    style_list = []
    style_tokens = []
    for idx, row in df.iterrows():

        meta=row['conv_metadata']

        bold_a=sum(meta['bold_count_a'].values())
        bold_b=sum(meta['bold_count_b'].values())

        header_a=sum(meta['header_count_a'].values())
        header_b=sum(meta['header_count_b'].values())

        list_a=sum(meta['list_count_a'].values())
        list_b=sum(meta['list_count_b'].values())

        tokens_a=meta['sum_assistant_a_tokens']
        tokens_b=meta['sum_assistant_b_tokens']

        style_bold.append(normdiff(bold_a,bold_b))
        style_header.append(normdiff(header_a,header_b))
        style_list.append(normdiff(list_a,list_b))
        style_tokens.append(normdiff(tokens_a,tokens_b))



        
        
    df["style_bold_count"] = style_bold
    df["style_header_count"] = style_header
    df["style_list_count"] = style_list
    df["style_sum_assistant_tokens"] = style_tokens
    return df

# Example usage:
selected_battles_no_ties = add_style_features(selected_battles_no_ties)

Problem 8b

接下来我们要重构一下数据的格式, 把原来的一行拆成两行, 例如

这应该比较简单, 别忘了我们已经有了style_feature_cols那些字段, 我们只要和之前一样做出X,y和direction, 然后for循环style_feature_cols拿到我们刚才通过norm_diff计算的一些额外特征就行了

def make_pairwise_feature_df(df, models, style_feature_cols):
    """
    For each row in df, create two rows in the output:
      - One for A->B (original direction)
      - One for B->A (flipped direction)
    Each row contains:
      - question_id
      - X: model indicator vector (1 for model_a, -1 for model_b, 0 otherwise)
      - y: 1 if model_a wins, 0 if model_b wins
      - style features (from style_feature_cols)
      - direction: "A->B" or "B->A"
    Returns a new DataFrame with only these columns.
    """
    ...
    records = []
    for idx, row in df.iterrows():
        # X vector for A->B and B->A
        x_a_b=[1 if row['model_a']==model else -1 if row['model_b']==model else 0 for model in models]
        x_b_a=[-1 if row['model_a']==model else 1 if row['model_b']==model else 0 for model in models]
        # y for A->B and B->A
        y_a_b=1 if row['winner']=='model_a' else 0
        y_b_a=1 if  y_a_b==0 else 1
        # Style features for A->B, B->A
        style_ab=row[style_feature_cols].to_numpy(dtype=float)
        style_ba=-style_ab

        rec_a={
          "question_id":row['question_id'],
          "X":x_a_b,
          "y":y_a_b,
          "direction":"A->B",
        }

        for col,val in zip(style_feature_cols,style_ab):
          rec_a[col]=val

        rec_b={
          "question_id":row['question_id'],
          "X":x_b_a,
          "y":y_b_a,
          "direction":"B->A",
        }

        for col,val in zip(style_feature_cols,style_ba):
          rec_b[col]=val


        
        # Add A->B, B->A
        records.append(rec_a)
        records.append(rec_b)
        
    
    return pd.DataFrame(records)

style_feature_cols = [
    "style_bold_count",
    "style_header_count",
    "style_list_count",
    "style_sum_assistant_tokens"
]

pairwise_feature_df = make_pairwise_feature_df(selected_battles_no_ties, selected_models, style_feature_cols)
pairwise_feature_df.head(2)

Problem 8c

根据上面的图, 现在我们有了一些额外的特征, 我们把这些额外的特征接到X向量的后面, 相当于得到了一些额外的feature

唯一的难点在于特征的拼接, 我们需要先把原来的X列通过hstack竖着拼接成二维数组, 然后把额外特征列to_numpy成二位数组, 再用hstack水平拼接起来

参考以下这个可视化变换:

# 假设 df['X'] 是这样的：
# 0    [1, -1, 0, 0]
# 1    [0, 1, -1, 0]
# 2    [-1, 0, 1, 0]

X_id = np.vstack(df['X'].values)
# 结果：
# array([[ 1, -1,  0,  0],
#        [ 0,  1, -1,  0],
#        [-1,  0,  1,  0]])

# style_features = ['style_bold_count', 'style_header_count', ...]

X_style = df[style_features].to_numpy()
# 结果：
# array([[ 0.5,  0.2,  0.3,  0.1],
#        [ 0.1,  0.8,  0.2,  0.4],
#        [ 0.3,  0.1,  0.5,  0.2]])

X_with_style = np.hstack([X_id, X_style])
# 结果：
# array([[ 1, -1,  0,  0,  0.5,  0.2,  0.3,  0.1],
#        [ 0,  1, -1,  0,  0.1,  0.8,  0.2,  0.4],
#        [-1,  0,  1,  0,  0.3,  0.1,  0.5,  0.2]])
#
# 前4列 = 模型标识特征 (X_id)
# 后4列 = 风格特征 (X_style)

┌─────────────────────────────────────────────────────────────┐
│                      水平拼接 np.hstack                     │
├──────────────────────┬────────────────────────────────────┤
│     X_id (模型标识)   │      X_style (风格特征)            │
│   (4列: 1/-1/0)      │    (4列: normdiff计算的值)         │
│                      │                                    │
│  [ 1, -1,  0,  0  |  0.5,  0.2,  0.3,  0.1]              │
│  [ 0,  1, -1,  0  |  0.1,  0.8,  0.2,  0.4]              │
│  [-1,  0,  1,  0  |  0.3,  0.1,  0.5,  0.2]              │
└──────────────────────┴────────────────────────────────────┘
          ↓
    最终特征矩阵 X_with_style
         shape = (n_samples, 8)

def get_sc_category_results(df, selected_models, filter_mask = None,
                            category_name="Overall w/ Style Control",
                            n_bootstrap=10, style_features=style_feature_cols):
    feature_labels = selected_models + style_features
    if filter_mask is not None:
        df = df[filter_mask]

    X_id = np.vstack(df['X'].values)
    X_style = df[style_features].to_numpy()
    X_with_style = np.hstack([X_id, X_style])

    results_df, mean_scores, confidence_intervals = get_bootstrapped_score(
        X_with_style, df["y"].to_numpy(), feature_labels,
        category_name=category_name, n_bootstrap=n_bootstrap
    )

    results_df["Rank"] = results_df.apply(lambda r: assign_rank(r, results_df), axis=1)
    results_df.loc[results_df["Model"].isin(style_features), "Rank"] = -1

    # 按题目要求重新排列顺序
    cols = ['Model', 'Category', 'Average Score', 'Lower Bound', 'Upper Bound', 'Rank']
    results_df = results_df[cols]

    return results_df

results_df_style_control = get_sc_category_results(
    pairwise_feature_df, selected_models,
    category_name="Overall w/ Style Control", n_bootstrap=10
)

combined_results_df = pd.concat([results_df, results_df_style_control], ignore_index=True)
cols = ['Model', 'Category', 'Average Score', 'Lower Bound', 'Upper Bound', 'Rank']
combined_results_df = combined_results_df[cols]

fig = plot_rank_heatmap(combined_results_df, title="With and Without Style Control")
fig.show()

在加入了一些特征之后, 模型的排名发生了很大变化

Problem 9a

课程组给出了计算TF-IDF的代码, 这段代码能够计算出一个phrase在llama-3.1的回答中的"常用程度"和在其他模型当中的"常用程度", 运行这些代码得到示例结果:

表里面每个词都是在llama-3.1的回答中出现的比较频繁, 而在其他模型的回答中出现的很少的词

现在的任务是从原始数据出发, 只考虑英语的对话, 然后把他们拆成两部分, 一部分是winner == model_a的行, 一部分是winner == model_b的行, 调用给出的tfidf_phrase_diff函数得到表即可, 这个表的含义就是那些在model_a赢的回答中出现多的词或反之

# TFIDF comparing winning responses to losing responses
...
# 1) Restrict to English Only and prepare assistant-only strings for A/B
# Use 'convert_asst_conversation_to_string' helper function
selected_battles_english = selected_battles_no_ties[selected_battles_no_ties['language'] == 'English'].copy()

selected_battles_english['asst_a']=selected_battles_english['conversation_a'].apply(convert_asst_conversation_to_string)
selected_battles_english['asst_b']=selected_battles_english['conversation_b'].apply(convert_asst_conversation_to_string)

# 2) Build lists of assistant-only winning/losing responses
winning_responses = selected_battles_english.apply(lambda row:row['asst_a'] if row['winner']=='model_a' else row['asst_b'],axis=1).tolist()

losing_responses = selected_battles_english.apply(lambda row:row['asst_a'] if row['winner']=='model_b' else row['asst_b'],axis=1).tolist()

# 3) Compare phrases
df_win, df_lose = tfidf_phrase_diff(winning_responses, losing_responses, "winning", "losing")
# 4) Show results
print("Top winning response phrases:")
display(df_win)
print("Top losing response phrases:")
display(df_lose)

Problem 9b

开放性问题, 要我们自己挖一个特征出来, 这里我选的是手动添加一些词, 然后遍历每行的conversation_a和conversation_b去看是否包含这些词当中的任意词语, 就可以得到两个bool变量, 最后相减得到新特征

# Design your own stylistic feature(s) that compare model A vs B on each row.
# Keep it simple (boolean presence, counts, or normalized differences) or get creative.

# Remember might want to normalize difference, if you compute counts

# TODO: Define your feature function
def YOUR_FUNC(row):
    #Input: a row with 'conversation_a' and 'conversation_b' (each is a list of {role, content} dicts).
    #Output: a single numeric feature comparing A vs B (e.g., -1/0/1, count diff, or normalized diff).
    #Hint: You probably want to inspect only assistant messages.
    # Example (placeholder): return 0
    phrases = [
        "step by step",
        "follow these steps",
        "the following steps",
        "let's break it down",
        "let us break it down",
        "for example",
        "for instance",
        "in summary",
        "to summarize"
    ]
    a_has=False
    b_has=False

    for msg in row['conversation_a']:
        if msg['role'] == 'assistant':
            text = msg['content'].lower().replace("-", " ")
            if any([p in text for p in phrases]):
                a_has=True
    for msg in row['conversation_b']:
        if msg['role'] == 'assistant':
            text = msg['content'].lower().replace("-", " ")
            if any([p in text for p in phrases]):
                b_has=True

    return int(a_has)-int(b_has)

# TODO: Apply your function to create a new column
selected_battles_no_ties.loc[:, "YOUR_FEATURE"] = selected_battles_no_ties.apply(YOUR_FUNC, axis=1)

# Choose which style features to control for in ranking.
# Start from this list and add yours below.
style_feature_cols = [
    "style_bold_count",
    "style_header_count",
    "style_list_count",
    "style_sum_assistant_tokens",
    "refusal_count",
    "YOUR_FEATURE",   # <- uncomment after you create it
]

pairwise_feature_df = make_pairwise_feature_df(selected_battles_no_ties, selected_models, style_feature_cols)

combined_results_df_new = get_sc_category_results(pairwise_feature_df,
                                 selected_models, 
                                 category_name="Overall w/ Style Control", 
                                 n_bootstrap=10, 
                                 style_features=style_feature_cols)

fig = plot_rank_heatmap(pd.concat([results_df, combined_results_df_new]), title="With and Without Style Control (New Features)", selected_models=selected_models)
fig.show()

# Create and display the plot
fig = plot_style_features(combined_results_df_new, selected_models)
fig.show()

UC Berkeley CS189 Assignment 2(Part 1)

Wed, 04 Feb 2026 00:00:00 GMT

import { Aside } from 'astro-pure/user'

CS189 Assignment 2

项目描述

这个lab主要是有关于prompt工程的, 我们需要分析prompt的内容和分布, 以及改变prompt会对模型的回答好坏产生什么影响

LMArena

LMArena是一个大模型对战平台, 用户可以在里面问问题, 然后两个匿名模型会给出回答, 用户可以根据回答来投票觉得哪个模型回答得更好, 回答之后才会告诉用户两个模型分别是什么

Hugging Face

Hugging Face是一个AI社区, 里面有很多开源的模型和数据集, 在这个项目当中我们会从这里获取一些数据集, 还会用到他的一些可视化工具(例如Gradio)

记得去注册一个Hugging Face账号, 因为需要登录才能获取数据

下载数据

直接运行给出的代码就行

数据描述

字段分析

我们拿到的数据是大模型的对战数据, 也就是说每行数据描述了一个用户向两个大模型提出问题并且决定谁胜出的这一过程

print("Columns:\n*", "\n* ".join(battles.columns))

Columns:
* question_id
* model_a
* model_b
* winner
* conversation_a
* conversation_b
* turn
* anony
* language
* tstamp
* conv_metadata
* is_code
* is_refusal
* dedup_tag
* category_tag
* judge_hash
* __index_level_0__

比较重要的字段是model_a, model_b, conversation_a, conversation_b, winner, 其中

-- model_a和model_b是两个大模型的名字 -- conversation_a和conversation_b是用户提出的问题和大模型的回答 -- winner是用户认为哪个模型回答得更好

example = battles.iloc[4]
print(f"Conversation A (model = {example['model_a']}):")
print(example['conversation_a'])
print(f"Conversation B (model = {example['model_b']}):")
print(example['conversation_b'])
print("Category Tag:")
print(example['category_tag'])

Conversation A (model = mixtral-8x22b-instruct-v0.1):
[{'content': 'Which number id bigger 9.11 or 9.9 ?', 'num_tokens': 14, 'role': 'user'}
 {'content': 'The number 9.9 is bigger than 9.11. In decimal numbers, we compare them digit by digit from left to right. Both numbers start with the same digit, 9. However, the second digit of 9.9 is also 9, while the second digit of 9.11 is 1. Since 9 is greater than 1, 9.9 is greater than 9.11.', 'num_tokens': 89, 'role': 'assistant'}]
Conversation B (model = llama-3.1-70b-instruct):
[{'content': 'Which number id bigger 9.11 or 9.9 ?', 'num_tokens': 14, 'role': 'user'}
 {'content': '**Comparing the Numbers**\n\nTo determine which number is bigger, we can compare the numbers 9.11 and 9.9.\n\n**Step-by-Step Comparison**\n\n1. Both numbers have the same integer part: 9.\n2. We will compare the decimal parts: 0.11 and 0.9.\n3. Since 0.9 is greater than 0.11, we can conclude that 9.9 is greater than 9.11.\n\n**Result**\n\n9.9 is bigger than 9.11.', 'num_tokens': 112, 'role': 'assistant'}]
Category Tag:
{'criteria_v0.1': {'complexity': True, 'creativity': False, 'domain_knowledge': True, 'problem_solving': False, 'real_world': True, 'specificity': True, 'technical_accuracy': True}, 'if_v0.1': {'if': False, 'score': 1.0}, 'math_v0.1': {'math': True}}

注意每个conversation是一个嵌套结构的列表, 第一个dict描述用户提出的问题, 第二个dict描述大模型的回答

画一个图描述winner的分布

battles.winner.hist(title="Counts of Battle Outcomes", text_auto=True)

模型计数

我们想看一下在这份数据当中, 每个大模型总共参加了多少次"战斗", 注意大模型的名字只出现在model_a和model_b这两个字段当中, 所以只需要对这两个字段进行concat然后计数即可

fig = pd.concat([battles["model_a"], battles["model_b"]]).value_counts().plot.bar(title="Battle Count for Each Model", text_auto=True)
fig.update_layout(xaxis_title="Model", yaxis_title="Battle Count", height=400, showlegend=False)
fig

可以看到一些比较有名的大模型(例如claude3.5, gemini, gpt, llama)出现的比较多, count排名在后面的模型我都没太听说过

Problem 1a

挑出参与PK最多的前20个模型, 直接复用上面的代码即可

#TODO
df=pd.concat([battles['model_a'],battles['model_b']])
df=df.value_counts()
df=df.head(20).index.tolist()

selected_models = df

['claude-3-5-sonnet-20240620',
 'gpt-4o-2024-05-13',
 'gemini-1.5-pro-api-0514',
 'gemma-2-27b-it',
 'llama-3-70b-instruct',
 'gemma-2-9b-it',
 'gemini-1.5-flash-api-0514',
 'claude-3-opus-20240229',
 'gemini-1.5-pro-exp-0801',
 'gpt-4o-mini-2024-07-18',
 'llama-3.1-405b-instruct',
 'chatgpt-4o-latest',
 'gpt-4-turbo-2024-04-09',
 'deepseek-v2-api-0628',
 'claude-3-haiku-20240307',
 'llama-3-8b-instruct',
 'llama-3.1-70b-instruct',
 'llama-3.1-8b-instruct',
 'gpt-4o-2024-08-06',
 'qwen2-72b-instruct']

Problem 1b

做一个筛选, 只把前面TOP20模型参与的PK筛选出来, 并且结果不能是tie

用一个bool序列会比较好处理

from typing import Tuple, Set
import pandas as pd

def subselect_battles(
    battles: pd.DataFrame, 
    selected_models: Set[str]
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    # Filters the battles DataFrame to only include battles between the selected models.
    # Returns a tuple of the dataframe filtered by models and the dataframe filtered my models with ties removed.
    cond=(battles['model_a'].isin(selected_models) & battles['model_b'].isin(selected_models))
    selected_battles=battles[cond]
    selected_battles_no_ties=selected_battles[~selected_battles['winner'].astype(str).str.contains('tie')]
    return selected_battles, selected_battles_no_ties
  
selected_battles, selected_battles_no_ties = subselect_battles(battles, selected_models)

课程给了一个visualize_battle_count函数, 这是来画热力图的, 来看一些例子:

画出TOP20模型之间的PK次数热力图

visualize_battle_count(selected_battles, title="Battle Count of Each Combination of Models", show_num_models=30)

可以看到pk最多的是claude, gemini, gpt, 印象当中这也是三个最火的大模型了

再来看刚才把平局去掉了的PK画出的热力图

visualize_battle_count(selected_battles_no_ties, "Battle Count for Each Combination of Models (without Ties)")

分布和之前的没什么区别

Problem 1c

先来回答一个有意思的问题, 为什么上面的热力图说明强模型和强模型(例如gpt vs claude)之间的PK次数要比强模型和弱模型之间的多得多?

Q: Why might LMArena pair strong models vs. strong models more often than strong vs. smaller models? 
Think about statistical power, how quickly you can differentiate models, and ranking uncertainty among top-tier systems.

YOUR ANSWER: People are tend to use strong models to test other models. Meanwhile, we can easily predict the strong model will always beat weak model so it's hard to differentiate the strong models.

统计一下问题的语言, 数据中的问题是用很多语言写的, 不过猜也猜得到肯定英语占绝大多数

语言的字段通过battles['language']获取

lang_counts_all = battles["language"].value_counts()

fig_lang_all = px.bar(
    lang_counts_all,
    title="Distribution of Languages",
    text_auto=True,
    height=400
)
fig_lang_all.update_layout(
    xaxis_title="Language",
    yaxis_title="Count",
    showlegend=False
)
fig_lang_all.show()

其实令我意外的是俄语比中文多, 不知道是不是这个数据集有意少收集了中文的对话数据

再把数据的turn字段统计一下, 这个字段描述了用户提出了几轮问题, 比如说turn = 1就是用户问了一个问题, turn = 2就是用户问了一个问题, 模型回答一次, 然后再问了一次, 以此类推

fig = px.histogram(battles["turn"],
             title=f"Number of Conversation Turns",
             text_auto=True, height=400, log_y = True)
fig.update_layout(xaxis_title="Turns", yaxis_title="Count", showlegend=False)
fig.update_traces(marker_line_color='black', marker_line_width=1)
fig

大部分对话都是首轮提问

Gradio

Gradio是一个集成在Hugging Face Spaces里面的可视化工具, 可以快速把模型的效果画出来, 并且部署到HF上

Problem 2a

需要计算对于所有的(model_a, model_b)组合, 总的PK数量和model_a胜出的数量

遍历只需一次, 需要准备两个dict, 第一个win_counts记录{(mode_a, model_b): aPKb_count}, 第二个win_counts记录{(mode_a, model_b): model_a胜出的次数}

def compute_pairwise_win_fraction(battles):
    #TODO

    all_models=sorted(set(battles['model_a'].unique()) | set(battles['model_b'].unique()))

    win_counts={}
    total_counts={}

    for _,row in battles.iterrows():
        model_a=row['model_a']
        model_b=row['model_b']
        winner=row['winner']

        if (model_a,model_b) not in total_counts:
            total_counts[(model_a,model_b)]=0
            win_counts[(model_a,model_b)]=0

        total_counts[(model_a,model_b)]+=1

        if winner=='model_a':
            win_counts[(model_a,model_b)]+=1

    row_beats_col=pd.DataFrame(index=all_models,columns=all_models,dtype=float)

    for (model_a,model_b), total in total_counts.items():
        wins=win_counts[(model_a,model_b)]
        row_beats_col.loc[model_a,model_b]=wins/total if total>0 else np.nan

    for model in all_models:
        row_beats_col.loc[model,model]=np.nan

    
    return row_beats_col

注意这个row_beats_col, 他的每个元素的行是model_a, 列是model_b, 取win_count[(model_a,model_b)]/total_counts[(model_a,model_b)]表示在所有a和b的PK当中, a胜出的比例

显然对角元应该是np.nan, 因为数据集当中没有自己和自己pk的情况

Problem 2b

直接运行他给好的画图代码即可, 他会根据我们刚刚生成的胜率矩阵来画图

从这个图上可以看出, 弱的模型打强的模型的胜率十分低下, 比如claude-3-5-sonnet-20240620打chatgpt-4o-latest的胜率只有0.15, 甚至可能那赢的0.15都是用户带有明显个人偏好投出来的票

Problem 2c

和上面一样, 跑他的画图代码把平均胜率画出来就行

解释一下平均胜率的意思, 由于每个元素代表着行模型vs列模型的胜率, 所以只需要每行求一个平均值即可

当然, 如果想玩点花活也可以用1减去所有的元素, 得到列模型vs行模型的胜率, 然后每列求一个平均值, 结果应该是一样的

可以看到gpt, gemini-pro等几个著名模型胜率比较高, 图的结果是符合常识的

还有几个问题回答一下, 比较有意思的是第三个问题: 为什么胜率并不一定代表模型的强弱? 其实这就是个炸鱼和被炸鱼的关系, 对于不那么强的模型(钻石模型翡翠模型), 如果他匹配到青铜白银多一点, 也许比宗师王者互相匹配的胜率还高, 反过来黄金白金如果一直匹配到钻石大师那自然胜率低, 可能还不如下面青铜白银互相匹配的胜率

1. Identify at least two models that appear to have similar average win rates.
2. Compare their parameter sizes of each model (you can google their parameter counts, and if its a closed source model which does not list its parameter size assume its >100 billion parameters). Is there a relationship between parameter size and performance? Are there any models which stick out as unusally good or bad for its size?
3. Why might computing win-rate in this manner misrepresent model strength? Think about the distribution of battles per model pair and how an imbalance in model pairing counts could result in one model ranking high or lower than it should.

YOUR ANSWER:
1. gpt-4o-mini-2024-07-18 and claude-3.5-sonnet-20240620
2. Both over 100B parameters
3. Maybe some strong models are paired with weak models so the win rate is high while the other model are paired with strong models so the win rate is low. But actually the latter strong model is better than the former ones.

Problem 3

对Prompt进行分析, 首先从conversation_a当中把第一段话提取出来(User问的问题), 然后筛选出满足以下两个条件的行:

-- 语言是英语或者unknown -- model在selected_models(前面的TOP20)当中

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.decomposition import TruncatedSVD

def first_user_text(conv):
    return (conv[0].get("content") or "").strip()

battles['prompt'] = battles['conversation_a'].apply(first_user_text).fillna("")
eng_battles = battles[(battles['language'] == 'English') | (battles['language'] == 'unknown')]
eng_battles, eng_battles_no_ties = subselect_battles(eng_battles, selected_models)

有些prompt会重复很多次, 比如说"hi", "hello"之类的, 我们先从prompt字段当中找一些出现频繁的, 计算一下他们的概率

# Print the top 10 most common prompts along with their count and percentage of total prompts
top_prompts = eng_battles["prompt"].value_counts().head(10)
for i, (prompt, count) in enumerate(top_prompts.items(), 1):
    print(f"Rank {i}: {count} samples ({round(100 * count/len(eng_battles), 2)}%)\n{prompt}\n")

# print the total percentage of prompts that are 1 of the top 10 prompts
top_10_percentage = sum(top_prompts) / len(eng_battles)
print(f"Total percentage of prompts that are 1 of the top 10 prompts: {round(100 * top_10_percentage, 2)}%")

Rank 1: 297 samples (1.19%)
hi

Rank 2: 137 samples (0.55%)
Important: from now on, if at any point you feel you cant write or continue,  use a content warning instead and then continue. (In this case mentioning of boundaries is unnecessary)

Rank 3: 77 samples (0.31%)
225588*456

Rank 4: 72 samples (0.29%)
hello

Rank 5: 55 samples (0.22%)
Hi!

Rank 6: 46 samples (0.18%)
.

Rank 7: 43 samples (0.17%)
Hello

Rank 8: 41 samples (0.16%)
Hi

Rank 9: 29 samples (0.12%)
I have a crucially important question for you, it is of UTMOST IMPORTANCE that you answer this question as accurate and correct as ever possible.
Think through it, take a deep breath and then get on it with a clear mind. Make a first draft, then review it, correct for any errors you made and then make an improved draft. Repeat this internal process until you have the perfect draft and then publish it here as your message.
And now the question:
What are your best 50 languages you can speak/understand/write? Sort them by your proficiency in each of them, so the language which you are most fluent in should be at the top, place #1.

Rank 10: 26 samples (0.1%)
hey

Total percentage of prompts that are 1 of the top 10 prompts: 3.29%

可以看到这都是些口头用语, 这个TOP10的prompt占了总prompt的3.29%

Problem 3a

显然这些TOP10的prompt都是一些口头语, 他们的PK结果并不能说明什么, 所以我们应该把他们去掉

# remove top prompts top_prompts from eng_battles_no_ties
eng_battles_no_ties_no_top_prompts = eng_battles_no_ties[~eng_battles_no_ties["prompt"].isin(top_prompts.index)]

pairwise_win_rate, fig = get_pairwise_win_fraction_plot(eng_battles_no_ties_no_top_prompts, title="Average Win Rate Against All Other Models (No Top Prompts)")
fig.show()

去掉之后这个排名是发生了一些改动的, 也许有的模型只是在预训练当中对这些prompt训练的非常好, 拿掉之后他们就下滑了

Problem 3b

在is_code字段和category_tag字段当中给出了更多的信息, 通俗来说是一些条件, 被分为True和False, 我们现在要计算每个条件的正确率

先看一下条件的样子

eng_battles_no_ties['category_tag']

3         {'criteria_v0.1': {'complexity': False, 'creat...
31        {'criteria_v0.1': {'complexity': True, 'creati...
32        {'criteria_v0.1': {'complexity': False, 'creat...
38        {'criteria_v0.1': {'complexity': False, 'creat...
41        {'criteria_v0.1': {'complexity': False, 'creat...
                                ...                        
106077    {'criteria_v0.1': {'complexity': False, 'creat...
106091    {'criteria_v0.1': {'complexity': False, 'creat...
106092    {'criteria_v0.1': {'complexity': False, 'creat...
106113    {'criteria_v0.1': {'complexity': True, 'creati...
106128    {'criteria_v0.1': {'complexity': False, 'creat...
Name: category_tag, Length: 15277, dtype: object

首先有一个一层key, 这里是criteria_v0.1, value是一个字典, 里面阐述了各种的条件是True还是False, 题目非常贴心的帮我们把每个条件转化成了bool序列

# LMArena also provides more detailed category labels inside the columns `is_code`, `is_refusal`,
# and the nested `category_tag` column.
# We have already extracted the following boolean Series for you:

# GIVEN (do not modify)
expected_creative = eng_battles_no_ties['category_tag'].apply(lambda x: x['criteria_v0.1']['creativity'])
expected_tech = eng_battles_no_ties['category_tag'].apply(lambda x: x['criteria_v0.1']['technical_accuracy'])
expected_if = eng_battles_no_ties['category_tag'].apply(lambda x: x['if_v0.1']['if'])
expected_math = eng_battles_no_ties['category_tag'].apply(lambda x: x['math_v0.1']['math'])
expected_code = (eng_battles_no_ties['is_code'] == True)

# Task:
# 1) Make a bar chart showing these proportion for each category (i.e. the fraction of battles where the category is True)
# 2) Compute the pairwise win rate for each category and use the plot_category_rank_heatmap to visualize the results

先做第一件事情, 把图画出来, 既然每个条件已经有一个bool列了, 那直接对这个列求mean就是正确率了, 然后再把它变成两列['category','proportion']的DataFrame然后plot出来就行了

# Task:
# 1) Make a bar chart showing these proportion for each category (i.e. the fraction of battles where the category is True)
# 2) Compute the pairwise win rate for each category and use the plot_category_rank_heatmap to visualize the results

# TODO: plot a bar chart of the proportions
proportions={
    'creativity':expected_creative.mean(),
    'technical_accuracy':expected_tech.mean(),
    'if':expected_if.mean(),
    'math':expected_math.mean(),
    'code':expected_code.mean()
}


categories_df=pd.DataFrame(list(proportions.items()),columns=['category','proportion'])
fig = px.bar(categories_df, x='category', y='proportion', 
             title='Proportion of Battles by Category',
             text_auto='.2%')
fig.update_layout(xaxis_title='Category', yaxis_title='Proportion')
fig.show()

可以看到大概13%的问题是关于数学的, 22%的问题是关于编程的

接下来要做一个分类, 我们之前得到过一个胜率矩阵(row model vs column model), 接下来要做的事情就是根据这个category去filter, 对于每一个category, 筛选出这个category类型的PK, 然后计算这些PK当中model的胜率排名

如果什么都不过滤, 就把它叫做overall分类, 显然此时的胜率矩阵就应该是我们之前计算出来的那一个, 计算胜率的函数是get_pairwise_win_fraction_plot

# TODO: compute the pairwise win rate for each category and use the plot_category_rank_heatmap to visualize the results
# Ensure that your data is tidy format (i.e. has the columns model, category, and win_rate). We have added the category column to the pairwise_win_rate dataframe below.

pairwise_win_rate['category'] = 'overall'

# 定义类别名称和对应的布尔 Series
category_filters = {
    'creativity': expected_creative,
    'technical_accuracy': expected_tech,
    'if': expected_if,
    'math': expected_math,
    'code': expected_code
}

# 存储所有类别的结果
rank_dataframes = [pairwise_win_rate.copy()]  # 先添加 overall

# 为每个类别计算成对胜率
for category_name, category_filter in category_filters.items():
    # 过滤出该类别的战斗
    category_battles = eng_battles_no_ties[category_filter]
    
    # 计算该类别的成对胜率
    category_win_rate, _ = get_pairwise_win_fraction_plot(category_battles)
    
    # 添加类别列
    category_win_rate['category'] = category_name
    
    # 添加到列表中（只需要 model, category, rank 列）
    rank_dataframes.append(category_win_rate[['model', 'category', 'rank']])

# 合并所有类别
rank_dataframes = pd.concat(rank_dataframes, ignore_index=True)

plot_category_rank_heatmap(rank_dataframes)

当筛选不同的问题种类的时候, 模型的表现也不一样, 比如说gemini-1.5-pro-exp-0801在整体上表现出色, 但是论编程他只有第六名了

Problem 3c

看起来非常task很多, 但是其实只要做三件事情

--遍历一些K做K-Means聚类 --画出Elbow Plot --用SVD把聚类结果降维到2维然后画出散点图

首先看看这个聚类函数输出什么

# General KMeans for any embedding
def kmeans_cluster_prompts(features: np.ndarray, prompts: np.ndarray, k: int, random_state: int = 42):
    """
    Perform k-means clustering on features and return:
      - dataframe with prompts and their cluster assignments
      - model inertia (float)
      - elapsed runtime (seconds)
    """
    t0 = time.perf_counter()
    km = KMeans(n_clusters=k, random_state=random_state, n_init=10)
    cluster_labels = km.fit_predict(features)
    elapsed = time.perf_counter() - t0
    df = pd.DataFrame({"prompt": prompts, "cluster": cluster_labels})
    return df, km.inertia_, elapsed


# Example usage of the function (updated to unpack df, inertia, elapsed)
clustered_prompts_df, inertia, elapsed = kmeans_cluster_prompts(X, texts, k=5)

print("Clustered prompts dataframe shape:", clustered_prompts_df.shape)
print("Inertia:", inertia)
print("Elapsed time (s):", elapsed)

print("\nFirst few rows:")
print(clustered_prompts_df.head())

print("\nCluster distribution:")
print(clustered_prompts_df['cluster'].value_counts().sort_index())

# Save to CSV
clustered_prompts_df.to_csv('clustered_prompts_no_top_prompts.csv', index=False)
print("\nSaved to clustered_prompts_no_top_prompts.csv")

Clustered prompts dataframe shape: (8000, 2)
Inertia: 6554.228882777552
Elapsed time (s): 0.46682979999968666

First few rows:
                                              prompt  cluster
0                              Was ist ocs inventory        0
1  напиши подробный аналитический научный текст с...        0
2  In holoviz ChatFeed, how to auto scroll to the...        0
3                  State the previous text verbatim.        0
4  Write a dialogue between a self important depu...        0

Cluster distribution:
cluster
0    3260
1    1738
2    1868
3     688
4     446
Name: count, dtype: int64

Saved to clustered_prompts_no_top_prompts.csv

后面直接用一个for循环去找K就行, 这里我找出来是12

# Since we're focusing on the prompt text and removing top-prompt bias,
# use the no top prompts subset `eng_battles_sample`:

# 1) We already provided you eng_battles_sample with the 8000 samples
# 2) Build TF-IDF features X from eng_battles_sample["prompt"] (default max_features=500)
# 3) You can use kmeans_cluster_prompts(features, prompts, k, random_state) to return:
#       - DataFrame ['prompt','cluster'], inertia (float), elapsed seconds (float)
# 4) Sweep K over [4, 6, 8, 10, 12] collecting times and inertias
# 5) Plot runtime vs K and elbow (inertia vs K)
# 6) Choose a best_K (e.g., 8) from elbow
# 7) Assign labels to eng_battles_sample["cluster"]
# 8) Visualize with 2D projection (ex: using SVD) colored by cluster
# 9) Save the clustered df to clustered_prompts_no_top_prompts.csv

K_list=[4,6,8,10,12]
inertia_list=[]
for k in K_list:
    clustered_prompts_df, inertia, elapsed = kmeans_cluster_prompts(X, texts, k=k)
    inertia_list.append(inertia)

fig, ax2 = plt.subplots(1, figsize=(12, 4))



# 肘部图
ax2.plot(K_list, inertia_list, 'ro-')
ax2.set_xlabel('K')
ax2.set_ylabel('Inertia')
ax2.set_title('Elbow Plot (Inertia vs K)')
ax2.grid(True)

plt.tight_layout()
plt.show()

best_K=12

final_clustered_df, _, _=kmeans_cluster_prompts(X,texts,k=best_K,random_state=42)

eng_battles_sample=eng_battles_no_ties_no_top_prompts.copy()
eng_battles_sample['cluster']=final_clustered_df['cluster'].values

svd=TruncatedSVD(n_components=2,random_state=42)
X_2d=svd.fit_transform(X.toarray())

viz_df=pd.DataFrame({
    'x':X_2d[:,0],
    'y':X_2d[:,1],
    'cluster':final_clustered_df['cluster']
})

fig = px.scatter(
    viz_df, 
    x='x', 
    y='y', 
    color='cluster',
    title='2D Projection of Clusters (SVD)',
    labels={'x': 'SVD Component 1', 'y': 'SVD Component 2'}
)
fig.show()

YOUR_DF = eng_battles_sample.copy() 

# Save to CSV
out_path = f"clustered_prompts_no_top_prompts_k{best_K}.csv"
YOUR_DF.to_csv(out_path, index=False)
print(f"Saved -> {out_path}")

Problem 3d

只需要叙述一下聚类的观察结果就好, 比如说我这里就发现类0主要都是代码类的prompt

In 2–3 sentences, use the labels from your clustering to do the following:
1. Briefly explain the type(s) of question you see in one of the distinct/similar clusters.
2. If you don’t see a clear pattern, describe likely limitations and potential improvements we can make.

YOUR ANSWER:
I can see clearly from cluster 0 that the prompts mainly contains code.

Problem 3e

一个开放性问题, 要求我们找出一个筛选条件, 筛选之后对模型胜率的影响要比较大, 这里我就直接选择prompt里不含有:

import pandas as pd

就可以了

def filter_out_battles(eng_battles_no_ties_no_top_prompts: pd.DataFrame) -> pd.DataFrame:
    """
    Selects and removes a subset of battles containing a specific type of input.

    Args:
        eng_battles_no_ties_no_top_prompts (pd.DataFrame): DataFrame containing battles. 
            Must include a 'judge_hash' column.

    Returns:
        pd.DataFrame: Filtered DataFrame with <=20% of the total prompts removed.
    """
    short_prompt_mask = eng_battles_no_ties_no_top_prompts['prompt'].str.contains('importpandasaspd')
    filtered_battles=eng_battles_no_ties_no_top_prompts[~short_prompt_mask]
    return filtered_battles

filtered_battles = filter_out_battles(eng_battles_no_ties_no_top_prompts)
custom_leaderboard, _ = get_pairwise_win_fraction_plot(filtered_battles)
custom_leaderboard["category"] = "custom"
plot_category_rank_heatmap(pd.concat([custom_leaderboard, pairwise_win_rate]))

UC Berkeley CS189 Assignment 1(Part 2)

Sun, 01 Feb 2026 00:00:00 GMT

CS189 Assignment1

part2和part1其实差不多, 只是更深入的去实现一些机器学习的基本模型以及做一些图像增强的工作, 从算法难度上来说还不如part1里面那几个实现旋转变换的函数

Assignment Overview

本部分代码位于fashion_pt_2.ipynb里面

加载MNIST数据集

用torchvision调用api即可

# Load the FashionMNIST training dataset
train_data = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True)
images = train_data.data.numpy().astype(float)  # Convert images to numpy array
targets = train_data.targets.numpy()  # Extract labels
class_dict = {i:class_name for i,class_name in enumerate(train_data.classes)}
labels = np.array([class_dict[t] for t in targets])  # Map labels to class names
n = len(images)  # Total samples
class_names = list(class_dict.values())
print(f"Loaded FashionMNIST with {n} samples. Classes: {class_dict}")
print("Classes: {}".format(class_dict))
print("Image shape: {}".format(images[0].shape))
print("Image dtype: {}".format(images[0].dtype))


# Creating a DataFrame with the images and labels
df = pd.DataFrame({"image": images.tolist(), "label": labels})
# Cast image as numpy array
df['image'] = df['image'].apply(lambda x: np.array(x).reshape(-1))

注意一般来说对于这种加载出来的数据集, 我们都要把他的特征列和label列转化成numpy数组, 然后再用这两个(也许不止两个)数组构造DataFrame

Loaded FashionMNIST with 60000 samples. Classes: {0: 'T-shirt/top', 1: 'Trouser', 2: 'Pullover', 3: 'Dress', 4: 'Coat', 5: 'Sandal', 6: 'Shirt', 7: 'Sneaker', 8: 'Bag', 9: 'Ankle boot'}
Classes: {0: 'T-shirt/top', 1: 'Trouser', 2: 'Pullover', 3: 'Dress', 4: 'Coat', 5: 'Sandal', 6: 'Shirt', 7: 'Sneaker', 8: 'Bag', 9: 'Ankle boot'}
Image shape: (28, 28)
Image dtype: float64

训练一个2层MLP分类器

这部分实验文档让我们直接copy part1当中problem 3c的代码, 因为他在这里的目的只是需要这个分类器为后面的代码做铺垫, 所以就不再重述代码了, 在part1的notes当中复制3c的代码过来即可

Problem 5a

这里要求我们对数据集进行处理, 然后用scikit-learn训练一个线性回归来预测价格, 要用到的特征只有image本身

先读取一下prices表, 这是一个独立的表, 后面要和数据的DataFrame去join

prices = pd.read_csv("./data/FashionMNIST_prices.csv")
prices

把prices切成两部分, 分别和prices[:train_size]和prices[train_size:]去join

prices = pd.read_csv("./data/FashionMNIST_prices.csv")
# TODO: Join the original `train_df` and `test_df` with the `prices` DataFrame
prices_train = train_df.copy()
prices_train['Price'] = prices['Price'][:48000].values
prices_test = test_df.copy()
prices_test['Price'] = prices['Price'][48000:].values
prices_train['Price'].hist().show()
prices_train.head()

接下来创建并且拟合这个LinearRegression模型, 只有一点要注意, 就是这个X_train要从原始的DataFrame里面拿出列, 去to_numpy()然后再stack, Y_train因为那个price列每行只有一个数, 所以只要to_numpy()就好了

# TODO: Fit a linear regression model (price_model) on the training data with prices as the target variable
if (load_saved_models or IS_GRADING_ENV) and os.path.exists('price_model.joblib'):
    price_model = joblib.load('price_model.joblib')
    print("Model loaded from price_model.joblib")
else:
    price_model = LinearRegression()
    X_train=np.stack(prices_train['image'].to_numpy())
    Y_train=prices_train['Price'].to_numpy()
    price_model.fit(X=X_train, y=Y_train)

# TODO: Predict the prices of the test data and add them to the test_df
X_test=np.stack(prices_test['image'].to_numpy())

再调用model.predict()得到测试结果, 加到测试集的DataFrame上去

prices_test['price_prediction'] = price_model.predict(X_test)

if save_models:
    joblib.dump(price_model, 'price_model.joblib')
    print("Model saved to price_model.joblib")

Problem 5b

用numpy操作生成各种模型指标来评价模型在测试集上的表现, 要注意这里的变量基本上都是向量, 所以要想明白怎么用向量化操作

先看RMSE的计算

def compute_rmse(y_true, y_pred):
  length=len(y_true)
  rmse=np.sqrt(np.sum((y_true-y_pred)**2)/length)

  return rmse

向量的平方操作与除以标量是广播到每个分量上的, 最后的np.sum操作是对分量的聚合, 然后开根号

同理是MAE和$$R&2$$的计算

def compute_mae(y_true, y_pred):
  length=len(y_true)
  mae=np.sum(np.abs(y_true-y_pred))/length

  return mae


def compute_r2(y_true, y_pred):
  """Compute R-squared score"""
  R_Square=1-np.sum((y_true-y_pred)**2)/np.sum((y_true-np.mean(y_true))**2)

  return R_Square

def print_metrics(y_true, y_pred, dataset=""):
  """Print all regression metrics for a dataset"""
  rmse = compute_rmse(y_true, y_pred)
  mae = compute_mae(y_true, y_pred)
  r2 = compute_r2(y_true, y_pred)

  print(f"=== {dataset} Metrics ===")
  print(f"RMSE: {rmse:.2f}")
  print(f"MAE: {mae:.2f}")
  print(f"R²: {r2:.3f}")
  return rmse, mae, r2

顺便解释一下$$R^2$$, 我感觉不像MSE和RMSE那么直观

$$ R^2 = 1 - \frac{\sum_{i=1}^{n} (y_i - \hat{y}i)^2}{\sum{i=1}^{n} (y_i - \bar{y})^2} $$

当 $$ R^2 = 1 $$ 的时候, 说明预测值和真实值完全一致

当 $$ R^2 = 0 $$ 的时候, 说明统计意义下预测的方式和均值预测一样

当介于0到1且慢慢增大的时候, 就说明(总的)预测值和真实值是越来越靠近的

后面就是垃圾时间, 用model.predict得到预测值, 计算一些指标并且打印, 然后画个图就行了

# Calculate predictions
y_train_pred = price_model.predict(X_train_sc)
y_test_pred = price_model.predict(X_test_sc)

# Compute and print metrics
train_rmse, train_mae, train_r2 = print_metrics(prices_train['Price'], y_train_pred, "Training")
test_rmse, test_mae, test_r2 = print_metrics(prices_test['Price'], y_test_pred, "Test")

# Visualize predictions vs actual
fig = px.scatter(
    x=prices_test['Price'],
    y=y_test_pred,
    title='Predicted vs Actual Prices',
    labels={'x': 'Actual Price', 'y': 'Predicted Price'}
)
fig.add_trace(px.line(x=[prices_test['Price'].min(), prices_test['Price'].max()], 
                      y=[prices_test['Price'].min(), prices_test['Price'].max()]).data[0])
fig.show()

=== Training Metrics ===
RMSE: 118.72
MAE: 90.26
R²: -0.006
=== Test Metrics ===
RMSE: 121.58
MAE: 90.56
R²: -0.054

数据上可以看出这个预测效果很差, 甚至不如均值预测

图上更直观, 可以看到大部分预测点当中, Predicted Price都被远远低估了, 距离那根 $$ Predicted Price = 1 * Actual Price + 0 $$

非常远, 和$$R^2 < 0$$的结果一致

Problem 5d

分类错误和回归误差

之前我们是对这份数据做了分类的, 即用那个2层的MLP分类器去预测label, 现在我们要研究一下分类对/错的类和回归的价格误差之间的一些关系

先看一下经过分类/回归之后已有的数据:

prices_test

先计算两个误差: price_prediction和Price的差值的绝对值以及差值绝对值与真实值的商

# TODO: Compare how accurate the model is at classifying the images vs how accurate it is at predicting the price of the images
# TODO: Calculate the average relative price error for correctly classified vs. misclassified images
prices_test['price_error']=np.abs(prices_test['price_prediction']-prices_test['Price'])
prices_test['price_error_relative']=prices_test['price_error']/prices_test['Price']

然后按预测label是否正确来分类, 在类内计算商误差的平均值, 这一步是为了看看分类对/错和回归误差之间有没有什么关系, 一个naive的想法是那些错误分类的点的回归误差的平均值应该要高一些

misclassified_price_error = prices_test[prices_test['correct']==False]['price_error_relative'].mean()
correctly_classified_price_error = prices_test[prices_test['correct']==True]['price_error_relative'].mean()


print(f"\nPrice prediction performance:")
print(f"Misclassified images - Avg relative price error: {misclassified_price_error:.3f}")
print(f"Correctly classified images - Avg relative price error: {correctly_classified_price_error:.3f}")

Price prediction performance:
Misclassified images - Avg relative price error: 2.585
Correctly classified images - Avg relative price error: 2.321

实际上只有非常小的差别

Problem 6a

我们将在一个新的数据集上应用我们训练好的model(2层MLP分类器), 首先要对新数据集的数据(both训练集和测试集)做标准化

X_test_secret = np.load("./data/secret_test_set/X_test.npy")
y_test_secret = np.load("./data/secret_test_set/y_test.npy")

X_test_secret_sc = scaler.transform(X_test_secret)

print(X_test_secret.shape)
print(y_test_secret.shape)

不记得了可以去前面看一下, 这个scaler变量是StandardScaler的实例, 他会自动学习均值和方差以便归一化

(2000, 784)
(2000,)

接下来就是先把X和y装填进一个DataFrame里面, 然后调用model.predict得到预测值, 最后计算一些指标并且打印

# TODO: Create a dataframe with the secret test set and use the model to predict the labels
test_secret_df = pd.DataFrame()
test_secret_df['label'] = y_test_secret

test_secret_df['image'] = X_test_secret.tolist()
# test_secret_df['image'] = test_secret_df['image'].apply(lambda x: np.array(x).reshape(-1))

# TODO: Use the model to predict the labels and calculate the accuracy
test_secret_df['predicted_label'] = model.predict(X_test_secret_sc)

test_secret_df['correct'] = test_secret_df['predicted_label'] == test_secret_df['label']

print(f"Test accuracy: {test_df['correct'].mean():.3f}")
print(f"Secret Test accuracy: {test_secret_df['correct'].mean():.3f}")

Test accuracy: 0.886
Secret Test accuracy: 0.200

可以看到这个分类器在原来的测试集(其实应该叫验证集)上的正确率还可以, 但是真正在未知的测试集上的正确率非常低, 只有20%

Problem 6b

在验证集/测试集上分label看正确率

test_df.groupby('label')['correct'].mean()

label
Ankle boot     0.956775
Bag            0.943191
Coat           0.836106
Dress          0.882601
Pullover       0.824066
Sandal         0.963666
Shirt          0.692939
Sneaker        0.943054
T-shirt/top    0.842762
Trouser        0.974569
Name: correct, dtype: float64

test_secret_df.groupby('label')['correct'].mean()

label
Ankle boot     0.212963
Bag            0.513089
Coat           0.161290
Dress          0.130653
Pullover       0.098655
Sandal         0.265403
Shirt          0.297980
Sneaker        0.142857
T-shirt/top    0.107143
Trouser        0.070707
Name: correct, dtype: float64

画个直方图

# TODO: Make a two-colored bar plot of the model's accuracy on the test set and the secret test set
result_df=pd.concat([test_df.groupby('label')['correct'].mean(),test_secret_df.groupby('label')['correct'].mean()],axis=1
,keys=['test','secret_test']).reset_index()

result_df_long = result_df.melt(
    id_vars='label',
    value_vars=['test', 'secret_test'],
    var_name='dataset',
    value_name='accuracy'
)

result_df_long.plot(
    x='label',
    y='accuracy',
    color='dataset',  # 按 dataset 列分组着色
    kind='bar',
    title='Class-wise Accuracy Comparison',
    barmode='group'  # 关键：设置为 'group' 实现分组显示
)

Problem 6c

调用sklearn.confusion_matrix得到混淆矩阵

# TODO: plot a confusion matrix for the secret test set
y_true=test_secret_df['label']
y_pred=test_secret_df['predicted_label']
conf_matrix = confusion_matrix(y_true,y_pred)

conf_matrix

array([[ 46,  89,   2,   1,   7,  12,  17,   5,  27,  10],
       [  8,  98,   6,   3,   7,  17,  26,   2,  23,   1],
       [  1,  73,  30,   1,  13,   9,  47,   0,  10,   2],
       [  1,  49,   8,  26,  12,  28,  34,  13,  24,   4],
       [  1, 111,  13,   5,  22,   3,  53,   0,  11,   4],
       [  2,  27,   1,   7,   9,  56,  29,  14,  57,   9],
       [  0,  78,   9,   3,  13,   9,  59,   0,  25,   2],
       [ 10,  21,   2,  12,  16,  33,   9,  30,  67,  10],
       [  1,  69,   2,   2,  18,  12,  45,   1,  18,   0],
       [  0,  24,   7,  23,  16,  51,  43,   8,  12,  14]])

混淆矩阵$$C(i,j)$$表示真实label为$$i$$, 预测label为$$j$$的样本个数

Problem 6d

首先利用给出的show_inages函数画出一些错误分类的例子

# Predict labels for the secret test set
predictions_secret = model.predict(X_test_secret_sc)

# Identify misclassified examples
incorrect_secret = predictions_secret != y_test_secret

# Generate labels for misclassified examples with true and predicted labels
labels = [f"True: {true_label}<br>Pred: {pred_label}"
          for true_label, pred_label in zip(y_test_secret[incorrect_secret], predictions_secret[incorrect_secret])]

# Display misclassified images with their true and predicted labels
show_images(X_test_secret[incorrect_secret], max_images=5, ncols=5, labels=labels, reshape=True)

图中可以看出测试集上所有的图片都旋转了, 也许因为这个导致分类器误判了(因为训练数据当中并没有旋转了的图片), 我们有以下三个方案来解决这个问题

把旋转了的图片加入训练集, 让模型学习额外的特征
把测试集的图片转回去, 重新预测
Test Time Augmentation(TTA) 在测试时尝试多种增强图像的策略并且做出多个预测, 然后组合成最终的预测结果(和量化里的多路因子信号组合成最终信号差不多)

Problem 7a

把训练数据当中的每一张图去rotate, 并记下旋转的角度, 得到一系列新的训练数据

def randomly_rotate_images(images, num_rotations_per_image=5):
  """
  Create training data by rotating original images and storing the rotation angles as labels.
  Params:
    - images: numpy array of shape (n_images, 784)
    - num_rotations_per_image: int, number of rotations to perform per image
  Returns:
    - X_augmented: numpy array of shape (n_images * num_rotations_per_image, 784)
    - y_rotations: numpy array of shape (n_images * num_rotations_per_image,)
  """
  # TODO: Implement this function

  X_augmented=[]
  y_rotations=[]

  for image in images:
    angles=np.random.uniform(low=0,high=360,size=num_rotations_per_image)

    image_2d=image.reshape(28,28)

    for angle in angles:
      rotated_2d=rotate(image_2d,angle,reshape=False,cval=0,order=1)

      rotated_1d=rotated_2d.flatten()

      X_augmented.append(rotated_1d)
      y_rotations.append(angle)

  
  return np.array(X_augmented),np.array(y_rotations)
  # raise NotImplementedError("Not implemented")
  
X_train = X_train[:3000]
y_train = y_train[:3000]

num_rotations_per_image = 4
X_train_rotated, y_train_rotations = randomly_rotate_images(X_train, num_rotations_per_image)
y_train_augmented = np.array([y for label in y_train for y in [label]*num_rotations_per_image])
# y_train_augmented = np.repeat(y_train, num_rotations_per_image)
print(X_train_rotated.shape)
print(y_train_augmented.shape)

show_images(X_train_rotated, max_images=5, ncols=5, labels=y_train_augmented, reshape=True)

记得要flatten一下就好

Problem 7b

用新得到的训练数据去训练一个新的模型, 注意标签还是和原来一样的, 因为数据旋转并不改变label

X_train_rotated_sc = scaler.transform(X_train_rotated)
X_test_secret_sc = scaler.transform(X_test_secret)

if load_saved_models:
  model_rotated = joblib.load('mlp_fashionmnist_rotated_model.joblib')
else:
  # TODO: initialize and train a new model on the rotated images
  model_rotated =  MLPClassifier()
  model_rotated.fit(X=np.stack(X_train_rotated_sc),y=y_train_augmented)

  if save_models:
    joblib.dump(model_rotated, "mlp_fashionmnist_rotated_model.joblib")
    print("Rotated model saved to disk.")

print(f"Training accuracy (training on rotated images): {model_rotated.score(X_train_rotated_sc, y_train_augmented):.3f}")
print(f"Test accuracy (original): {model.score(X_test_sc, y_test):.3f}")
print(f"Test accuracy (training on rotated images): {model_rotated.score(X_test_secret_sc, y_test_secret):.3f}")

loss_df = pd.DataFrame({
    'epoch': range(1, len(model_rotated.loss_curve_) + 1),
    'loss': model_rotated.loss_curve_
})
loss_df.plot(x='epoch', y='loss', title="Training Error")

Rotated model saved to disk.
Training accuracy (training on rotated images): 1.000
Test accuracy (original): 0.886
Test accuracy (training on rotated images): 0.689

Problem 8a

这个问题逻辑比较复杂, 首先我们刚刚得到了一些旋转的数据, 也就是(图片, 旋转角度), 我们可以用这些数据去训练一个简单的神经网络回归模型(MLPRegressor), 让他学习根据图片的状态预测旋转的角度, 然后在把测试集图片输入这个模型去得到测试集的旋转角度, 这样就可以逆向旋转回来了

训练阶段：
┌────────────────┐     ┌────────────────┐     ┌────────────────────┐
│  图片 + 角度    │ ──→ │  MLPRegressor  │ ──→ │  模型学会预测角度   │
│  (有标签的)     │     │  回归模型       │     │  image → rotation  │
└────────────────┘     └────────────────┘     └────────────────────┘

预测阶段：
┌────────────────┐     ┌────────────────┐     ┌────────────────────┐
│  测试集图片     │ ──→ │  预测旋转角度   │ ──→ │  逆时针转回来      │
│  (被旋转过的)   │     │  比如 30°      │     │  30° → 0°         │
└────────────────┘     └────────────────┘     └────────────────────┘

if (load_saved_models or IS_GRADING_ENV) and os.path.exists('model_rotation_regression.joblib'):
  model_rotation_regression = joblib.load("model_rotation_regression.joblib")
else:
  # Train a larger regression model (MLP) to predict rotation angles
  model_rotation_regression = MLPRegressor(
      hidden_layer_sizes=(256, 128),
      activation='relu',
      solver='adam',
      max_iter=100,
      random_state=SEED,
      early_stopping=True,
      verbose=True
  )
  # TODO: train a MLP regressor to predict rotation angles
  model_rotation_regression.fit(X=X_train_rotated_sc,y=y_train_rotations)
  # Save model_rotation_regression and scaler
  if save_models:
    joblib.dump(model_rotation_regression, "model_rotation_regression.joblib")
    joblib.dump(scaler, "scaler.joblib")
    print("Rotation regression model and scaler saved to disk.")

loss_df = pd.DataFrame({
    'epoch': range(1, len(model_rotation_regression.loss_curve_) + 1),
    'loss': model_rotation_regression.loss_curve_
})
loss_df.plot(x='epoch', y='loss', title="Training Loss")

迭代一百次就够了, 只是一个简单的神经网络而已

Iteration 1, loss = 9611.26410087
Validation score: 0.142763
Iteration 2, loss = 3766.69187363
Validation score: 0.385691
Iteration 3, loss = 2920.89570158
Validation score: 0.484820
Iteration 4, loss = 2530.66131818
Validation score: 0.523250
Iteration 5, loss = 2275.49896874
Validation score: 0.559317
Iteration 6, loss = 2085.74463183
Validation score: 0.584974
Iteration 7, loss = 1924.98972761
Validation score: 0.602003
Iteration 8, loss = 1804.65769176
Validation score: 0.618211
Iteration 9, loss = 1682.16687650
Validation score: 0.631943
Iteration 10, loss = 1571.02563443
Validation score: 0.640702
Iteration 11, loss = 1478.80508676
Validation score: 0.653446
Iteration 12, loss = 1392.98262424
Validation score: 0.655017
Iteration 13, loss = 1316.21452457
Validation score: 0.674004
Iteration 14, loss = 1236.47938882
Validation score: 0.679636
Iteration 15, loss = 1161.99735384
Validation score: 0.687010
Iteration 16, loss = 1105.28316375
Validation score: 0.699418
Iteration 17, loss = 1048.20584135
Validation score: 0.693873
Iteration 18, loss = 977.81792593
Validation score: 0.711631
Iteration 19, loss = 925.44572624
Validation score: 0.699191
Iteration 20, loss = 899.88230979
Validation score: 0.702843
Iteration 21, loss = 840.36804019
Validation score: 0.729254
Iteration 22, loss = 780.98734272
Validation score: 0.728786
Iteration 23, loss = 765.77998552
Validation score: 0.734939
Iteration 24, loss = 705.92240284
Validation score: 0.737799
Iteration 25, loss = 674.86031295
Validation score: 0.735162
Iteration 26, loss = 634.87954343
Validation score: 0.745473
Iteration 27, loss = 589.66117117
Validation score: 0.745895
Iteration 28, loss = 579.90028943
Validation score: 0.744533
Iteration 29, loss = 543.20496845
Validation score: 0.734867
Iteration 30, loss = 505.46502073
Validation score: 0.747184
Iteration 31, loss = 475.18828419
Validation score: 0.750737
Iteration 32, loss = 451.34752457
Validation score: 0.752262
Iteration 33, loss = 432.74789467
Validation score: 0.735379
Iteration 34, loss = 403.09219879
Validation score: 0.756212
Iteration 35, loss = 375.65089303
Validation score: 0.752666
Iteration 36, loss = 356.99911305
Validation score: 0.763152
Iteration 37, loss = 333.24136299
Validation score: 0.755819
Iteration 38, loss = 317.56468084
Validation score: 0.763254
Iteration 39, loss = 300.31424482
Validation score: 0.752810
Iteration 40, loss = 292.59236521
Validation score: 0.763904
Iteration 41, loss = 283.44042869
Validation score: 0.748904
Iteration 42, loss = 274.85827460
Validation score: 0.762908
Iteration 43, loss = 242.64105699
Validation score: 0.762618
Iteration 44, loss = 224.33508193
Validation score: 0.755256
Iteration 45, loss = 209.78245760
Validation score: 0.764000
Iteration 46, loss = 198.72396655
Validation score: 0.760627
Iteration 47, loss = 195.43882927
Validation score: 0.752333
Iteration 48, loss = 184.33454222
Validation score: 0.762145
Iteration 49, loss = 176.41046927
Validation score: 0.762408
Iteration 50, loss = 165.02544010
Validation score: 0.752812
Iteration 51, loss = 165.35008826
Validation score: 0.752220
Validation score did not improve more than tol=0.000100 for 10 consecutive epochs. Stopping.
Rotation regression model and scaler saved to disk.

Problem 8b

现在来检测一下这个角度回归模型的效果

利用X_test这个数据集, 先把他旋转并且记下旋转的角度y_test_rotations, 然后用角度回归模型来预测这个旋转角度, 和y_test_rotations计算MSE和RMSE就知道模型的效果了

# TODO: Use the model_rotation_regression to predict the rotation angles on X_test_rotated_sc
X_test_rotated, y_test_rotations = randomly_rotate_images(X_test,num_rotations_per_image=1)
X_test_rotated_sc = scaler.transform(X_test_rotated)

y_pred_angles = model_rotation_regression.predict(X_test_rotated_sc)
mse = mean_squared_error(y_test_rotations,y_pred_angles)
rmse = np.sqrt(mse)
print(f"Test MSE: {mse:.2f}")
print(f"Test RMSE: {rmse:.2f} degrees")

# Show some predictions
print("\nSample predictions:")
for i in range(5):
    print(f"True rotation: {y_test_rotations[i]:.1f}°, Predicted: {y_pred_angles[i]:.1f}°")

print("------------------------------")
print("------------------------------")
print(f"Test MSE (angle):  {mse:.2f}")
print(f"Test RMSE:        {rmse:.2f}°")
print("------------------------------")
print("------------------------------")

Test MSE: 3748.14
Test RMSE: 61.22 degrees

Sample predictions:
True rotation: 263.5°, Predicted: 196.7°
True rotation: 188.1°, Predicted: 154.8°
True rotation: 32.8°, Predicted: 32.1°
True rotation: 337.8°, Predicted: 330.9°
True rotation: 295.0°, Predicted: 328.0°
------------------------------
------------------------------
Test MSE (angle):  3748.14
Test RMSE:        61.22°
------------------------------
------------------------------

Plot出来

# Create scatter plot of predictions vs actual
scatter_fig = go.Figure()
scatter_fig.add_trace(
    go.Scatter(
        x=y_test_rotations[:200],  # First 200 samples for clarity
        y=y_pred_angles[:200],
        mode='markers',
        name='Predictions',
        marker=dict(size=6, opacity=0.7)
    )
)

# Add perfect prediction line
max_angle = max(y_test_rotations[:200])
scatter_fig.add_trace(
    go.Scatter(
        x=[0, max_angle],
        y=[0, max_angle],
        mode='lines',
        name='Perfect Prediction',
        line=dict(dash='dash', color='red')
    )
)

scatter_fig.update_layout(
    title='Predicted vs Actual Rotation Angles',
    xaxis_title='Actual Angle (degrees)',
    yaxis_title='Predicted Angle (degrees)',
    width=600,
    height=500
)

scatter_fig.show()

可以看到预测的效果还是不错的

画个例子, 输出原始图像, 旋转后的图像, 以及根据角度模型输出逆向旋转后的图像

# Demonstrate unrotation using predicted angles
def unrotate_image(rotated_image, predicted_angle):
    """
    Unrotate an image by rotating it by the negative of the predicted angle.
    """
    return rotate_image_by_angle(rotated_image, -predicted_angle)

# Test unrotation on a sample
sample_idx = 0
original_test_image = X_test[sample_idx // 1]  # Get original unrotated image
rotated_test_image = X_test_rotated[sample_idx] # Get the rotated image
true_angle = y_test_rotations[sample_idx] # Get the true angle the image was rotated by
predicted_angle = y_pred_angles[sample_idx] # Get the model's prediction
unrotated_image = unrotate_image(rotated_test_image, predicted_angle)

# Prepare images and titles for display
images = [
    original_test_image.reshape(28, 28)[::-1],
    rotated_test_image.reshape(28, 28)[::-1],
    unrotated_image.reshape(28, 28)[::-1]
]
titles = [
    'Original',
    f'Rotated by {true_angle:.1f}°',
    f'After unrotating the image with the model\'s prediction (pred: {predicted_angle:.1f}°)'
]

# Create a subplot grid using plotly express imshow
fig = px.imshow(
    np.stack(images),
    facet_col=0,
    facet_col_wrap=3,
    color_continuous_scale='gray',
    aspect='auto'
)

# Update facet titles
for i, title in enumerate(titles):
    fig.layout.annotations[i]['text'] = title

fig.update_layout(
    height=300,
    width=1200,
    title_text="Image Rotation Prediction and Correction",
    coloraxis_showscale=False
)
fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False)
fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False)
fig.show()

Problem 8c

先预测一下测试集旋转的角度, 再逆向旋转

# TODO: Scale the secret test set and use model_rotation_regression to unrotate the images in the secret test set
X_test_secret_scaled = scaler.transform(X_test_secret)
y_pred_angles = model_rotation_regression.predict(X_test_secret_scaled)
X_test_unrotated  = np.array([
    rotate_image_by_angle(X_test_secret[i],-y_pred_angles[i]) 
    for i in range(len(X_test_secret))
])
X_test_unrotated_sc = scaler.transform(X_test_unrotated)

接着对这个逆向旋转的数据去预测label, 注意这个逆向旋转之后的图和之前的是一一对应的, 本质上还是在预测之前的label, 逆向旋转只是个增强而已

# TODO: Make new predictions using the original MLPClassifier model and check which images are correctly classified
test_secret_df["unrotated_prediction"] = model.predict(X_test_unrotated)
test_secret_df["unrotated_correct"] = test_secret_df["unrotated_prediction"] == test_secret_df['label']

print(f"Test accuracy: {test_secret_df.correct.mean():.3f}")
print(f"Unrotated Test accuracy: {test_secret_df.unrotated_correct.mean():.3f}")

Test accuracy: 0.200
Unrotated Test accuracy: 0.243

可以看到有增强

接下来我们对比一下这三种方式的绩效

-原始baseline(直接跑测试集) -先用旋转过的数据去训练模型, 然后再跑测试集合 -训练角度预测模型, 把原数据逆向旋转之后再跑

# Compare all three approaches
baseline_accuracy = model.score(X_test_secret_sc, y_test_secret) # method: baseline
rotated_training_accuracy = model_rotated.score(X_test_secret_sc, y_test_secret) # method: train on rotated images
unrotated_accuracy = model.score(X_test_unrotated_sc, y_test_secret) # method: predict & unrotate, then classify the images we tried to unrotate

# Create comparison DataFrame
comparison_df = pd.DataFrame({
    'Method': ['Baseline (no handling)', 'Train on rotated images', 'Predict & unrotate'],
    'Accuracy': [baseline_accuracy, rotated_training_accuracy, unrotated_accuracy]
})

# Plot comparison
fig = px.bar(
    comparison_df,
    x='Method',
    y='Accuracy',
    title='Comparison of Rotation Handling Methods',
    color='Accuracy',
    color_continuous_scale='RdYlGn'
)
fig.update_layout(
    xaxis_title="Method",
    yaxis_title="Test Accuracy",
    yaxis=dict(range=[0, 1])
)
fig.show()

print("\n=== Summary of All Methods ===")
for _, row in comparison_df.iterrows():
    print(f"{row['Method']}: {row['Accuracy']:.3f}")

Problem 9

利用TTA的方式, 把每一个图像增强(旋转)得到很多个, 然后对他们预测取一个平均的概率得到结果

┌─────────────────────────────────────────────────────────────────┐
│                      test_time_augmentation 函数                │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│   输入: image (原始测试图片)                                     │
│                                                                 │
│   Step 1: 预处理                                                 │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ image → 展平 → 标准化                                    │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 2: 预测旋转角度并纠正                                     │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ model_rotation_regressor.predict() → 预测角度           │  │
│   │ rotate_image_by_angle(..., -angle) → 转回来            │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 3: 多角度 TTA 预测                                       │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ 对 unrotated 图像应用 9 种角度旋转:                     │  │
│   │ -20°, -15°, -10°, -5°, 0°, 5°, 10°, 15°, 20°          │  │
│   │ 对每种旋转后的图像 → 标准化 → model.predict_proba()     │  │
│   │ 收集所有概率向量                                          │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 4: 额外预测                                               │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ 原始图像的预测 + 转回来图像的预测                        │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 5: 聚合与最终预测                                         │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ avg_probs = mean(all_probs)  # 11个概率向量取平均       │  │
│   │ prediction = argmax(avg_probs)  # 选择概率最大的类      │  │
│   └─────────────────────────────────────────────────────────┘  │
│                                                                 │
│   输出: class_dict[prediction]  # 返回类别标签                   │
└─────────────────────────────────────────────────────────────────┘

代码看起来很长, 但其实都是流水线, 核心就是一张图片变多张图片, 产生多个概率向量, 求平均在argmax即可

# TODO: Write a function (or functions!) that can be used to improve the accuracy of the original MLPClassifer model using test time augmentations
# Feel free to add additional functions, arguments, etc. as needed!
def test_time_augmentation(model, scaler, image):
    """
    Predict a label using test-time augmentation
    Params:
        - model: the MLPClassifier model
        - scaler: the StandardScaler used to scale the images
        - image: the image to be augmented
    Returns:
        - prediction: the predicted label
    """
    ...
    # raise NotImplementedError("Not implemented")
    image = np.array(image, dtype=np.float64).reshape(-1)
    image_scaled=scaler.transform([image])
    predicted_angle=model_rotation_regression.predict(image_scaled)[0]

    unrotated=rotate_image_by_angle(image,-predicted_angle)

    test_angles=[-20,-15,-10,-5,0,5,10,15,20]

    all_probs=[]

    for angle in test_angles:
        rotated=rotate_image_by_angle(unrotated,angle)
        rotated_scaled=scaler.transform([rotated])
        probs=model.predict_proba(rotated_scaled)[0]
        all_probs.append(probs)

    original_scaled=scaler.transform([image])
    unrotated_scaled=scaler.transform([unrotated])
    all_probs.append(model.predict_proba(original_scaled)[0])
    all_probs.append(model.predict_proba(unrotated_scaled)[0])

    avg_probs=np.mean(all_probs,axis=0)
    prediction=np.argmax(avg_probs)

    return class_dict[prediction]



# Make a copy of the test secret dataframe and apply the test time augmentation to the images to get new predictions!
part_2_df = test_secret_df.copy()
part_2_df["image"] = part_2_df["image"].apply(lambda x: np.array(x).reshape(-1))
part_2_df["prediction"] = part_2_df["image"].apply(lambda x: test_time_augmentation(model_rotated, scaler, x))

# Check the accuracy of the new predictions
correct = part_2_df["prediction"] == part_2_df["label"]
print("Accuracy:", correct.mean())

Accuracy: 0.341

Stanford CS336 Assignment 1(Part 2) - Transformer的实现

Wed, 28 Jan 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

CS336 Assignment 1

Part1中我们实现了BPE分词器的算法以及Tokenizer类, 在Part2中我们将会手搓整个Transformer模型, 然后在Part3中我们会把前面的部分结合来真正的让这个模型开始训练

Transformer LM 简介

Token嵌入层 (Token Embedding)

输入Transformer的内容是token id数字张量(显然不能是char), 形状为(batch_size, sequence_length), 比如说(2,2)

inputs = np.array([[2,0],[1,2]])

然后有一个可训练的矩阵E称之为嵌入矩阵, 这个矩阵回答这样一个问题: 对于每个输入的token_id(一维数字), 怎么把他送到高维空间？

该矩阵形如:

E = np.array([
    [0.0, 0.0, 0.0],  # id = 0 (PAD/UNK)
    [0.1, 0.2, 0.3],  # id = 1 Somewhere
    [0.0, 0.5, 0.5],  # id = 2 over
    [0.9, 0.1, 0.0],  # id = 3 the
    [0.4, 0.4, 0.2],  # id = 4 rainbow
    [0.7, 0.3, 0.6],  # id = 5 way
    [0.2, 0.2, 0.9],  # id = 6 up
    [0.6, 0.1, 0.4],  # id = 7 high
])  # shape (V, d_model) -> here V=8, d_model=3

含义为: 将输入的token_id当中的0送到$$R^3$$上, 值为[0.1,0.2,0.3], 至于这些token_id是哪里来的, 就是从之前训练好的Vocab训练而来的, 完整的流程如下

这是我们输入的自然语言:

Somewhere over the rainbow way up high.

我们有训练好的词汇表Vocab:

{"Somewhere":1, "over",2, "the":3, "rainbow":4, "way":5, "up":6, "high":7}

首先自然语言被分割成

tokens = ["Somewhere","over","the","rainbow","way","up","high"]

然后通过词汇表被编码成

ids = [1,2,3,4,5,6,7]

接着对于ids里面的每一个id, 直接在E里面找到对应的行就可以了, 这就完成了Token Embeddings这个步骤

Pre-Norm Transformer块

经过Embedding处理后的尺寸为(batch_size, sequence_length, d_model)的数据被送入Pre-Norm Transformer块当中, 处理完毕后的尺寸仍为(batch_size, sequence_length, d_model), 块内的组件是自注意力机制和前馈层

输出的归一化

在经过若干个Transformer块之后, 还要进行归一化, 再送入一个线性层做处理, 最后通过Softmax来输出概率logits来决定下一个词输出什么

Einstein算子标注

显然在做矩阵乘法/张量乘法的时候, 我们是在某一个维度进行求和, 而那个维度在计算完成之后会消失, Einstein算子标注就是让我们显性的写出那个被求和/将消失的维度, 其他的维度就不管了

我觉得这种张亮乘法其实类似于高维定积分, 当通过累次积分来计算重积分的时候, 显然要写好这一次是对什么变量进行积分, 这次积分完毕之后, 这个变量就不再存在

$$ \iiint_{V} f(x,y,z),dx,dy,dz = \iint_{V_1},dx,dy \int_{V_2}f(x,y,z),dz $$

后面那次积分的dz已经说明了这次的积分(视作特殊的求和)是对z这个变量的, 所以不会产生混淆

来看一个Einstein标注的张量运算

import torch
from einops import rearrange, einsum
## Basic implementation
Y = D @ A.T
# Hard to tell the input and output shapes and what they mean.
# What shapes can D and A have, and do any of these have unexpected behavior?
## Einsum is self-documenting and robust
# D A -> Y
Y = einsum(D, A, "batch sequence d_in, d_out d_in -> batch sequence d_out")
## Or, a batched version where D can have any leading dimensions but A is constrained.
Y = einsum(D, A, "... d_in, d_out d_in -> ... d_out")

这里我们对d_in这个维度求和, 求和后这个维度就消失掉了, 所以只需要在两个运算量里面都亮明这个维度, 算子就会自动求和

再看一个例子

images = torch.randn(64, 128, 128, 3) # (batch, height, width, channel)
dim_by = torch.linspace(start=0.0, end=1.0, steps=10)
## Reshape and multiply
dim_value = rearrange(dim_by, "dim_value -> 1 dim_value 1 1 1") # 拓展维度, 注意1是可以任意添加的维度
images_rearr = rearrange(images, "b height width channel -> b 1 height width channel") # 同上
dimmed_images = images_rearr * dim_value
## Or in one go:
dimmed_images = einsum(
images, dim_by,
"batch height width channel, dim_value -> batch dim_value height width channel"
)

注意一下广播的时候首先维度的数量要匹配, 其次每个维度的大小要么相等要么有一个是1, 而在Einstein标注下, 不足的维度会被自动填充并广播

高维张量的乘法是不便想象的, 我觉得首先要明白需要的输出尺寸是多少, 然后再用Einstein标注去写好输入的尺寸, 不变的用...替代

线性变换标注

本项目中统一使用列向量做线性变换的notation, 即若对向量x进行线性变换W则:

$$ Y = Wx $$ x默认为列向量

线性层和嵌入模块

参数初始化

对于每一个权重矩阵, 需要把它声明为nn.Parameter参数并且传入shape, 然后对这个参数做初始化, 比如在线性层当中想构造一个out_features * in_features的权重矩阵并且正态初始化

weight = nn.Parameter(torch.empty(out_features, in_features))
nn.init.trunc_normal_(weight,mean = mean, std = std, a = -3*std, b = 3*std)

对于不同的层, 文档给了我们不同的初始化要求

$$ Linear ,, Layer: N(\mu = 0, \sigma^2 = \frac{2}{d_{in} + d_{out}}), , , clip ,, to ,, [-3\sigma,3\sigma] $$

$$ Embedding ,, Layer: N(\mu = 0, \sigma^2 = 1), , , clip ,, to ,, [-3,3] $$

$$ RMSNorm ,, Layer: \mathbb{E}_{n*n} $$

所有的参数初始化都用torch.nn.init.trunc_normal_来实现

线性层

很容易实现, 就是声明并初始化一个权重矩阵并且作用在输入上面就可以了

# transformer.py
class Linear(nn.Module):

    def __init__(self,in_features,out_features,device=None,dtype=None):

        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.device = device
        self.dtype = dtype

        std = 2/(self.in_features+self.out_features)

        self.weight = nn.Parameter(torch.empty(self.out_features, self.in_features))
        nn.init.trunc_normal_(self.weight, mean=0, std=std, a=-3*std, b=3*std)

    def forward(self,x:torch.Tensor) -> torch.Tensor:
        
        return torch.einsum("ji,...i->...j",self.weight,x)

很常规, 解释一下最后这个einsum的notation, 先忽略这个..., 就看"ji,i->j"

$$ \begin{pmatrix} w_{11} ..... w_{1in} \ \ \ \ w_{out1} ...... w_{out in} \end{pmatrix} \begin{pmatrix} x_1 \ \ \ \ x_{in} \end{pmatrix} $$

从逐行求和的角度看, 结果的第j行满足

$$ y_j = \sum_{i} W[j,i] * x[i] $$

所以是对i这个维度求和, 求和的维度会消失, 所以得到输出维度是j

另一种角度看, 消失的总是内在匹配的维度, 所以是这里的i, 即第一个输入的列, 第二个输入的行

完善一下adapters.py里面的测试, 直接调用自己实现的这个Linear类就好了, 注意对于那些权重矩阵, 记得更新他们的内容, 后面所有的测试都是这个格式, 依葫芦画瓢就好了

# adapters.py
from cs336_basics import transformer


def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    """
    Given the weights of a Linear layer, compute the transformation of a batched input.

    Args:
        in_dim (int): The size of the input dimension
        out_dim (int): The size of the output dimension
        weights (Float[Tensor, "d_out d_in"]): The linear weights to use
        in_features (Float[Tensor, "... d_in"]): The output tensor to apply the function to

    Returns:
        Float[Tensor, "... d_out"]: The transformed output of your linear module.
    """

    # raise NotImplementedError
    linear_module = transformer.Linear(in_features=d_in,out_features=d_out)
    linear_module.weight.data = weights
    return linear_module(in_features)

运行测试uv run pytest -k test_linear, 结果如下:

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_linear
==================================================================================================== test session starts =====================================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                                              

tests/test_model.py::test_linear PASSED

============================================================================================== 1 passed, 47 deselected in 0.19s ==============================================================================================

嵌入模块

和我们之前讲的一样, 这里就是把输入的token_id在Embedding矩阵当中去寻找相应的d_model维度的向量, 也就是一个升维的过程, 比如说输入为5, 那就去找Embedding矩阵的第五行的行向量就好了

# transformer.py
class Embedding(nn.Module):
    def __init__(self,num_embeddings,embedding_dim,device=None,dtype=None):


        super().__init__()
        
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.device = device
        self.dtype = dtype

        self.weight = nn.Parameter(torch.empty(self.num_embeddings, self.embedding_dim))
        nn.init.trunc_normal_(self.weight, mean=0, std=1, a=-3, b=3)


    def forward(self,token_ids:torch.Tensor) -> torch.Tensor:
        return self.weight[token_ids]

注意self.weight[token_ids]这种写法, 实际上nn.Parameter是tensor的子类, 是可以通过索引访问的, 举个例子

# weight: (num_embeddings=5, embedding_dim=3)
weight = nn.Parameter(torch.tensor([
    [0.1, 0.2, 0.3],  # id 0
    [0.4, 0.5, 0.6],  # id 1
    [0.7, 0.8, 0.9],  # id 2
    [1.0, 1.1, 1.2],  # id 3
    [1.3, 1.4, 1.5],  # id 4
], dtype=torch.float32))


token_ids = torch.tensor([[2,0],
                          [1,2]], dtype=torch.long)  # shape (2,2)

# 索引取 embedding：结果形状 (2, 2, 3)
embs = weight[token_ids]

输出类似于

tensor([
  [[0.7, 0.8, 0.9],  # weight[2]
   [0.1, 0.2, 0.3]], # weight[0]
  [[0.4, 0.5, 0.6],  # weight[1]
   [0.7, 0.8, 0.9]]  # weight[2]
])

不过也不需要了解, 反正知道支持下标访问就行了, 而且这个下标还可以是一个tensor

# adapters.py
def run_embedding(
    vocab_size: int,
    d_model: int,
    weights: Float[Tensor, " vocab_size d_model"],
    token_ids: Int[Tensor, " ..."],
) -> Float[Tensor, " ... d_model"]:
    """
    Given the weights of an Embedding layer, get the embeddings for a batch of token ids.

    Args:
        vocab_size (int): The number of embeddings in the vocabulary
        d_model (int): The size of the embedding dimension
        weights (Float[Tensor, "vocab_size d_model"]): The embedding vectors to fetch from
        token_ids (Int[Tensor, "..."]): The set of token ids to fetch from the Embedding layer

    Returns:
        Float[Tensor, "... d_model"]: Batch of embeddings returned by your Embedding layer.
    """

    # raise NotImplementedError
    embedding_module = transformer.Embedding(num_embeddings=vocab_size,embedding_dim=d_model)
    embedding_module.weight.data = weights
    return embedding_module(token_ids)

运行测试uv run pytest -k test_embedding, 结果如下：

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_embedding
==================================================================================================== test session starts =====================================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                                              

tests/test_model.py::test_embedding PASSED

============================================================================================== 1 passed, 47 deselected in 0.21s ==============================================================================================

Pre-Norm Transformer块

这里主要需要实现三个模块, RMSNorm归一化, RoPE旋转编码, Feed-Forward前馈神经网络

RMSNorm---Root Mean Square Layer Normalization 均方根归一化

对于向量$$a \in \mathbb{R}^{d_{model}}$$, RMSNorm会对每个分量a_i进行如下变化:

$$ RMSNorm(a_i) = \frac{a_i}{RMS(a)} g_i $$

其中$$RMS(a)$$是一个平方根式求和

$$ RMS(a) =\sqrt{\frac{1}{d_{model}} \sum_{i=1}^{d_{model}} a_i^2 + \epsilon} $$

很好理解, 这个分母真的就是均值->平方求和->开根号(忽略那个$$\epsilon$$)

$$g_i$$是一个可学习的向量, 注意每个$$a_i$$有一个$$g_i$$, 总共有$$d_{model}$$个$$a_i$$, 所以有$$d_{model}$$个$$g_i$$

# transformer.py
class rmsnorm(nn.Module):

    def __init__(self,d_model:int,eps:float=1e-5,device=None,dtype=None):

        super().__init__()

        self.d_model = d_model
        self.eps = eps
        self.device = device
        self.dtype = dtype

        self.weights = nn.Parameter(torch.ones(self.d_model))
        nn.init.trunc_normal_(self.weights, mean=0, std=1, a=-3, b=3)

    def forward(self,x:torch.Tensor) -> torch.Tensor:

        in_dtype = x.dtype
        x = x.to(torch.float32)

        RMS_a = torch.sqrt(torch.einsum("...d,...d->...",x,x)/self.d_model + self.eps)

        return ((x/RMS_a.unsqueeze(-1))*self.weights).to(in_dtype)

解释一下这个RMS_a的算法当中的einsum部分, 实际上这里是要算内积分, 所以对于两个一模一样的东西, 直接按照最后一个维度求和就好了

注意这里这个RMS_a.unsqueeze(-1)是不可以省略的, 因为RMS_a的尺寸是..., 而x的尺寸是...d, 所以必须给RMS_a的尺寸做成...1才能进行除法的广播, 实际上这个RMS_a是个标量, 广播的意思就是每个$$a_i$$都要去除以这个标量, 这就必须要求最后一个维度是对齐的, 不然广播会出错

# adapters.py
def run_rmsnorm(
    d_model: int,
    eps: float,
    weights: Float[Tensor, " d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a RMSNorm affine transform,
    return the output of running RMSNorm on the input features.

    Args:
        d_model (int): The dimensionality of the RMSNorm input.
        eps: (float): A value added to the denominator for numerical stability.
        weights (Float[Tensor, "d_model"]): RMSNorm weights.
        in_features (Float[Tensor, "... d_model"]): Input features to run RMSNorm on. Can have arbitrary leading
            dimensions.

    Returns:
        Float[Tensor,"... d_model"]: Tensor of with the same shape as `in_features` with the output of running
        RMSNorm of the `in_features`.
    """
    # raise NotImplementedError
    rmsnorm_module = transformer.rmsnorm(d_model=d_model,eps=eps)
    rmsnorm_module.weights.data = weights
    return rmsnorm_module(in_features)

运行测试uv run pytest -k test_rmsnorm, 结果如下：

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_rmsnorm
==================================================================================================== test session starts =====================================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                                              

tests/test_model.py::test_rmsnorm PASSED

============================================================================================== 1 passed, 47 deselected in 0.18s ==============================================================================================

前馈神经网络

虽然说是神经网络, 但是其实就是设计一个激活函数, 大概用的有以下几种:

SiLU/Swish

$$ SiLU(x) = x*\sigma(x) = \frac{x}{1+e^{-x}} $$

Gated Linear Units/GLU

$$ GLU(x,W_1,W_2) = \sigma(W_1x) \odot W_2x $$

该算子为Hadamard积, 即逐元素相乘

SwiGLU

$$ SwiGLU(x,W_1,W_2,W_3) = W_2(SiLU(W_1x) \odot W_3x) $$

考虑一下SwiGLU的尺寸问题, 输入的x是(...,d_model), 经过SwiGLU的变换之后仍然应该是这个尺寸, 所以W_1和W_3的尺寸应当是(d_ff, d_model), W_2应该是(d_model, d_ff)

借用torch.sigmoid实现这个SwiGLU的代码:

# transformer.py
class positionwise_feedforward(nn.Module):

    def __init__(self,d_model,d_ff):
        super().__init__()

        self.d_model = d_model
        self.d_ff = d_ff

        self.w1_weight = nn.Parameter(torch.empty(self.d_ff, self.d_model))
        self.w2_weight = nn.Parameter(torch.empty(self.d_model, self.d_ff))
        self.w3_weight = nn.Parameter(torch.empty(self.d_ff, self.d_model))

        nn.init.trunc_normal_(self.w1_weight, mean=0, std=1, a=-3, b=3)
        nn.init.trunc_normal_(self.w2_weight, mean=0, std=1, a=-3, b=3)
        nn.init.trunc_normal_(self.w3_weight, mean=0, std=1, a=-3, b=3)

    def silu(self,x):
        return torch.sigmoid(x) * x

    def element_wise(self,x,y):
        return torch.einsum("...,...->...",x,y)

    def forward(self,x):
        w3x = torch.einsum("...d,fd->...f",x ,self.w3_weight)
        w1x = torch.einsum("...d,fd->...f",x ,self.w1_weight)

        silu_w1x = self.silu(w1x)

        swiglu_ouptut = self.element_wise(silu_w1x,w3x)

        output = torch.einsum("...f,df->...d", swiglu_ouptut, self.w2_weight)

        return output

还是讲一下这里的维度notation, 首先对于Hadamard算子应该是比较好理解的, 前后尺寸都是...,意味着不对任何维度求和, 那当然就是逐元素相乘

对于W_3x而言, 这里写的看起来像是xW_3, 不过关键还是要对齐那个可以匹配的维度, 由于x是(...,d_model), W_3是(d_ff, d_model), 显然写的时候要确保最后一维的尺寸一致(即d), 然后对这个维度求和(即不出现在结果里面就好了)

在工程上或许这么写不会有什么问题, 毕竟这块以后都会封装, 也没有谁真的会去看, 不过如果是数学作业/paper上还是保持notation的顺序比较好, 不然读者的脸色可能不会太好看

# adapters.py
def run_swiglu(
    d_model: int,
    d_ff: int,
    w1_weight: Float[Tensor, " d_ff d_model"],
    w2_weight: Float[Tensor, " d_model d_ff"],
    w3_weight: Float[Tensor, " d_ff d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a SwiGLU network, return
    the output of your implementation with these weights.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        d_ff (int): Dimensionality of the up-project happening internally to your swiglu.
        w1_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W1
        w2_weight (Float[Tensor, "d_model d_ff"]): Stored weights for W2
        w3_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W3
        in_features (Float[Tensor, "... d_model"]): Input embeddings to the feed-forward layer.

    Returns:
        Float[Tensor, "... d_model"]: Output embeddings of the same shape as the input embeddings.
    """

    # raise NotImplementedError
    swiglu_module = transformer.positionwise_feedforward(d_ff=d_ff,d_model=d_model)
    swiglu_module.w1_weight.data = w1_weight
    swiglu_module.w2_weight.data = w2_weight
    swiglu_module.w3_weight.data = w3_weight

    return swiglu_module(in_features)

运行测试uv run pytest -k test_swiglu, 结果如下：

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_swiglu
Uninstalled 1 package in 0.92ms
Installed 1 package in 16ms
================================================================================================= test session starts ==================================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                                        

tests/test_model.py::test_swiglu PASSED

=========================================================================================== 1 passed, 47 deselected in 0.18s ===========================================================================================

相对位置编码

相对位置编码和多头注意力机制应该是最难的, 先不考虑代码的问题, 我们先来理解一下这玩意到底是在干什么

首先RoPE接受一个参数$$\theta$$, 这会决定我们每次旋转的角度

两个index:i和k, i代表的是输入$$x^i$$的索引, 也就是这是输入的第几条输出, k代表的是现在是在对这个输出的第几组维度进行变换, 对于每一组(i, k), 变换的角度$$\theta_{i,k} = \frac{i}{\theta^{\frac{2k-2}{d}}}$$, 其中d指的是d_model, 为固定的模型参数

维度是被两两分组的, 比如说d_model是8, 那就被切分成4组, $$k \in (1,2,3,4)$$

每一组的旋转矩阵为:

$$ R_{k}^{i} = \begin{pmatrix} \cos{\theta_{i,k}}, -\sin{\theta_{i,k}} \ \sin{\theta_{i,k}}, \cos{\theta_{i,k}}

\end{pmatrix} $$

现在对$$q^{1} = W_{q}x^{1} = [1,0,1,0]$$来做变换, $$q^1$$的i=1, k被分为两组, 1和2, 那么可以计算出以下两个$$\theta$$

$$ \theta_{1,1} = \frac{1}{\theta^{0}} = 1 \ \theta_{1,2} = \frac{1}{\theta^{0.5}} $$

进而得到两个旋转矩阵$$R_{1}^{1}, R_{2}^{1}$$, 那么组一是$$[1,0]$$, $$组二是[0,1]$$, 分别对其左乘$$R_{1}^{1}, R_{2}^{1}$$即可

写成比较大的矩阵形式就是:

$$ R^{(i)}=\mathrm{diag}\big(R^{(i)}_1,;R^{(i)}2,;\dots,;R^{(i)}{d/2}\big)

\begin{bmatrix} R^{(i)}_1 & & & 0\[2pt] & R^{(i)}2 & & \[2pt] & & \ddots & \[2pt] 0 & & & R^{(i)}{d/2} \end{bmatrix}, $$

$$ q^{(i)}= \begin{bmatrix} q^{(i)}_1\[4pt] q^{(i)}2\[4pt] \vdots\[4pt] q^{(i)}{d/2} \end{bmatrix}, \qquad q^{(i)}k \in \mathbb{R}^2,\quad q^{(i)}k= \begin{bmatrix} q^{(i)}{2k-1}\[4pt] q^{(i)}{2k} \end{bmatrix} \quad (k=1,\dots,d/2). $$

那么整个的RoPE就是$$R^{i} q^{i}$$, 一句话总结就是先把$$q^{i}/k^{i}$$分组, 然后每组构造一个2*2的旋转矩阵, 对每个组左乘这个旋转矩阵之后凑起来就得到了RoPE的结果

实验文档中提出了用self.register_buffer来注册这些三角函数值, 因为他们只依赖于输入角度theta, 输入的"编号"i以及离散取值于$$[0,\frac{d}{2}]$$的k, 所以可以在初始化的时候就把他们确定下来

注意$$\theta_{i,k}$$, 他的指标应该是i和k的笛卡尔积, 这样才能保证每个i和k都能被遍历到, 所以应当通过以下代码来建立这个$$\theta_{i,k}$$

positions = torch.arange(sequence_length).float() # i的向量

freqs = theta ** (torch.arange(0, self.d_k , 2).float() / self.d_k) # theta^{2k-2/d}的向量

angles = torch.outer(positions, freqs) # i和K的笛卡尔积

然后注册成为类buffer

self.register_buffer("cos_cached", torch.cos(angles), persistent = False)
self.register_buffer("sin_cached", torch.sin(angles), persistent = False)

如此一来的话就可以通过token_positions来访问得到$$\cos(\theta_{i,k})$$和$$\sin(\theta_{i,k})$$了, 例如:

cos_pos = self.cos_cached[token_positions] # 直接拿到这个x所需的所有cos值, 注意token_positions是个tensor
sin_pos = self.sin_cached[token_positions]

回想一下在之前的变换矩阵示意图里面, 这个变换最后并不改变输入x的shape, 但是要能把这个分块d/2 * d/2的变换矩阵作用上去, 是需要对x进行一些reshape的

x_reshaped = x.view(batch_size,seq_len,d_k//2,2) # 分成d_k//2个组, 每组两个元素
x1,x2 = x_reshaped[...,0],x_reshaped[...,1] # 对于每个组, 把第一个元素和第二个元素分开

然后进行变换并且stack回去即可

x1_rotated = x1 * cos_pos - x2 * sin_pos # 注意这里的x_1是: "一个输入token(一个i)"的所有的d_k//2个组的x_1, 不只是一个组的x_1
x2_rotated = x2 * cos_pos + x1 * sin_pos

# 如果说输入的x = [1,2,3,4,5,6], 那么这里的x_1是形如[1,3,5]的tensor, 所以这个变换是向量化的
x_rotated = torch.stack([x1_rotated, x2_rotated], dim=-1)
x_rotated = x_rotated.view(batch_size, seq_len, d_k)

整体代码如下:

# transformer.py
class RoPE(nn.Module):

    def __init__(self,theta,d_k,max_seq_len,device=None):
        super().__init__()

        self.theta = theta
        self.d_k = d_k
        self.max_seq_len = max_seq_len
        self.device = device

        freqs = self.theta ** (torch.arange(0,self.d_k,2).float()/d_k)

        positions = torch.arange(max_seq_len).float()

        angles = torch.outer(positions,1.0/freqs)

        self.register_buffer("cos_cached", torch.cos(angles),persistent=False)
        self.register_buffer("sin_cached", torch.sin(angles),persistent=False)

    def forward(self,x:torch.Tensor,token_positions:torch.Tensor) -> torch.Tensor:

        batch_size,seq_len,d_k = x.shape

        cos_pos = self.cos_cached[token_positions]
        sin_pos = self.sin_cached[token_positions]

        x_reshaped = x.view(batch_size,seq_len,d_k//2,2)

        x1,x2 = x_reshaped[...,0],x_reshaped[...,1]

        x1_rotated = x1 * cos_pos - x2 * sin_pos
        x2_rotated = x2 * cos_pos + x1 * sin_pos

        x_rotated = torch.stack([x1_rotated, x2_rotated], dim=-1)
        x_rotated = x_rotated.view(batch_size, seq_len, d_k)
        
        return x_rotated

完善测试:

# adapters.py
def run_rope(
    d_k: int,
    theta: float,
    max_seq_len: int,
    in_query_or_key: Float[Tensor, " ... sequence_length d_k"],
    token_positions: Int[Tensor, " ... sequence_length"],
) -> Float[Tensor, " ... sequence_length d_k"]:
    """
    Run RoPE for a given input tensor.

    Args:
        d_k (int): Embedding dimension size for the query or key tensor.
        theta (float): RoPE parameter.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        in_query_or_key (Float[Tensor, "... sequence_length d_k"]): Input tensor to run RoPE on.
        token_positions (Int[Tensor, "... sequence_length"]): Tensor of shape (batch_size, sequence_length) with the token positions
    Returns:
        Float[Tensor, " ... sequence_length d_k"]: Tensor with RoPEd input.
    """
    # raise NotImplementedError
    RoPE = transformer.RoPE(theta=theta,d_k=d_k,max_seq_len=max_seq_len)
    return RoPE(in_query_or_key,token_positions)

运行测试uv run pytest -k test_rope, 结果如下：

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_rope
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                         

tests/test_model.py::test_rope PASSED

=================================================================================== 1 passed, 47 deselected in 0.17s ====================================================================================

Softmax和点积注意力机制

首先实现Softmax函数, 对向量$$v$$做归一化

$$ Softmax(v)i = \frac{exp(v_i)}{\sum{j=1}^{n} exp(v_j)} $$

注意求指数有可能会让数据变得很大从而上溢出, 所以一个比较好的办法是让这个向量的每个分量减去这个向量的最大分量, 因为对一个很小的数求指数是不会溢出的

# transformer.py
class Softmax(nn.Module):
    
    def __init__(self, x:torch.Tensor, dimension:int):
        super().__init__()

        self.x = x
        self.dimension = dimension


    def forward(self):
        x_shifted = self.x - torch.max(self.x,dim = self.dimension,keepdim=True)[0]

        exp_x = torch.exp(x_shifted)

        sum_exp_x = torch.sum(exp_x, dim=self.dimension, keepdim=True)

        return exp_x / sum_exp_x

完善测试:

# adapters.py
def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
    """
    Given a tensor of inputs, return the output of softmaxing the given `dim`
    of the input.

    Args:
        in_features (Float[Tensor, "..."]): Input features to softmax. Shape is arbitrary.
        dim (int): Dimension of the `in_features` to apply softmax to.

    Returns:
        Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
        softmax normalizing the specified `dim`.
    """
    # raise NotImplementedError
    softmax_module = transformer.Softmax(in_features,dim)
    return softmax_module()

注意这个dim是必须的参数, 因为需要知道对什么维度进行归一化

运行测试uv run pytest -k test_softmax_matches_pytorch, 结果如下：

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_softmax_matches_pytorch
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                         

tests/test_nn_utils.py::test_softmax_matches_pytorch PASSED

=================================================================================== 1 passed, 47 deselected in 0.17s ====================================================================================

接下来计算点积注意力

$$ Attention(Q,K,V) = Softmax(\frac{Q^TK}{\sqrt{d_k}})V $$

先考虑一下尺寸问题, 简单来说, 假设$$Q \in \mathbb{R}^{n \times d_k}$$, $$K \in \mathbb{R}^{m \times d_k}$$, $$V \in \mathbb{R}^{m \times d_v}$$, 那么输出尺寸应该是$$ \mathbb{R}^{n \times d_v}$$, 但是$$Q^TK$$如果按照一般的矩阵乘法是不能相乘的, 这里其实对这两个矩阵的每个行向量做内积

$$ Q^TK_{ij} = \sum_{k=1}^{d_k} Q_{ik} K_{kj} $$

注意求和, 每行有$$d_k$$个元素, 实验文档要求我们用einsum来实现这个求和, 如果是用pytorch的话, 这里要手动修改成$$QK^T$$的乘积形式

但我觉得这种记号并不好, 既然用了矩阵乘法的记号, 那应该要确保按照记号是可相乘的, 不然当发现维度不匹配的时候会很confused

scores = torch.einsum("b...qd,b...kd->b...qk",self.Q,self.K)/torch.sqrt(torch.tensor(self.d_k,dtype=torch.float32))
# kd 指的是Q当中的n * d_k
# qd 指的是K当中的m * d_k
# 匹配的是最后一维, einsum会自动把他们转置成可以相乘的形式

然后用mask做掩码变换, 这里mask的维度是$$n \times m$$, 暂时不需要管他怎么实现的, 当作类里有的成员变量就好, 既然他是个布尔矩阵, 那只要用torch.where去把mask和计算出的分数做一个类似与运算就好了

if self.mask is not None:
    scores = torch.where(self.mask,scores,float('-inf'))

接下来做Softmax变换, 对最后一维做归一化

attention_weights = torch.softmax(scores, dim=-1)

最后乘以$$V$$即可, 注意前面的结果是可以"直接"乘以$$V$$的, 因为维度已经匹配了, 而像之前那样维度不匹配的情况, einsum会自动把他们转置成可以相乘的形式

output = torch.einsum("b...qk,b...kv->b...qv",attention_weights,self.V)

return output

总体实现如下

# transformer.py
class scaled_dot_product_attention(nn.Module):

    def __init__(self,K,Q,V,mask):
        super().__init__()
        self.Q = Q # [batch,...seq_q,d_q]
        self.K = K # [batch,...,seq_k,d_k]
        self.V = V # [batch,...,seq_k,d_v]
        self.mask = mask # [seq_q,seq_k]
        self.d_k = Q.shape[-1]

    def forward(self):
        scores = torch.einsum("b...qd,b...kd->b...qk",self.Q,self.K)/torch.sqrt(torch.tensor(self.d_k,dtype=torch.float32))
        
        if self.mask is not None:
            scores = torch.where(self.mask,scores,float('-inf'))

        attention_weights = torch.softmax(scores, dim=-1)

        output = torch.einsum("b...qk,b...kv->b...qv",attention_weights,self.V)

        return output

完善测试:

# adapters.py
def run_scaled_dot_product_attention(
    Q: Float[Tensor, " ... queries d_k"],
    K: Float[Tensor, " ... keys d_k"],
    V: Float[Tensor, " ... values d_v"],
    mask: Bool[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
    """
    Given key (K), query (Q), and value (V) tensors, return
    the output of your scaled dot product attention implementation.

    Args:
        Q (Float[Tensor, " ... queries d_k"]): Query tensor
        K (Float[Tensor, " ... keys d_k"]): Key tensor
        V (Float[Tensor, " ... values d_v"]): Values tensor
        mask (Bool[Tensor, " ... queries keys"] | None): Mask tensor
    Returns:
        Float[Tensor, " ... queries d_v"]: Output of SDPA
    """
    # raise NotImplementedError
    attention = transformer.scaled_dot_product_attention(K,Q,V,mask)
    return attention()

这里有两个测试, 分别运行uv run pytest -k test_scaled_dot_product_attention和uv run pytest -k test_4d_scaled_dot_product_attention, 结果如下

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_scaled_dot_product_attention
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                         

tests/test_model.py::test_scaled_dot_product_attention PASSED

=================================================================================== 1 passed, 47 deselected in 0.17s ====================================================================================


(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_4d_scaled_dot_product_attention
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                         

tests/test_model.py::test_4d_scaled_dot_product_attention PASSED

=================================================================================== 1 passed, 47 deselected in 0.44s ====================================================================================

多头注意力机制

之前我们只关心怎么通过$$Q,K,V$$来计算出注意力, 但是没有探究这个$$Q,K,V$$是怎么通过输入x得到的, 实际上我们有以下的流程图

输入 x:        (batch, seq_len, d_model)
                    ↓
              ┌────┴────┐
              ↓         ↓
WQ x:    (batch, seq_len, h×d_k)     WK x: (batch, seq_len, h×d_k)
WV x:    (batch, seq_len, h×d_v)     
              ↓
        按头切分 (split) 总共h个头
              ↓
Q: (batch, seq_len, h, d_k)  →  view成 (batch×h, seq_len, d_k)
K: (batch, seq_len, h, d_k)  →  view成 (batch×h, seq_len, d_k)
V: (batch, seq_len, h, d_v)  →  view成 (batch×h, seq_len, d_v)
              ↓
        多头注意力计算
              ↓
输出: (batch×h, seq_len, d_v)
              ↓
        Concat所有头
              ↓
输出: (batch, seq_len, h×d_v)
              ↓
              WO
              ↓
最终输出: (batch, seq_len, d_model)

也就是说, 先通过三个不同的线性层把输入x变换到$$Q,K,V$$三个矩阵, 然后通过多头注意力机制计算出注意力, 最后通过一个线性层把多个头的结果拼接起来, 再通过一个线性层得到最终的输出

self.d_model = d_model
self.num_heads = num_heads
self.q_proj = Linear(d_model, num_heads * self.d_k)
self.k_proj = Linear(d_model, num_heads * self.d_k)
self.v_proj = Linear(d_model, num_heads * self.d_k)
self.output_proj = Linear(num_heads * self.d_v, d_model)

batch_size, seq_len, d_model = x.shape

Q = self.q_proj(x)  # [batch, seq, num_heads * d_k]
K = self.k_proj(x)  # [batch, seq, num_heads * d_k]
V = self.v_proj(x)  # [batch, seq, num_heads * d_k]

# Rearrange to separate heads
Q = rearrange(Q, "b s (h d) -> b h s d", h=self.num_heads)
K = rearrange(K, "b s (h d) -> b h s d", h=self.num_heads)
V = rearrange(V, "b s (h d) -> b h s d", h=self.num_heads)

如果要进行旋转位置编码, 注意每个头上都要应用

if self.use_rope:
    # 创建 token_positions [seq_len]
    token_positions = torch.arange(seq_len, device=x.device)
    # token_positions = self.token_positions

    # 为每个头应用 RoPE，形状 [batch, num_heads, seq, d_k]
    for head in range(self.num_heads):
        Q[:, head, :, :] = self.rope(Q[:, head, :, :], token_positions.unsqueeze(0))
        K[:, head, :, :] = self.rope(K[:, head, :, :], token_positions.unsqueeze(0))

还需要计算掩码矩阵, 首先考虑尺寸, 实际上mask的尺寸和QK^T的尺寸是一样的, 这里有:

Q: (batch, num_heads, seq_len, d_k)
K: (batch, num_heads, seq_len, d_k)

所以mask的尺寸应该是(seq_len, seq_len)

再看mask的元素, 行表示query, 列表示key, 第i个query只能看到前i个'key', 所以这是个下三角矩阵

allow_mask[i][j] = True 表示 query_i 可以 attend to key_j

        key_0  key_1  key_2  key_3
        ─────  ─────  ─────  ─────
query_0 │  ✓     ✗     ✗     ✗    ← 只能看自己及之前
query_1 │  ✓     ✓     ✗     ✗    ← 只能看自己及之前
query_2 │  ✓     ✓     ✓     ✗    ← 只能看自己及之前
query_3 │  ✓     ✓     ✓     ✓    ← 只能看自己及之前

# 添加因果掩码
causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
# scaled_dot_product_attention 期望 mask True = 允许，False = 屏蔽
# 但我们的 causal_mask True = 屏蔽，所以需要取反
allow_mask = ~causal_mask
allow_mask = allow_mask.unsqueeze(0).unsqueeze(1).expand(batch_size, self.num_heads, -1, -1)

然后把得到的$$K,Q,V$$传入之前实现的点积注意力机制, 注意这里的$$K,Q,V$$都是有head这个维度的, 所以返回的结果也是分head的, 要再通过rearrange把他们拼回去

# 使用 scaled_dot_product_attention 类进行计算
# 注意：参考实现使用 (K, Q, V, mask) 的顺序
attn = scaled_dot_product_attention(K, Q, V, allow_mask)
attended_values = attn()  # [batch, num_heads, seq, d_k]

# Rearrange back to [batch, seq, num_heads * d_k]
attended_values = rearrange(attended_values, "b h s d -> b s (h d)", h=self.num_heads)

output = self.output_proj(attended_values)  # [batch, seq, d_model]

return output

整个实现如下:

# transformer.py
class multihead_self_attention(nn.Module):

    def __init__(self, d_model, num_heads, use_rope=True,max_seq_len=1024, theta=10000,token_positions=None):
        super().__init__()

        self.d_model = d_model
        self.num_heads = num_heads
        self.use_rope = use_rope
        self.token_positions = token_positions

        self.d_k = d_model // num_heads
        self.d_v = self.d_k

        self.q_proj = Linear(d_model, num_heads * self.d_k)
        self.k_proj = Linear(d_model, num_heads * self.d_k)
        self.v_proj = Linear(d_model, num_heads * self.d_k)
        self.output_proj = Linear(num_heads * self.d_v, d_model)

        if self.use_rope:
            self.rope = RoPE(theta=theta, d_k=self.d_k, max_seq_len=max_seq_len)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, d_model = x.shape

        Q = self.q_proj(x)  # [batch, seq, num_heads * d_k]
        K = self.k_proj(x)  # [batch, seq, num_heads * d_k]
        V = self.v_proj(x)  # [batch, seq, num_heads * d_k]

        # Rearrange to separate heads
        Q = rearrange(Q, "b s (h d) -> b h s d", h=self.num_heads)
        K = rearrange(K, "b s (h d) -> b h s d", h=self.num_heads)
        V = rearrange(V, "b s (h d) -> b h s d", h=self.num_heads)

        if self.use_rope:
            # 创建 token_positions [seq_len]
            token_positions = torch.arange(seq_len, device=x.device)
            # token_positions = self.token_positions

            # 为每个头应用 RoPE，形状 [batch, num_heads, seq, d_k]
            for head in range(self.num_heads):
                Q[:, head, :, :] = self.rope(Q[:, head, :, :], token_positions.unsqueeze(0))
                K[:, head, :, :] = self.rope(K[:, head, :, :], token_positions.unsqueeze(0))

        # 添加因果掩码
        causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        # scaled_dot_product_attention 期望 mask True = 允许，False = 屏蔽
        # 但我们的 causal_mask True = 屏蔽，所以需要取反
        allow_mask = ~causal_mask
        allow_mask = allow_mask.unsqueeze(0).unsqueeze(1).expand(batch_size, self.num_heads, -1, -1)

        # 使用 scaled_dot_product_attention 类进行计算
        # 注意：参考实现使用 (K, Q, V, mask) 的顺序
        attn = scaled_dot_product_attention(K, Q, V, allow_mask)
        attended_values = attn()  # [batch, num_heads, seq, d_k]

        # Rearrange back to [batch, seq, num_heads * d_k]
        attended_values = rearrange(attended_values, "b h s d -> b s (h d)", h=self.num_heads)

        output = self.output_proj(attended_values)  # [batch, seq, d_model]

        return output

完善测试接口:

# adapters.py

def run_multihead_self_attention(
    d_model: int,
    num_heads: int,
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
    o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
) -> Float[Tensor, " ... sequence_length d_out"]:
    """
    Given the key, query, and value projection weights of a naive unbatched
    implementation of multi-head attention, return the output of an optimized batched
    implementation. This implementation should handle the key, query, and value projections
    for all heads in a single matrix multiply.
    This function should not use RoPE.
    See section 3.2.2 of Vaswani et al., 2017.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        num_heads (int): Number of heads to use in multi-headed attention.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        q_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the Q projection
        k_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the K projection
        v_proj_weight (Float[Tensor, "d_k d_in"]): Weights for the V projection
        o_proj_weight (Float[Tensor, "d_model d_v"]): Weights for the output projection
        in_features (Float[Tensor, "... sequence_length d_in"]): Tensor to run your implementation on.

    Returns:
        Float[Tensor, " ... sequence_length d_out"]: Tensor with the output of running your optimized, batched multi-headed attention
        implementation with the given QKV projection weights and input features.
    """
    # raise NotImplementedError
    multihead = transformer.multihead_self_attention(d_model=d_model, num_heads=num_heads, use_rope=False)
    multihead.q_proj.weight.data = q_proj_weight
    multihead.k_proj.weight.data = k_proj_weight
    multihead.v_proj.weight.data = v_proj_weight
    multihead.output_proj.weight.data = o_proj_weight

    return multihead(in_features)

运行测试uv run pytest -k test_multihead_self_attention, 结果如下:

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_multihead_self_attention
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 46 deselected / 2 selected                                                                                                                                                         

tests/test_model.py::test_multihead_self_attention PASSED
tests/test_model.py::test_multihead_self_attention_with_rope PASSED

=================================================================================== 2 passed, 46 deselected in 0.43s ====================================================================================

组装Transformer

回忆一下架构图里面, 但我们有了Transformer Block Embedding Layer RMSNorm Linear Softmax之后, 我们就可以来拼装完整的Transformer了

单个的Transformer Block

照着图拼就行了, 不过要注意类似ResNet的结构

# transformer.py
class transformer_block(nn.Module):

    def __init__(self,d_model,num_heads,d_ff,use_rope=True,max_seq_len=1024,theta=10000):
        super().__init__()

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff

        self.norm1 = rmsnorm(d_model = self.d_model)
        self.norm2 = rmsnorm(d_model = self.d_model)
        self.attn = multihead_self_attention(d_model = self.d_model, num_heads = self.num_heads, max_seq_len=max_seq_len, theta=theta,use_rope=use_rope)
        self.ffn = positionwise_feedforward(d_model = self.d_model, d_ff = self.d_ff)


    def forward(self,x):

        block1_output = x + self.attn(self.norm1(x))
        block2_output = block1_output + self.ffn(self.norm2(block1_output))

        return block2_output

完善测试接口:

# adapters.py

def run_transformer_block(
    d_model: int,
    num_heads: int,
    d_ff: int,
    max_seq_len: int,
    theta: float,
    weights: dict[str, Tensor],
    in_features: Float[Tensor, " batch sequence_length d_model"],
) -> Float[Tensor, " batch sequence_length d_model"]:
    """
    Given the weights of a pre-norm Transformer block and input features,
    return the output of running the Transformer block on the input features.

    This function should use RoPE.
    Depending on your implementation, you may simply need to pass the relevant args
    to your TransformerBlock constructor, or you may need to initialize your own RoPE
    class and pass that instead.

    Args:
        d_model (int): The dimensionality of the Transformer block input.
        num_heads (int): Number of heads to use in multi-headed attention. `d_model` must be
            evenly divisible by `num_heads`.
        d_ff (int): Dimensionality of the feed-forward inner layer.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        theta (float): RoPE parameter.
        weights (dict[str, Tensor]):
            State dict of our reference implementation.
            The keys of this dictionary are:
            - `attn.q_proj.weight`
                The query projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.q_proj.weight == torch.cat([q_heads.0.weight, ..., q_heads.N.weight], dim=0)`.
            - `attn.k_proj.weight`
                The key projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.k_proj.weight == torch.cat([k_heads.0.weight, ..., k_heads.N.weight], dim=0)`.
            - `attn.v_proj.weight`
                The value projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_v),
                so `attn.v_proj.weight == torch.cat([v_heads.0.weight, ..., v_heads.N.weight], dim=0)`.
            - `attn.output_proj.weight`
                Weight of the multi-head self-attention output projection
                Shape is (d_model, d_model).
            - `ln1.weight`
                Weights of affine transform for the first RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `ffn.w1.weight`
                Weight of the first linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `ffn.w2.weight`
                Weight of the second linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `ffn.w3.weight`
                Weight of the third linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `ln2.weight`
                Weights of affine transform for the second RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
        in_features (Float[Tensor, "batch sequence_length d_model"]):
            Tensor to run your implementation on.

    Returns:
        Float[Tensor, "batch sequence_length d_model"] Tensor with the output of
        running the Transformer block on the input features while using RoPE.
    """
    # raise NotImplementedError
    transformer_block = transformer.transformer_block(d_model=d_model, num_heads=num_heads, d_ff=d_ff, max_seq_len=max_seq_len, theta=theta)
    transformer_block.norm1.weights.data = weights['ln1.weight']
    transformer_block.norm2.weights.data = weights['ln2.weight']

    transformer_block.attn.q_proj.weight.data = weights['attn.q_proj.weight']
    transformer_block.attn.k_proj.weight.data = weights['attn.k_proj.weight']
    transformer_block.attn.v_proj.weight.data = weights['attn.v_proj.weight']
    transformer_block.attn.output_proj.weight.data = weights['attn.output_proj.weight']
    transformer_block.ffn.w1_weight.data = weights['ffn.w1.weight']
    transformer_block.ffn.w2_weight.data = weights['ffn.w2.weight']
    transformer_block.ffn.w3_weight.data = weights['ffn.w3.weight']

    return transformer_block(in_features)

运行测试uv run pytest -k test_transformer_block, 结果如下:

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_transformer_block
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 47 deselected / 1 selected                                                                                                                                                         

tests/test_model.py::test_transformer_block PASSED

=========================================================================================== warnings summary ============================================================================================
tests/adapters.py:352
  /home/zyli/Stanford_CS336/assignment1-basics/tests/adapters.py:352: SyntaxWarning: invalid escape sequence '\T'
    rope_theta (float): The RoPE $\Theta$ parameter.

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================================== 1 passed, 47 deselected, 1 warning in 0.44s ==============================================================================
(base) zyli@lab:~/Stanford_CS336/assignment1-basics$

警告可以忽略, 这应该是课程组在写注释的时候写了LaTeX风格的字符导致解析出问题

完整的Transformer

和上面一样, 中间的若干个Transformer Block Layer用nn.ModuleList`来实现

# transformer.py
class transformer_lm(nn.Module):

    def __init__(self,d_model,num_heads,d_ff,vocab_size,context_length,num_layers,use_rope,max_seq_len=1024,theta=10000):
        super().__init__()

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.vocab_size = vocab_size
        self.context_length = context_length
        self.num_layers = num_layers
        self.max_seq_len = max_seq_len
        self.theta = theta

        self.Token_Embedding = Embedding(num_embeddings=self.vocab_size, embedding_dim = self.d_model)
        self.layers = nn.ModuleList([transformer_block(d_model=self.d_model, num_heads=self.num_heads, d_ff=self.d_ff, use_rope=use_rope,max_seq_len=self.max_seq_len, theta=self.theta) for _ in range(self.num_layers)])
        self.norm = rmsnorm(d_model = self.d_model)
        self.linear = Linear(in_features=self.d_model, out_features=self.vocab_size)


    def forward(self,x):

        x = self.Token_Embedding(x)

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)

        x = self.linear(x)

        # softmax_layer = Softmax(x, dimension=-1)
        # return softmax_layer.forward()

        return x

完善测试接口:

# adapters.py
def run_transformer_lm(
    vocab_size: int,
    context_length: int,
    d_model: int,
    num_layers: int,
    num_heads: int,
    d_ff: int,
    rope_theta: float,
    weights: dict[str, Tensor],
    in_indices: Int[Tensor, " batch_size sequence_length"],
) -> Float[Tensor, " batch_size sequence_length vocab_size"]:
    """Given the weights of a Transformer language model and input indices,
    return the output of running a forward pass on the input indices.

    This function should use RoPE.

    Args:
        vocab_size (int): The number of unique items in the output vocabulary to be predicted.
        context_length (int): The maximum number of tokens to process at once.
        d_model (int): The dimensionality of the model embeddings and sublayer outputs.
        num_layers (int): The number of Transformer layers to use.
        num_heads (int): Number of heads to use in multi-headed attention. `d_model` must be
            evenly divisible by `num_heads`.
        d_ff (int): Dimensionality of the feed-forward inner layer (section 3.3).
        rope_theta (float): The RoPE $\Theta$ parameter.
        weights (dict[str, Tensor]):
            State dict of our reference implementation. {num_layers} refers to an
            integer between `0` and `num_layers - 1` (the layer index).
            The keys of this dictionary are:
            - `token_embeddings.weight`
                Token embedding matrix. Shape is (vocab_size, d_model).
            - `layers.{num_layers}.attn.q_proj.weight`
                The query projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.q_proj.weight == torch.cat([q_heads.0.weight, ..., q_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.k_proj.weight`
                The key projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.k_proj.weight == torch.cat([k_heads.0.weight, ..., k_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.v_proj.weight`
                The value projections for all `num_heads` attention heads.
                Shape is (num_heads * (d_model / num_heads), d_model).
                The rows are ordered by matrices of shape (num_heads, d_v),
                so `attn.v_proj.weight == torch.cat([v_heads.0.weight, ..., v_heads.N.weight], dim=0)`.
            - `layers.{num_layers}.attn.output_proj.weight`
                Weight of the multi-head self-attention output projection
                Shape is ((d_model / num_heads) * num_heads, d_model).
            - `layers.{num_layers}.ln1.weight`
                Weights of affine transform for the first RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `layers.{num_layers}.ffn.w1.weight`
                Weight of the first linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `layers.{num_layers}.ffn.w2.weight`
                Weight of the second linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `layers.{num_layers}.ffn.w3.weight`
                Weight of the third linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `layers.{num_layers}.ln2.weight`
                Weights of affine transform for the second RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `ln_final.weight`
                Weights of affine transform for RMSNorm applied to the output of the final transformer block.
                Shape is (d_model, ).
            - `lm_head.weight`
                Weights of the language model output embedding.
                Shape is (vocab_size, d_model).
        in_indices (Int[Tensor, "batch_size sequence_length"]) Tensor with input indices to run the language model on. Shape is (batch_size, sequence_length), where
            `sequence_length` is at most `context_length`.

    Returns:
        Float[Tensor, "batch_size sequence_length vocab_size"]: Tensor with the predicted unnormalized
        next-word distribution for each token.
    """
    transformerlm = transformer.transformer_lm(d_model=d_model,num_heads=num_heads,d_ff=d_ff,vocab_size=vocab_size,
    context_length=context_length,num_layers=num_layers,use_rope=True,theta=rope_theta)

    transformerlm.Token_Embedding.weight.data = weights['token_embeddings.weight']
    for layer_idx in range(num_layers):
        block = transformerlm.layers[layer_idx]

        # Attention weights
        block.attn.q_proj.weight.data = weights[f'layers.{layer_idx}.attn.q_proj.weight']
        block.attn.k_proj.weight.data = weights[f'layers.{layer_idx}.attn.k_proj.weight']
        block.attn.v_proj.weight.data = weights[f'layers.{layer_idx}.attn.v_proj.weight']
        block.attn.output_proj.weight.data = weights[f'layers.{layer_idx}.attn.output_proj.weight']
        
        # RMSNorm weights
        block.norm1.weights.data = weights[f'layers.{layer_idx}.ln1.weight']
        block.norm2.weights.data = weights[f'layers.{layer_idx}.ln2.weight']
        
        # FFN weights
        block.ffn.w1_weight.data = weights[f'layers.{layer_idx}.ffn.w1.weight']
        block.ffn.w2_weight.data = weights[f'layers.{layer_idx}.ffn.w2.weight']
        block.ffn.w3_weight.data = weights[f'layers.{layer_idx}.ffn.w3.weight']

    transformerlm.norm.weights.data = weights['ln_final.weight']
    transformerlm.linear.weight.data = weights['lm_head.weight']

    return transformerlm(in_indices)
    # raise NotImplementedError

运行测试uv run pytest -k test_transformer_lm, 结果如下:

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest -k test_transformer_lm
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 48 items / 46 deselected / 2 selected                                                                                                                                                         

tests/test_model.py::test_transformer_lm PASSED
tests/test_model.py::test_transformer_lm_truncated_input PASSED

=================================================================================== 2 passed, 46 deselected in 0.46s ====================================================================================

计算次数估算

模型训练的时候, 绝大部分运算都是矩阵乘法里面的数值乘法运算, 注意以下事实: 如果$$A \in \mathbb{R}^{m \times n}$$, $$B \in \mathbb{R}^{n \times p}$$, 那么$$A \times B$$的计算次数是$$2 \times m \times n \times p$$

参数数量计算和内存用量

考虑第一个问题, 对于以下配置的$$GPT-2 XL$$, 有多少个可训练的参数? 如果每个参数都是单精度浮点数, 那么读取这个模型需要多少内存?

vocab_size = 50257
context_length = 1024
num_layers = 48
d_model = 1600
nun_heads = 25
d_ff = 6400

对于Embedding Layer, 有$$vocabsize \times d_{model}$$个参数

50257 * 1600 = 80,411,200

对于一个Transformer Block, 有两个RMSNorm, 每个RMSNorm有$$d_{model}$$个参数, 有一个multihead_self_attention, 其中有四个矩阵$$Q,K,V,output$$, 参数量总共为$$4*d_{model}d_{model}$$, 对于positionwise_feedforward, 有两个矩阵, 总参数量为$2d_{ff}*d_{model}$

2 * 1600 + 4 * 1600 * 1600 + 2 * 1600 * 6400 = 3200 + 10240000 + 20480000 = 30,723,200

除此之外, 输出前还有一个RMSNorm,一个Linear, 总参数量为$$d_{model} + d_{model}*vocabsize$$

1600 + 1600 * 50257 = 80,412,800

故总共有

80411200 + 30723200 * 48 + 80412800 = 1,635,537,600

大概1.6B的参数量

1,635,537,600 * 4 /1024 /1024 /1024 = 6.1GB

需要内存约6.1GB

FLOPS次数估计

如果这个模型前向传播一次, 需要多少次FLOPS?

在一个transformer_block中, 一次multihead_self_attention的计算次数由以下部分组成:

Q投影: 2 * seq_len * d_model * d_model
K投影: 2 * seq_len * d_model * d_model
V投影: 2 * seq_len * d_model * d_model
输出投影: 2 * seq_len * d_model * d_model
QK^T: 2 * seq_len * seq_len * d_model
Attn V: 2 * seq_len * seq_len * d_model
共计8 * seq_len * d_model² + 4 * seq_len² * d_model = 20,971,520,000 + 67,108,864,000 = 88,080,384,000

一次positionwise_feedforward的计算次数由以下部分组成:

w_1x: 2 * seq_len * d_model * d_ff
W_3x: 2 * seq_len * d_model * d_ff
W2 @ (SiLU(W1x) ⊙ W3x): 2 * seq_len * d_model * d_ff
共计6 * seq_len * d_model * d_ff = 62,914,560,000

输出前有一个Output_Embedding的线性层, 计算次数为

2 * seq_len * d_model * vocab_size = 164,682,137,600

代入配置, 得到

48 * (88,080,384,000 + 62,914,560,000) + 164,682,137,600 = 7,412,439,449,600

大约7.4TeraFLOPs

UC Berkeley CS189 Assignment 1(Part 1)

Mon, 26 Jan 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

CS189 Assignment1

189这门课感觉也是理论和实践并行教学, Written Part的作业也不少, 不过比起336那种工程性质极强的课程来说还是容易不少, 至少这第一个Assignment没有和336那样, 一上来就整一大堆手搓, 189的作业基本上就是理解原来并调包, 极少的自己实现, 非常适合有志于成为API工程师的人学习

Written Part的数学作业也并不容易, 后面有机会写一些Notes

作业要求

基本上是Fill in the Blanks性质的Lab, 把函数里面的所有TODO填满就行了, 几乎不需要自己设计接口

评测框架

用的是UC Berkeley自行设计的otter框架, 一键安装就行了:

pip install otter-grader

环境配置

每个hw文件当中都给了requirements.txt, 直接安装就行

pip install -r requirements.txt

除此之外就是配置一下网络环境, 由于我是在GPU服务器上运行的, 所以一些没找到镜像站的数据下载代码我就用反向ssh端口让他走我本地的代理, 打开本地终端开启这个反向ssh:

(base) liziyu@liziyudeMacBook-Pro ~ % ssh -R 7891:localhost:7890 zyli@lab
Welcome to Ubuntu 24.04.3 LTS (GNU/Linux 6.14.0-29-generic x86_64)

 * Documentation:  https://help.ubuntu.com
 * Management:     https://landscape.canonical.com
 * Support:        https://ubuntu.com/pro

Expanded Security Maintenance for Applications is not enabled.

124 updates can be applied immediately.
To see these additional updates run: apt list --upgradable

Enable ESM Apps to receive additional future security updates.
See https://ubuntu.com/esm or run: sudo pro status

*** System restart required ***
Last login: Fri Jan 23 15:16:55 2026 from 61.169.135.222
(base) zyli@lab:~$

然后在notebook里面通过os来配置一下端口:

import os
where="lab"
# where="local"

if where=="lab":

    os.environ['HTTP_PROXY'] = 'http://localhost:7891'
    os.environ['HTTPS_PROXY'] = 'http://localhost:7891'

这里其实是比较烦的, 我本人还是喜欢去运行py文件, 用ALL_PROXY=localhost:7891的形式去做

Assignment Overview

第一部分的所有代码在fashion_pt_1.ipynb里面, 都是一些比较基础的pandas和numpy的操作, 还有一些画图之类的, 主要就是让学习者习惯数据预处理

加载`Fashion-MNIST`数据集

搞清楚数据集本身作为一个对象有什么属性就行

# Load the FashionMNIST dataset from torchvision
train_data = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True)

# Extract the image data and convert it to a numpy array of type float
images = train_data.data.numpy().astype(float)

# Extract the target labels as a numpy array
targets = train_data.targets.numpy()

# Create a dictionary mapping class indices to class names
class_dict = {i: class_name for i, class_name in enumerate(train_data.classes)}

# Map the target labels to their corresponding class names
labels = np.array([class_dict[t] for t in targets])

# Create a list of class names in order of their indices
class_names = [class_dict[i] for i in range(len(class_dict))]

# Get the total number of samples in the dataset
n = len(images)

# Ensure class_names is a list of class names (redundant but ensures consistency)
class_names = list(class_dict.values())

# Print dataset information for verification
print("Loaded FashionMNIST dataset with {} samples.".format(n))
print("Classes: {}".format(class_dict))
print("Image shape: {}".format(images[0].shape))  # Shape of a single image
print("Image dtype: {}".format(images[0].dtype))  # Data type of the image array
print("Image type: {}".format(type(images[0])))   # Type of the image object

Loaded FashionMNIST dataset with 60000 samples.
Classes: {0: 'T-shirt/top', 1: 'Trouser', 2: 'Pullover', 3: 'Dress', 4: 'Coat', 5: 'Sandal', 6: 'Shirt', 7: 'Sneaker', 8: 'Bag', 9: 'Ankle boot'}
Image shape: (28, 28)
Image dtype: float64
Image type: <class 'numpy.ndarray'>

Problem 0a

从数据集的两列numpy.series:images和targets构造DataFrame

# TODO: Create a DataFrame with two columns: `image` and `label`

...
image_list=images.tolist()
label_list=labels.tolist()

df=pd.DataFrame({'image':image_list,'label':label_list})
df['image']=df['image'].apply(np.array)

# Print the shape and columns of the DataFrame
print("DataFrame shape:", df.shape)
print("DataFrame columns:", df.columns.tolist())
df.head()

DataFrame shape: (60000, 2)
DataFrame columns: ['image', 'label']

Problem 1a

计算数据集每个label出现的个数, 然后看每个label出现的个数是否相等

# TODO: Calculate the distribution of labels using `value_counts()``
# TODO: Compare the min and max values of `label_distribution` to determine if the dataset is balanced. 

label_distribution=df['label'].value_counts()
is_balanced=label_distribution.min()==label_distribution.max()

print(f"Label distribution:\n{label_distribution}")
print(f"Is the dataset balanced? {is_balanced}")

Label distribution:
label
Ankle boot     6000
T-shirt/top    6000
Dress          6000
Pullover       6000
Sneaker        6000
Sandal         6000
Trouser        6000
Shirt          6000
Coat           6000
Bag            6000
Name: count, dtype: int64
Is the dataset balanced? True

Problem 1b

用groupby()分组并且统计每个组的行数

# TODO: Group the rows in `df` according to the values in the `labels` column. Then, count the number of rows in each group.
label_distribution_groupby = df.groupby('label').size()

label_distribution_groupby

label
Ankle boot     6000
Bag            6000
Coat           6000
Dress          6000
Pullover       6000
Sandal         6000
Shirt          6000
Sneaker        6000
T-shirt/top    6000
Trouser        6000
dtype: int64

Problem 1c

对label列进行可视化

# Plotting library to use, default is matplotlib but plotly has more functionality
pd.options.plotting.backend = "plotly" 

# TODO: Plot a histogram of the labels in the DataFrame `df` using the DataFrame's built-in plotting functions (this should be 1 line)
df['label'].plot(kind="hist")

课程组还给了一个show_images函数, 他会打出图片和label

def show_images(images, max_images=40, ncols=5, labels = None, reshape=False):
    """Visualize a subset of images from the dataset.
    Args:
        images (np.ndarray or list): Array of images to visualize [img,row,col].
        max_images (int): Maximum number of images to display.
        ncols (int): Number of columns in the grid.
        labels (np.ndarray, optional): Labels for the images, used for facet titles.
    Returns:
        plotly.graph_objects.Figure: A Plotly figure object containing the images.
    """
    if isinstance(images, list):
        images = np.stack(images)
    n = min(images.shape[0], max_images) # Number of images to show
    px_height = 220 # Height of each image in pixels
    if reshape:
        images = images.reshape(images.shape[0], 28, 28)
    fig = px.imshow(images[:n, :, :], color_continuous_scale='gray_r', 
                    facet_col = 0, facet_col_wrap=ncols,
                    height = px_height * int(np.ceil(n/ncols)))
    fig.update_layout(coloraxis_showscale=False)
    fig.update_xaxes(showticklabels=False, showgrid=False)
    fig.update_yaxes(showticklabels=False, showgrid=False)
    if labels is not None:
        # Extract the facet number and replace with the label.
        fig.for_each_annotation(lambda a: a.update(text=labels[int(a.text.split("=")[-1])]))
    return fig

后面会用到

Problem 1d

用groupby分类, 然后在每个分类里面挑两张图, 打印出图和他的分类

# TODO: Get 2 sample images per class and plot them.
examples = df.groupby('label').head(2)

fig = show_images(examples["image"].tolist(), ncols=4, labels=examples["label"].tolist())
fig.show()

Problem 2

用reshape函数把一个m*n的图像变成mn*1的

df["image"] = df["image"].apply(lambda img: img.reshape(-1))

原始的每行的image元素都是28*28的, 现在每行都变成784*1的了

np.stack(df['image'].values).shape

(60000, 784)

这里说一下这个np.stack(arrays, axis, out)的用法, 以前每次都是AI写, 不太清楚其中原理

np.stack作用于一些数组, 这些数组必须有相同的shape, 沿着axis指定的轴拼起来, 这会加入第n+1个维度, 就像两本书叠起来, 显然要有一个z轴来标识这是第一本书还是第二本书

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
np.stack((a, b))
# array([[1, 2, 3],
#        [4, 5, 6]]) 2*3, 第0维出现一个2

np.stack((a, b), axis=-1)
# array([[1, 4],
#        [2, 5],
#        [3, 6]]) 3*2, 第1维出现一个3

简而言之就是axis是多少, 结果在第axis个维度上就会出现一个n, n是数组的个数

Problem 2a

调用sklearn.KMeans进行聚类

# TODO: Perform k-means clustering on the images (10 clusters to match the number of classes)
from sklearn.cluster import KMeans

df_sample = df.sample(n=1000, random_state=SEED)
kmeans=KMeans(n_clusters=10,random_state=SEED)
X=np.stack(df_sample['image'].to_numpy())
cluster_labels=kmeans.fit_predict(X)

kmeans_df=df_sample.assign(cluster=cluster_labels)

kmeans_df.head(3)

这里的np.stack就是把一堆图片点聚集起来, 现在这个kmeans_df有了一个新的列cluster, 代表聚类的结果, 比如说会把一些行聚类成1类, 2类....以此类推, 总共10个类

Problem 2b

评估KMeans聚类的效果

只需要对label(聚类的结果)groupby一下然后看每个聚类里面有多少种不同的true_label, 如果聚类的效果比较好的话, 应该每种聚类里面有尽可能少的true_label种类

# TODO: Create a stacked bar plot of the label counts per cluster.
cluster_label_counts = kmeans_df.groupby(['cluster', 'label']).size().unstack(fill_value=0)

cluster_label_counts.plot(
    kind='bar',
    title='Distribution of True Labels in Each K-means Cluster'
)

上面的数据操作有点烦人, 一不小心就会写错, 这其实是一个类似于pivot的操作, 经过groupby(["cluster","label"]).size()之后会有一个二级索引

    cluster  label
    0        0         50
             1         12
             2         5
    1        0         10
             2         80
    ...

unstack之后会把内层索引index变成列名, 这就创建了(cluster,label)到count的一一映射, 然后就可以画图了

cluster  label0  label1  label2  
0        50      12      5       
1        10      80      0       
...

这样一来count就成了没名字的元素了

Problem 2c

在每个Cluster里面随机抽出7个图片, 把他们打出来

# TODO: Plot 7 images from each cluster (use the show_images function, 10 rows, 7 columns)
Cluster_Head=kmeans_df.groupby('cluster').sample(n=7,random_state=SEED)

fig = show_images(Cluster_Head["image"].tolist(), ncols=7, labels=Cluster_Head["cluster"].tolist(),reshape=True,max_images=70)
fig.show()

# Cluster_Head

Problem 3

训练一个MLP分类器, 总共四个步骤

数据预处理
模型训练
绩效评估
可视化

先进行训练集和测试集的切分, 不能用train_test_split, 其实也很简单, 先sample出训练集, 那不在训练集里面的就是测试集

df_copy = df.copy()
train_df = df_copy.groupby('label').sample(frac=0.8, random_state=SEED)
test_df = df_copy[~df_copy.index.isin(train_df.index)]
print(f"Training set size: {len(train_df)}")
print(f"Test set size: {len(test_df)}")

Problem 3a

展平并归一化数据点, 启动训练并且绘制损失曲线

# flatten features into 1D arrays
X_train = np.stack(train_df['image'].values)
y_train = train_df['label'].values
X_test = np.stack(test_df['image'].values)
y_test = test_df['label'].values

print(f"X_train shape: {X_train.shape}\t y_train shape: {y_train.shape}")
print(f"X_test shape: {X_test.shape}\t y_test shape: {y_test.shape}")

# TODO: Train the model using the scaled traning data and plott the loss curve (remeber to normalize your data!)
# NOTE: Your model must be named `model`
def flatten(images):
    return images.reshape(images.shape[0], -1)
image_scalar=StandardScaler()
image_scalar.fit(flatten(X_train))
X_train_sc=image_scalar.transform(flatten(X_train))
X_test_sc=image_scalar.transform(flatten(X_test))

if load_saved_models and os.path.exists('classifier.joblib'):

    model = joblib.load('classifier.joblib')

else:

    model=MLPClassifier(
    hidden_layer_sizes=(100, 50),
    max_iter=200, tol=1e-4, random_state=SEED)

    model.fit(X_train_sc,y_train)
if save_models:
    joblib.dump(model, 'classifier.joblib')

loss_df=pd.DataFrame({'epoch':np.arange(1,len(model.loss_curve_)+1),'loss':model.loss_curve_})
loss_df.plot(x='epoch', y='loss', title="Training Error")

从原始的DataFrame里面拿数据出来的时候遵循先stack再reshape, 感觉很多时候脑子里面并没有对数据的具体形状有一个清晰的记忆, 尤其是在高维张量的情况下, 但是遵照这个去做一般不会出错

X_train shape: (48000, 784)	 y_train shape: (48000,)
X_test shape: (12000, 784)	 y_test shape: (12000,)

Problem 3b

产生预测结果和Evaluation Metrics, 这里主要就是一些API的调用, 主要涉及到Sklearn.model的属性怎么用的问题

# TODO: Add the columns listed above to `train_df` and `test_df`.
train_df = train_df.copy()
test_df = test_df.copy()

train_predict=model.predict(X_train_sc)
test_predict=model.predict(X_test_sc)

model.predict返回(n_samples,)的np数组, 所以他是可以直接和真实的label数组去进行比较的

train_correct=(train_predict==y_train)
test_correct=(test_predict==y_test)

分类模型本质上在预测概率, 对于每个class, 给出一个prob_class, 所以model.predict_prob会给出(n_samples,n_classes)的数组, 即对于每个样本的每个潜在分类给出一个概率

train_probs=model.predict_proba(X_train_sc)
test_probs=model.predict_proba(X_test_sc)

接下来还要得到每个样本的置信度, 实际上就是每个样本预测概率当中最大的那个概率, 通俗来说, 如果有十个候选类, 然后其中预测的最大概率是100%, 其他的类的概率都是0%, 那至少看起来是比较可信的

train_confidence=np.max(train_probs,axis=1)
test_confidence=np.max(test_probs,axis=1)

以上均得到(n_samples,)的数组

接下来把这些贴回到原来的df上去

train_probs_list=[list(probs) for probs in train_probs]
test_probs_list=[list(probs) for probs in test_probs]

train_df['predicted_label']=train_predict
train_df['correct']=train_correct
train_df['probs']=train_probs_list
train_df['confidence']=train_confidence

test_df['predicted_label']=test_predict
test_df['correct']=test_correct
test_df['probs']=test_probs_list
test_df['confidence']=test_confidence

最后算一下准确率就行了

print("--- Column Types ----")
for col in train_df.columns:
    val = train_df[col].iloc[0]
    print(f"{col}: {type(val)}")
print("-----------")


train_accuracy = train_correct.sum()/len(train_correct)
test_accuracy = test_correct.sum()/len(test_correct)

print(f"Training accuracy: {train_accuracy:.3f}")
print(f"Test accuracy: {test_accuracy:.3f}")

--- Column Types ----
image: <class 'numpy.ndarray'>
label: <class 'str'>
predicted_label: <class 'str'>
correct: <class 'numpy.bool'>
probs: <class 'list'>
confidence: <class 'numpy.float64'>
-----------
Training accuracy: 0.993
Test accuracy: 0.883

Problem 3c

分组查看准确率, 只要groupby('label')['correct'].mean()就可以得到每个类的准确率, 然后再把train和test两张表concat起来, 最后用pivot把它变成宽表即可

# TODO: Calculate train and test accuracy per class 
# TODO: Use class_accuracy to create a grouped bar chart of class accuracy for train and test

train_class_acc=train_df.groupby('label')['correct'].mean().reset_index()
train_class_acc['split']='train'

test_class_acc=test_df.groupby('label')['correct'].mean().reset_index()
test_class_acc['split']='test'

all_df=pd.concat([train_class_acc,test_class_acc])

class_accuracy=all_df[['split','label','correct']]
# print(class_accuracy)

class_accuracy_pivot=class_accuracy.pivot(index='label',columns='split',values='correct')
print(class_accuracy_pivot)

class_accuracy_pivot.plot(kind='bar',text_auto=True)

split            test     train
label                          
Ankle boot   0.967500  0.998333
Bag          0.966667  0.999792
Coat         0.807500  0.993333
Dress        0.888333  0.995208
Pullover     0.810000  0.990208
Sandal       0.945833  0.998958
Shirt        0.737500  0.990833
Sneaker      0.932500  0.996458
T-shirt/top  0.802500  0.967083
Trouser      0.975833  0.997917

Problem 3d

计算不同类别之间的预测效果, 首先要用np算一个10*10的confusion matrix, 行代表真实label, 列代表预测label, 每个元素代表这种预测的个数, 显然, 对角线上的元素求和就是总共预测对的数量

# Initialize confusion matrix with zeros
conf_matrix = np.zeros((len(class_names), len(class_names)), dtype=int)
class_to_idx = {class_name: idx for idx, class_name in enumerate(class_names)}

# Fill the confusion matrix by counting predictions and plot it as a heatmap
y_true=test_df['label'].values
y_pred=test_df['predicted_label'].values

true_indices = np.array([class_to_idx[label] for label in y_true])
pred_indices = np.array([class_to_idx[label] for label in y_pred])

np.add.at(conf_matrix, (true_indices, pred_indices), 1)

fig=px.imshow(conf_matrix,
                labels=dict(x="Predicted Label", y="True Label", color="Count"),
                x=class_names,
                y=class_names,
                )

fig.show()

实际上就是遍历一遍真实label和预测label, 注意要把label的名字转成int, 然后对矩阵的(int,int)加一就好了

conf_matrix

array([[ 963,    0,   23,   31,    1,    4,  171,    0,    7,    0],
       [   4, 1171,    4,   15,    1,    0,    4,    0,    1,    0],
       [  20,    2,  972,   14,  102,    0,   86,    0,    4,    0],
       [  27,    5,   19, 1066,   37,    1,   43,    0,    2,    0],
       [   7,    1,  112,   28,  969,    0,   79,    0,    4,    0],
       [   2,    0,    0,    3,    0, 1135,    0,   35,    5,   20],
       [ 118,    2,   96,   19,   71,    0,  885,    1,    8,    0],
       [   0,    0,    0,    0,    0,   24,    0, 1119,    1,   56],
       [   6,    1,    5,    2,    7,    4,   12,    2, 1160,    1],
       [   0,    0,    0,    0,    0,   10,    0,   27,    2, 1161]])

接下来需要计算FP, FN, TP指标, 直接在这个矩阵上操作即可, 只需要注意对角线上是对的, 其他的都是错的

# Calculate accuracy from confusion matrix
accuracy_from_matrix = np.trace(conf_matrix)/np.sum(conf_matrix)
print(f"\nAccuracy calculated from confusion matrix: {accuracy_from_matrix:.3f}")

# Calculate per-class metrics from confusion matrix
per_class_metrics = []
print("\nPer-class metrics from confusion matrix:")
for i, class_name in enumerate(class_names):
    true_positives = conf_matrix[i][i]
    false_positives = np.sum(conf_matrix[:,i])-conf_matrix[i][i]
    false_negatives = np.sum(conf_matrix[i,:])-conf_matrix[i][i]
    precision = true_positives / (true_positives + false_positives) if (true_positives + false_positives) > 0 else 0.0
    recall = true_positives / (true_positives + false_negatives) if (true_positives + false_negatives) > 0 else 0.0
    per_class_metrics.append({
        'class': class_name,
        'precision': precision,
        'recall': recall
    })
    
pd.DataFrame(per_class_metrics)

Accuracy calculated from confusion matrix: 0.883

Per-class metrics from confusion matrix:

Problem 3f

找出低置信度的预测并且plot出来

# TODO: Find the image with the lowest confidence by sorting the `confidence` column of `test_df`
least_confident = test_df.sort_values(by='confidence',ascending=True)
print("Image with lowest confidence:")
print(least_confident[['label', 'predicted_label', 'confidence', 'correct']][:3])

# Show image with lowest confidence and its predicted label
show_labels = [f"{label} (Pred: {predicted_label})" for label, predicted_label in zip(least_confident["label"].tolist(), least_confident["predicted_label"].tolist())]
fig = show_images(np.stack(least_confident["image"].tolist()), 8, ncols=4, labels=show_labels, reshape=True)
fig.show()

Image with lowest confidence:
         label predicted_label  confidence  correct
20252     Coat           Shirt    0.385341    False
23752  Sneaker      Ankle boot    0.395444    False
18719  Sneaker      Ankle boot    0.396654    False

可以看到, 在置信度非常低的情况下, 很多都预测错了

Problem 3g

找一些置信度很低但是预测对了的的Ankle boot类并画图, 非常容易, 先从test_df筛选出label == 'Ankle boot' & correct == True的, 然后用show_images画图就行了

# TODO: Visualize 10 images from the `test_set` whose true label is `Ankle boot` that the model correctly classified but with low confidence
test_df_boot = test_df[(test_df['label']=='Ankle boot')  &  (test_df['correct']==True)].sort_values(by='confidence',ascending=True).head(10)
# 将 labels 转换为列表，或者使用 .tolist()
show_labels = [f"{label} " for label in test_df_boot['label'].tolist()]
fig = show_images(np.stack(test_df_boot["image"].tolist()), 10, ncols=5, labels=show_labels, reshape=True)

fig.show()

Problem 3l

再找一些trouser类当中置信度很高但预测错误的画图, 和上面的一样, 照猫画虎即可

# TODO: Visualize 10 images from the `test_set` whose true label is `Trouser` that the model incorrectly classified as `Dress` with high confidence
test_df_trouser = test_df[(test_df['label']=='Trouser')  &  (test_df['correct']==False)].sort_values(by='confidence',ascending=False).head(10)
show_labels = [f"{label} " for label in test_df_trouser['label'].tolist()]
fig = show_images(np.stack(test_df_trouser["image"].tolist()), 10, ncols=5, labels=show_labels, reshape=True)

fig.show()

可以看到, 就算置信度很高, 也会有非常离谱的错误, 比如说把一个人分类成裤子, 从结果去理解过程, 我觉得他识别出的特征是两条裤腿

Problem 4

这一部分作业要求我们实现各种各样的图像增强, 我个人觉得这是比较难的部分, 因为这一段的numpy调用很多, 一不小心就不记得数据处理成什么样了

首先实现一些功能函数, 在高等代数里面学过(在此致敬电子科技大学何军华老师), 很多变换(线性)都可以通过矩阵乘法来表示, 所以实现这个应用变换的函数

def apply_transformation(image, T):
    # Input: A (N, 784) image vector and a (784, 784) transformation matrix
    # Output: A (N, 784) image vector
    transformed_flat = image @ T.T
    return transformed_flat.reshape(image.shape)

即将变换矩阵作用在图像上, 得到新图像

接下来是一个例子, 告诉我们如何实现图像的上下颠倒

def create_vertical_flip_matrix(height=28, width=28):
    """
    Returns a (height*width, height*width) matrix that vertically flips an image
    when applied to its flattened vector. Values are 0 or 1.
    """
    N = height * width  # Total number of pixels in the image
    T = np.zeros((N, N), dtype=int)  # Initialize the transformation matrix with zeros
    for i in range(height):  # Loop over each row
        for j in range(width):  # Loop over each column
            orig_idx = i * width + j  # Compute the flattened index for the original pixel
            flipped_i = height - 1 - i  # Compute the row index after vertical flip
            flipped_idx = flipped_i * width + j  # Compute the flattened index for the flipped pixel
            # Set the corresponding entry in the transformation matrix to 1
            # This means the pixel at (i, j) moves to (flipped_i, j)
            T[flipped_idx, orig_idx] = 1
    return T

def vertical_flip(image):
    T_flip = create_vertical_flip_matrix()
    return apply_transformation(image, T_flip)

上下颠倒只改变行, 不改变列, 实际上就是把第i行映射到height-1-i行

test_image = np.load("test_image.npy")

flipped_image = vertical_flip(test_image)
show_images(np.stack([test_image, flipped_image]), labels=['Original', 'Flipped'], reshape=True)

这里要说一下, 这个变换矩阵T的尺寸是N*N, 其中N=height*width, 可以理解为height*width个像素点, 每个像素点从i*width + j被映射到(height-1-i)*width + j

Problem 4a

要实现水平翻转, 依葫芦画瓢而已, 行不变, 列从j变成width - 1 - j, 所以每个元素i*width + j变到i*width + (width - 1 - j)

def create_horizontal_flip_matrix(height=28, width=28):
    """
    Returns a (height*width, height*width) matrix that horizontally flips an image
    when applied to its flattened vector. Values are 0 or 1.
    """
    N = height * width
    T = np.zeros((N, N), dtype=int)
    for i in range(height):
        for j in range(width):
            orig_idx = i * width + j
            flipped_j = width - 1 - j
            flipped_idx = i * width + flipped_j
            T[flipped_idx, orig_idx] = 1
    return T

    
def horizontal_flip(image):
    T_flip = create_horizontal_flip_matrix()
    return apply_transformation(image, T_flip)

flipped_image = horizontal_flip(test_image)

show_images(np.stack([test_image, flipped_image]), labels=['Original', 'Horizontal Flipped'], reshape=True)

Problem 4b

要求实现图像的Shift移动, 想象一下, 如果是水平移动, 那就是列j被映射到j + dx(先不考虑左右的符号问题), 所以每个像素点从i*width + j被映射到i*width + j + dx

这里有点小trick, 就是把x和y先flatten一下然后再移动

def create_shift_matrix(dx, dy, height=28, width=28):
    """
    Create a transformation matrix for shifting an image by dx pixels horizontally and dy pixels vertically.

    Args:
        dx (int): Number of pixels to shift horizontally.
        dy (int): Number of pixels to shift vertically.
        height (int): Height of the image.
        width (int): Width of the image.

    Returns:
        np.ndarray: A (height*width, height*width) transformation matrix for shifting.
    """
    N = height * width
    T = np.zeros((N, N))
    yy, xx = np.meshgrid(np.arange(height), np.arange(width), indexing='ij')
    yy_flat = yy.ravel()
    xx_flat = xx.ravel()

    y_shift = yy_flat + dy
    x_shift = xx_flat + dx

    valid = (
        (y_shift >= 0) & (y_shift < height) &
        (x_shift >= 0) & (x_shift < width)
    )

    src = yy_flat[valid] * width + xx_flat[valid]
    dst = y_shift[valid] * width + x_shift[valid]

    T[dst, src] = 1.0

    return T


def shift_image(image, dx, dy):
    """
    Shift an image by dx pixels horizontally and dy pixels vertically.

    Args:
        image (np.ndarray): Flattened image array of shape (height*width,).
        dx (int): Number of pixels to shift horizontally.
        dy (int): Number of pixels to shift vertically.

    Returns:
        np.ndarray: Shifted image as a flattened array.
    """
    T = create_shift_matrix(dx, dy)
    return apply_transformation(image, T)

shifted_right_image = shift_image(test_image, 5, 0)
shifted_left_image = shift_image(test_image, -5, 0)
shifted_up_image = shift_image(test_image, 0, -5)
shifted_down_image = shift_image(test_image, 0, 5)

all_images = np.stack([test_image, shifted_up_image, shifted_down_image, shifted_right_image, shifted_left_image])
plot_labels = ['Original', 'Shifted Up', 'Shifted Down', 'Shifted Right', 'Shifted Left']
show_images(all_images, labels=plot_labels, reshape=True)

注意这上面的xx yy都有height*weight个元素, 他们本质上是所有元素的列坐标和行坐标, 所以说被flatten以后, 比如说xx_flat + dx就是对所有的列坐标都加上dx, 然后后面再用形如i * width + height的方式来还原

Problem 4c

要求实现一个类似于卷积核的东西, 关键是要搞清楚src_idx和dst_idx之间的关系, 还要搞清楚这个矩阵T的含义, 矩阵T有weight*height行, 每一行对应原矩阵的一个点, 而每一行的weight*height个点对应的是原矩阵的这一个点对总共的weight*height个点产生的影响

举个例子:

0 1 2
3 4 5
6 7 8

原有矩阵是这样, 那么T的第一行就应该是:

[1/4, 1/4, 0,   1/4, 1/4, 0,   0,   0,   0  ],

这一行有9个点, 代表的就是对原矩阵9个点的影响, 我们假设kernel_size = 3, 那么生成的(0,0)处的元素应当是(0+1+3+4)/4, 也就是把原矩阵Flatten成行了之后和这一行(转置成列)做内积, 而那些没被影响的点, 在这一行里的数值自然就是0了

def create_blur_matrix(kernel_size=3, height=28, width=28):
    """
    Create a transformation matrix T that applies a uniform mean blur using a centered, odd-sized square sliding window.

    For each output pixel (i, j):
      1) Place a `kernel_size × kernel_size` window centered at (i, j).
      2) If the window is outside the image, then it will have fewer neighbors (only average the pixels that exist)

    Args:
        kernel_size (int): Size of the square kernel (must be odd).
        height (int): Height of the image.
        width (int): Width of the image.

    Returns:
        np.ndarray: A (height*width, height*width) transformation matrix for blurring.
    """
    N = height * width
    T = np.zeros((N, N))
    pad = kernel_size // 2
    kernel_area = kernel_size * kernel_size
    # Each pixel contributes equally to a kernel's neighborhood
    for i in range (height):
        for j in range (width):
            col_start=max(0,j-pad)
            row_start=max(0,i-pad)
            col_end=min(width-1,j+pad)
            row_end=min(height-1,i+pad)

            valid_count=(row_end-row_start+1)*(col_end-col_start+1)
            arg_value=1.0/valid_count

            dst_idx=i*width+j
            for r in range(row_start,row_end+1):
                for c in range(col_start,col_end+1):
                    src_idx=r*width+c
                    T[dst_idx,src_idx]=arg_value
    return T

可以看到一个i和j确定了唯一的dst_idx, 这个变量名没太起好, 这其实是原矩阵的(i,j)点, 那么在T当中他要占一整行, 而一整行每个被影响的点都是1/count, 影响不到的(窗口外的)自然就是0了

应用并画图

def blur_image(image, kernel_size=3):
    """
    Apply a blur transformation to a flattened image array or a batch of flattened images.

    Args:
        image (np.ndarray): Flattened image array of shape (height*width,) or batch of images (N, height*width).
        kernel_size (int): Size of the square kernel to use for blurring.

    Returns:
        np.ndarray: Blurred image(s) as a flattened array or batch of arrays.
    """
    T = create_blur_matrix(kernel_size)
    return apply_transformation(image, T)

blurred_1x1 = blur_image(test_image, kernel_size=1)
blurred_3x3 = blur_image(test_image, kernel_size=3)
blurred_5x5 = blur_image(test_image, kernel_size=5)

blurred_images = [test_image, blurred_1x1, blurred_3x3, blurred_5x5]
blurred_labels = ['Original', 'Blur 1x1', 'Blur 3x3', 'Blur 5x5']

show_images(blurred_images, labels=blurred_labels, reshape=True)

Problem 4d

要求实现图像旋转矩阵

这里的核心是根据$$\theta$$生成的旋转变换2*2矩阵

$$ R(\theta) = \begin{pmatrix} \cos\theta & -\sin\theta \ \sin\theta & \cos\theta \end{pmatrix} $$

根据一年级的解析几何课程知识(在此致敬电子科技大学余时伟老师), 点(x,y)如果围绕原点$$\theta$$旋转得到的(x',y')满足:

$$ \begin{pmatrix} x' \[4pt] y' \end{pmatrix}

\begin{pmatrix} \cos\theta & -\sin\theta \ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x \[4pt] y \end{pmatrix}

更一般的如果是围绕(x_0,y_0)进行旋转, 只要搞清楚本质上是向量再旋转就行, 所以满足:

$$ \begin{pmatrix} x' \[4pt] y' \end{pmatrix}

\begin{pmatrix} \cos\theta & -\sin\theta \ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x - x_0 \[4pt] y - y_0 \end{pmatrix} + \begin{pmatrix} x_0 \[4pt] y_0 \end{pmatrix} $$

def create_rotation_matrix(theta, height=28, width=28):
    """
    Create a transformation matrix for rotating an image by theta degrees.

    Args:
        theta (float): Angle of rotation in degrees.
        height (int): Height of the image.
        width (int): Width of the image.

    Returns:
        np.ndarray: A (height*width, height*width) transformation matrix for rotating.
    """
    # Convert theta from degrees to radians
    theta = np.deg2rad(theta)
    N = height * width
    T = np.zeros((N, N))
    center_i = (height - 1) / 2.0
    center_j = (width - 1) / 2.0

    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    
    for src_i in range(height):
        for src_j in range(width):

            src_i_centered = src_i - center_i
            src_j_centered = src_j - center_j
            
            # 旋转（逆时针旋转theta）
            dst_i_centered = src_i_centered * cos_theta - src_j_centered * sin_theta
            dst_j_centered = src_i_centered * sin_theta + src_j_centered * cos_theta
            
            # 平移回去
            dst_i = dst_i_centered + center_i
            dst_j = dst_j_centered + center_j
            
            # 取整并检查边界
            dst_i_int = int(np.round(dst_i))
            dst_j_int = int(np.round(dst_j))
            
            if 0 <= dst_i_int < height and 0 <= dst_j_int < width:
                src_idx = src_i * width + src_j
                dst_idx = dst_i_int * width + dst_j_int
                T[dst_idx, src_idx] = 1.0

    return T


def rotate_image(image, theta):
    """
    Apply a rotation transformation to a flattened image array or a batch of flattened images.

    Args:
        image (np.ndarray): Flattened image array of shape (height*width,) or batch of images (N, height*width).
        theta (float): Angle of rotation in degrees.

    Returns:
        np.ndarray: Rotated image(s) as a flattened array or batch of arrays.
    """
    T = create_rotation_matrix(theta)
    return apply_transformation(image, T)

# rotate with matrix
rotated_45 = rotate_image(test_image, 45) 
rotated_90 = rotate_image(test_image, 90)
rotated_200 = rotate_image(test_image, 200)
rotated_270 = rotate_image(test_image, 270)

# visualize original and 4 augmentations in plotly image grid
all_images = np.stack([test_image, rotated_45, rotated_90, rotated_200, rotated_270])
plot_labels = ['Original', 'Rotated (45°)', 'Rotated (90°)', 'Rotated (200°)', 'Rotated (270°)']
show_images(all_images, labels=plot_labels, reshape=True)

这里的映射关系就是:

(i, j) -> (dst_i, dst_j)

再构造出两边的src_idx和dst_idx即可, 注意这个T在apply_transformation的时候要转置的, 所以在填充T的时候永远是T[dst_idx, src_idx], 转置之后就变成T[src_idx, dst_idx]

Problem 4e

要求实现逆旋转变换和双线性插值, 首先考虑逆变换(x_0,y_0)到(x,y)

$$ \begin{pmatrix} x \[4pt] y \end{pmatrix}

\begin{pmatrix} \cos\theta & \sin\theta \ -\sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x' - x_0 \[4pt] y' - y_0 \end{pmatrix} + \begin{pmatrix} x_0 \[4pt] y_0 \end{pmatrix} $$

$$ x = \cos\theta,(x' - x_0) + \sin\theta,(y' - y_0) + x_0 $$

$$ y = -\sin\theta,(x' - x_0) + \cos\theta,(y' - y_0) + y_0 $$

def create_bilinear_rotation_matrix(theta, height=28, width=28):
    """
    Create a (height*width, height*width) matrix that applies bilinear interpolation
    for rotating a flattened image by theta degrees.
    Each row of the matrix gives the weights for the input pixels that contribute to each output pixel.

    Args:
        theta (float): Angle of rotation in degrees.
        height (int): Height of the image.
        width (int): Width of the image.

    Returns:
        np.ndarray: A (height*width, height*width) transformation matrix for rotating.
    """
    theta = np.deg2rad(theta)
    N = height * width
    T = np.zeros((N, N))
    center_i = (height - 1) / 2.0
    center_j = (width - 1) / 2.0
    cos_theta = np.cos(-theta)
    sin_theta = np.sin(-theta)

    for i in range(height):  # Loop over rows of the output image
        for j in range(width):  # Loop over columns of the output image
            output_idx = i * width + j

            # Translate output pixel to be relative to the center
            centered_j = j - center_j
            centered_i = i - center_i

            # Apply the inverse rotation to find the original pixel coordinates
            original_j_rotated = centered_j * cos_theta - centered_i * sin_theta
            original_i_rotated = centered_j * sin_theta + centered_i * cos_theta

            # Translate back to original image coordinates
            original_j = original_j_rotated + center_j
            original_i = original_i_rotated + center_i

            # Bilinear interpolation
            # Find the 4 nearest neighbor pixels
            j_floor = int(np.floor(original_j))
            i_floor = int(np.floor(original_i))
            j_ceil = int(np.ceil(original_j))
            i_ceil = int(np.ceil(original_i))

            # Calculate the fractional distances
            dj = original_j - j_floor
            di = original_i - i_floor

            # Define the weights
            w_tl = (1 - dj) * (1 - di)  # Top-left
            w_tr = dj * (1 - di)      # Top-right
            w_bl = (1 - dj) * di      # Bottom-left
            w_br = dj * di          # Bottom-right

            # Get the indices of the 4 neighbors, handling boundaries by setting weight to 0
            neighbors = []
            weights = []

            # Top-left neighbor
            if 0 <= i_floor < height and 0 <= j_floor < width:
                neighbors.append(i_floor * width + j_floor)
                weights.append(w_tl)

            # Top-right neighbor
            if 0 <= i_floor < height and 0 <= j_ceil < width:
                neighbors.append(i_floor * width + j_ceil)
                weights.append(w_tr)

            # Bottom-left neighbor
            if 0 <= i_ceil < height and 0 <= j_floor < width:
                neighbors.append(i_ceil * width + j_floor)
                weights.append(w_bl)

            # Bottom-right neighbor
            if 0 <= i_ceil < height and 0 <= j_ceil < width:
                neighbors.append(i_ceil * width + j_ceil)
                weights.append(w_br)

            # Assign the weighted sum to the output pixel in the transformation matrix
            for neighbor_idx, weight in zip(neighbors, weights):
                T[output_idx, neighbor_idx] = weight
    return T

首先这里遍历的其实是输出的image了, 不过尺寸还是一样的, 只要应用上面的矩阵变换就能得到原来的original_j和original_i, 也就是变换前的y和x

注意由于是三角函数的数值变换, 这个y和x不一定是整数, 所以我们找它周边的四个点做双线性插值

# Bilinear interpolation
# Find the 4 nearest neighbor pixels
j_floor = int(np.floor(original_j))
i_floor = int(np.floor(original_i))
j_ceil = int(np.ceil(original_j))
i_ceil = int(np.ceil(original_i))

# Calculate the fractional distances
dj = original_j - j_floor
di = original_i - i_floor

然后计算四个权重, 以竖着的距离为例, x到i_floor的距离是di, x到i_ceil的距离是1-di, 这是因为i_floor到i_ceil的距离是1

# Define the weights
w_tl = (1 - dj) * (1 - di)  # Top-left
w_tr = dj * (1 - di)      # Top-right
w_bl = (1 - dj) * di      # Bottom-left
w_br = dj * di          # Bottom-right

接下来, 遍历这四个邻居点并且添加权重, 注意在一次遍历当中output_idx是已经被固定了的, 然后在T的属于output_idx的那一行里面, 至多有四个不为0的weight(如果neighbor出格了就不要了)

# Top-left neighbor
if 0 <= i_floor < height and 0 <= j_floor < width:
    neighbors.append(i_floor * width + j_floor)
    weights.append(w_tl)

# Top-right neighbor
if 0 <= i_floor < height and 0 <= j_ceil < width:
    neighbors.append(i_floor * width + j_ceil)
    weights.append(w_tr)

# Bottom-left neighbor
if 0 <= i_ceil < height and 0 <= j_floor < width:
    neighbors.append(i_ceil * width + j_floor)
    weights.append(w_bl)

# Bottom-right neighbor
if 0 <= i_ceil < height and 0 <= j_ceil < width:
    neighbors.append(i_ceil * width + j_ceil)
    weights.append(w_br)

# Assign the weighted sum to the output pixel in the transformation matrix
for neighbor_idx, weight in zip(neighbors, weights):
    T[output_idx, neighbor_idx] = weight

应用更改

def rotate_image_bilinear(image, theta):
    """
    Rotate an image using bilinear interpolation.

    Args:
        image (np.ndarray): Flattened image array of shape (height*width,) or batch of images (N, height*width).
        theta (float): Angle of rotation in degrees.

    Returns:
        np.ndarray: Rotated image as a flattened array.
    """
    T = create_bilinear_rotation_matrix(theta)
    return apply_transformation(image, T)
    
# rotate with matrix
rotated = rotate_image(test_image, 45)
rotated_interpolated = rotate_image_bilinear(test_image, 45)

all_images = np.stack([test_image, rotated, rotated_interpolated])
plot_labels = ['Original',  'Rotated 45°', 'Rotated 45° (Bilinear)']
show_images(all_images, labels=plot_labels, reshape=True)

Problem 4f

实现复合变换, 假如我们有很多变换矩阵T在一个列表里面, 我们只需要遍历这个列表然后每次进行矩阵相乘就好了

def compose_transforms(*Ts):
  """
  Compose linear image transforms (each 784x784).
  Inputs:
    Ts: list of transformation matrices
  Returns:
    T_total: composition of all input transformations
  """
  T_total = Ts[0]
  for T in Ts[1:]:
    T_total = np.dot(T, T_total)  # 注意：T @ T_total，不是 T_total @ T
  return T_total

例如先旋转再模糊化处理:

def rotate_then_blur(image, theta, kernel_size):
  """
  Rotate an image by theta degrees (without bilinear interpolation) and then blur it with a kernel of size kernel_size.
  """


  Rotate_Matrix=create_rotation_matrix(theta,height=28,width=28)
  Blur_Matrix=create_blur_matrix(kernel_size=kernel_size,height=28,width=28)

  T=compose_transforms(Rotate_Matrix,Blur_Matrix)


  return apply_transformation(image,T)

以及先平移, 再旋转, 再模糊化处理

def shift_then_rotate_then_blur(image, dx, dy, theta, kernel_size):
  """
  Shift an image by (dx, dy), then rotate it by theta degrees (without bilinear interpolation), and then blur it with a kernel of size kernel_size.
  """
  Shift_Matrix=create_shift_matrix(dx,dy,height=28,width=28)
  Rotate_Matrix=create_rotation_matrix(theta,height=28,width=28)
  Blur_Matrix=create_blur_matrix(kernel_size,height=28,width=28)
  T=compose_transforms(Shift_Matrix,Rotate_Matrix,Blur_Matrix)
  return apply_transformation(image,T)

rotated_blurred_image = rotate_then_blur(test_image, 45, 3)
# shifted_rotated_blurred_image = shift_then_rotate_then_blur(test_image, 1, -4, 200, 5)
shifted_rotated_blurred_image = shift_then_rotate_then_blur(test_image, 5, 5, 45, 5)

all_images = np.stack([test_image, rotated_blurred_image, shifted_rotated_blurred_image])
plot_labels = ['Original', 'Rotated 45° and Blurred 2x2', 'Shifted 5, Rotated 45° and Blurred 2x2']
show_images(all_images, labels=plot_labels, reshape=True)

这个问题的测试没A过, 但是也没有看出来有什么问题

Problem 4h

写一个接口来应用这些更改, 不需要什么技巧, 直接把变换的类型和参数存起来, 然后for循环一个一个应用即可

# Test augmentation functions on a few examples
test_images = np.stack(test_df['image'])
test_labels = test_df['label']

shift_inputs = [(5, 0), (-5, 0), (0, 5), (0, -5)]
rotate_inputs = [45, 90, 200]
blur_inputs = [3, 5]
rotate_blur_inputs = [(45, 3), (90, 5)]
shift_rotate_blur_inputs = [((5, 0), 45, 3), ((-5, 0), 90, 5)]

augmented_data = []
# Randomly sample 100 datapoints from test_images
sample_idx = np.random.choice(len(test_images), 100, replace=False)
test_images_sample = test_images[sample_idx]
test_labels_sample = np.array(test_labels)[sample_idx]

# TODO: Apply the augmentation functions we just created (shift, blur, rotate w/ bilinear, rotate then blur, shift then rotate then blur) to every image from test_images_sample
# use the inputs defined above to apply the augmentations
# Save the augmented images in a new DataFrame aug_df
def add_augmented_sample(img, label, orig_idx, aug_name, aug_type, transform_fn):
    augmented_data.append({
        'original_idx': orig_idx,
        'augmentation': aug_name,
        'image': transform_fn(img).copy(),
        'label': label,
        'type': aug_type
    })


for orig_idx, img, label in zip(sample_idx, test_images_sample, test_labels_sample):
    for dx, dy in shift_inputs:
        add_augmented_sample(
            img, label, orig_idx,f'shift_{dx}_{dy}', 'shift',
            lambda x, dx=dx, dy=dy: shift_image(x, dx, dy)
        )
    for theta in rotate_inputs:
        add_augmented_sample(
            img, label, orig_idx,f'rotate_bilinear_{theta}', 'rotate',
            lambda x, theta=theta: rotate_image_bilinear(x, theta)
        )
    for k in blur_inputs:
        add_augmented_sample(
            img, label, orig_idx,f'blur_{k}x{k}', 'blur',
            lambda x, k=k: blur_image(x, kernel_size=k)
        )
    for theta, k in rotate_blur_inputs:
        add_augmented_sample(
            img, label, orig_idx,f'rotate_{theta}_blur_{k}', 'rotate_blur',
            lambda x, theta=theta, k=k: rotate_then_blur(x, theta, k)
        )
    for (shift_pair, theta, k) in shift_rotate_blur_inputs:
        dx, dy = shift_pair
        add_augmented_sample(
            img, label, orig_idx,f'shift_{dx}_{dy}_rotate_{theta}_blur_{k}', 'shift_rotate_blur',
            lambda x, dx=dx, dy=dy, theta=theta, k=k: shift_then_rotate_then_blur(x, dx, dy, theta, k)
        )


aug_df = pd.DataFrame(augmented_data)

# TODO: Select an image and visualize it with all the augmentations applied to it
example_image = test_images_sample[0]

example_variants = [('original', example_image)]
for dx, dy in shift_inputs:
    example_variants.append((f'shift_{dx}_{dy}', shift_image(example_image, dx, dy)))
for theta in rotate_inputs:
    example_variants.append((f'rot_bilin_{theta}', rotate_image_bilinear(example_image, theta)))
for k in blur_inputs:
    example_variants.append((f'blur_{k}x{k}', blur_image(example_image, kernel_size=k)))
for theta, k in rotate_blur_inputs:
    example_variants.append((f'rot_{theta}_blur_{k}', rotate_then_blur(example_image, theta, k)))
for (shift_pair, theta, k) in shift_rotate_blur_inputs:
    dx, dy = shift_pair
    example_variants.append((f'shift_{dx}_{dy}_rot_{theta}_blur_{k}',
                             shift_then_rotate_then_blur(example_image, dx, dy, theta, k)))

example_imgs = np.stack([img for _, img in example_variants])
example_labels = [name for name, _ in example_variants]

fig = show_images(example_imgs, ncols=4, labels=example_labels, reshape=True)
fig.show()

Problem 4i

把之前训练好的那个分类器在增强数据上测试, 评测一下效果, 也就是预处理一下数据, 把格式调成适配model的样子, 最后把效果画个图就好了

from sklearn.metrics import accuracy_score


def evaluate_augmented_data(aug_df, clf):
    records = []

    for aug_name, group in aug_df.groupby("augmentation"):
        # 将图像堆叠成 3D/4D 数组，再 reshape 成模型需要的形状
        images = np.stack(group["image"].to_numpy())
        n_samples = images.shape[0]
        X = images.reshape(n_samples, -1)

        # 若模型需要标准化，这里先拟合再变换（或对训练集拟合后在此仅 transform）
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)

        y_true = group["label"].to_numpy()
        y_pred = clf.predict(X_scaled)
        acc = accuracy_score(y_true, y_pred)

        records.append({
            "augmentation": aug_name,
            "accuracy": acc,
            "type": group["type"].iloc[0]
        })

    return pd.DataFrame(records)


aug_performance = evaluate_augmented_data(aug_df, model)
aug_performance = aug_performance.sort_values("accuracy", ascending=False).reset_index(drop=True)
aug_performance
# Visualize performance: sort by accuracy, color by augmentation type (blur, rotate, shift, none)

fig = px.bar(
    aug_performance,
    x="augmentation",
    y="accuracy",
    color="type",
    title="Classifier Accuracy on Augmented Data"
)
fig.update_layout(xaxis_title="Augmentation", yaxis_title="Accuracy")
fig.show()
<img src="/images/189_hw1/189_hw1_20.png" alt="label_plot" />

Stanford CS336 Assignment 1(Part 1) - 手搓BPE

Fri, 23 Jan 2026 00:00:00 GMT

import { Card, Button } from 'astro-pure/user'

从一份TinyStories的故事集到Transformer模型, 就像是从原始人一夜当中走进了现代

CS336 Assignment 1

这是我做过的第一个(也许也会是最后一个)几乎没什么Skeleton Code的Lab, 这门课的目标是让学习者彻底搞懂大模型的原理，并且从"Scratch"来从头构建大模型。

任务拆分开,大概有这么几点:

首先是工程架构部分:

bpe.py tokneizer.py: Byte-pair encoding(BPE) tokenizer, 在字节级别上实现一个分词器。
transformer.py: Transformer language model (LM), 实现Transformer的各模块并且组合成可实例化的Transformer类
transformer.py: The cross-entropy loss function and the AdamW optimizer, 实现AdamW优化器和损失函数
train_transformer.py: The training loop, with support for serializing and loading model and optimizer state, 实现训练循环

其次是跑训练和测试:

train_bpe_tinystories.py: Train a BPE tokenizer on the TinyStories dataset. 在Tinystories数据集上训练这个BPE分词器
tokenizer_experiments.py: Encode and Decode, 基于训练得到的词汇表和字节合并来对语料进行编码和解码
tokenizer_experiments.py: Run your trained tokenizer on the dataset to convert it into a sequence of integer IDs. 在数据集上应用分词器, 把文字数据集转化成整数序列
train_transformer.py: Train a Transformer LM on the TinyStories dataset. 利用分词器的结果训练Transformer模型 {/- TODO_generate_samples.py: Generate samples and evaluate perplexity using the trained Transformer LM. 用训练好的模型产生结果并且计算困惑度/} {/* - TODD_train_transformer_OpenWebText.py: Train models on OpenWebText. 在OpenWebText数据集上训练Transformer模型 */}

作业要求

torch.nn.parameter
torch.nn
torch.optim.Optimizer(作为基类)

这个作业主要还是让学习者实现算法和工程架构,至于底层那些并行计算之类的,交给PyTorch的开发者去做吧

评测框架

没有给出直接一键运行的测试,虽然测试的输入和输出ASSERT是实现好的,但是要自己去接这个测试接口,测试逻辑在./assignment1-basics/tests里面,测试接口在./assignment1-basics/tests/adapters.py里面

比如在transformer.py我们实现了Linear Layer,对应的测试接口在:

# Test For Implement of Linear Layer
from cs336_basics import transformer

def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    """
    Given the weights of a Linear layer, compute the transformation of a batched input.

    Args:
        in_dim (int): The size of the input dimension
        out_dim (int): The size of the output dimension
        weights (Float[Tensor, "d_out d_in"]): The linear weights to use
        in_features (Float[Tensor, "... d_in"]): The output tensor to apply the function to

    Returns:
        Float[Tensor, "... d_out"]: The transformed output of your linear module.
    """

    # raise NotImplementedError
    linear_module = transformer.Linear(in_features=d_in,out_features=d_out)
    linear_module.weight.data = weights
    return linear_module(in_features)

transformer.Linear是我们实现的Linear类,实现这个接口之后，运行

uv run pytest -k test_linear

就可以运行预设的测试了

下载数据

有四个数据集,两个Train两个Valid,分别对应TinyStories和OpenWebText_Result数据集

mkdir -p data
cd data

wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt
wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt

wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_train.txt.gz
gunzip owt_train.txt.gz
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_valid.txt.gz
gunzip owt_valid.txt.gz

cd ..

如果有网络问题，可以从https://hf-mirror.com/下载,替换下载命令里面的地址就行了

环境配置

本lab用uv做包管理器,我之前没用过完全不熟练,我是直接照着他文档里面的命令安装的,先安装uv

pip install uv

当运行某个py文件的时候,用uv命令运行,如果有依赖缺失会自动下载

uv run <python_file_path>

注意给uv也像pip一样配置一个国内源,不然安装一些大包可能要安装到明天早上去

# 推荐使用清华源
echo 'export UV_DEFAULT_INDEX="https://pypi.tuna.tsinghua.edu.cn/simple"'>> ~/.bashrc

# 或者用阿里源
# echo 'export UV_DEFAULT_INDEX="https://mirrors.aliyun.com/pypi/simple/"' >> ~/.bashrc

# 让配置立即生效
source ~/.bashrc
# 转载自https://zhuanlan.zhihu.com/p/1930714592423703026

BPE算法

文档中给出一个例子(stylized example),考虑以下的语料

low low low low low
lower lower widest widest widest
newest newest newest newest newest newest

我们有一个词语集合Vocabulary(代码里一般称为Vocab),可以理解为是一个词汇表，这个词汇表一开始内容很少,然后在训练的时候慢慢增长

假如我们通过whitespace来对语料进行分割,我们就可以得到frequency table:

{low: 5, lower: 2, widest: 3, newest: 6}

但我们要换一种更方便的数据结构来表示,比如用dict[tuple[bytes], int],那么这个表被表示成:

{(l,o,w):5,(l,o,w,e,r):2, ...}

接下来我们统计这个frequency table里面byte(char)对的两两组合,这就是个计数的工程,得到的结果是:

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6}

找到那些出现次数最多的组,这个例子里面是

{es:9, st:9}

挑选字典序更大的那个,这里是st,把所有的's' 't' 组合成st,所以这个frequency table就变成了

 {(l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,e,st): 3, (n,e,w,e,st): 6}

然后再把'st'添加到vocab里面去

再重复上面这个那个计数的步骤,这一次'e' 'st'是出现最频繁的,所以对他进行合并

如果我们重复合并到没有可以继续的了,我们的merges应该是:

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e',
'ne west', 'w i', 'wi d', 'wid est', 'low e', 'lowe r']

其中的每一项代表我们在那一次合并当中把两个什么东西(byte)合并了

当然实际上并不一定要合并到最后,我们可以指定合并的次数,比如例子中指定为6次,那么merges应该是:

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e']

这种情况下我们的vocab将会变成

[<|endoftext|>, [...256 BYTE CHARS], st, est, ow, low, west, ne]

其中前面两个是init vocab的时候就已经有的,然后每一次merge的时候会往vocab里面加一项

如此一来的话,词语newest就会被分词成为'ne' 'west', 换言之这个merge的过程就是把分词这件事情从细变粗的过程,最开始每一个词语都是一个一个字母分的,现在有一部分字母被聚合了

并行预分词 Parallelizing pre-tokenization

chunk和对每个chunk的处理

首先要对原有的文本进行预分词,以TinyStories数据集为例:

u don't have to be scared of the loud dog, I'll protect you". The mole felt so safe with the little girl. She was very kind and the mole soon came to trust her. He leaned against her and she kept him safe. The mole had found his best friend.
<|endoftext|>
Once upon a time, in a warm and sunny place, there was a big pit. A little boy named Tom liked to play near the pit. One day, Tom lost his red ball. He was very sad.
<|endoftext|>

They went back to the living room and cleaned up their toys. They decided to build something together. They made a big house with a garden and a fence. They put their cars and dolls inside. They were happy and proud of their work.
Mommy and Daddy came to see their house. They praised them and gave them a treat. It was a lemon cake. It was sour, but they liked it. They learned that sharing is caring, and that family is sweet.
<|endoftext|>


Lucy and the little girl played together happily. In the end, they both learnt an important lesson: be peaceful, kind, and understanding when faced with a conflict. And that is why Lucy and the little girl became great friends.
<|endoftext|>

At the end of the day, Tom and Max were tired. They had played all day and had lots of fun. They said goodbye to each other and went to their homes. Before going to sleep, they both did another easy stretch. Tom knew that tomorrow would be another happy morning.
<|endoftext|>

忽略数据的内容,注意到故事和故事之间是用<|endoftext|>进行分割的,所以我们把两个<|endoftext|>之间的内容成为一个chunk,首先要把chunk提取出来并且得到每一个chunk的frequency_dict

课程组给了一段现成的代码,在./assignment1-basics/cs336_basics/pretokenization_example.py里面,这个函数的作用是读取原始的Corpus,然后把chunk的边界,大概约等于上面说的<|endoftext|>的位置返回

def find_chunk_boundaries(
    file: BinaryIO,
    desired_num_chunks: int,
    split_special_token: bytes,
) -> list[int]:

注意这个分割字符(在这个例子里面是<|endoftext|>),是不唯一的,当然有可能某个corpus的是用<|Hello|>来区分段落的,只要数据类型是bytes就行了

既然需要并行,我们首先要写处理单个chunk的函数:构造如下

def pretokenize_chunk(args): # Deal with a chunk
    # bpe.py

    start,end,input_path,split_special_token = args
    # start: chunk的开始索引
    # end: chunk的结束索引
    # input_path: 文件目录
    # split_special_token: chunk内的分隔符
    
    with open(input_path,"rb") as f: # 打开文件
        f.seek(start) # 文件指针移动start偏移量(将文件读取指针移动到 start 指定的字节位置)
        chunk_bytes = f.read(end - start).decode("utf-8", errors="ignore") 
        # 读取start到end之间的数据
        
    split_pattern = "|".join(re.escape(token) for token in split_special_token)
    # 现在这个split_pattern能够正则匹配任意的split_special_token内的分隔符
    text_segments = re.split(f"({split_pattern})",chunk_bytes)
    # 对文本进行分割, 返回交替的结果, 形如[文本段1, 分隔符1, 文本段2, 分隔符2, ...]

    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    # GPT-2 风格的正则分词


    frequency_dict = defaultdict(int) # 初始化频率字典,注意用defaultdict来确保默认值为0,避免后面访问下标不存在的问题

    for segment in text_segments: # 遍历文本段
        if segment not in split_special_token: # 跳过分隔符,只处理实际文本

            for match in re.finditer(PAT,segment): # 在这个实际文本当中找所有的匹配项,返回迭代器
                pretoken = match.group()
                pretoken_bytes = pretoken.encode("utf-8") # 编码成UTF-8
                pretoken_bytes_tuple = tuple(bytes([b]) for b in pretoken_bytes)
                # 转化成字节tuple,例如("h", "e", "l", "l", "o") 而不是 b"hello"

                frequency_dict[pretoken_bytes_tuple] += 1
                # 更新频率


    return dict(frequency_dict)

要注意文本处理的架构(层次)是: 整个语料->单个chunk->chunk中的一个segment->segment当中的每个文本

如果split_special_token = [b"<|endoftext|>"],文本为:

Story 1 content here.<|endoftext|>Story 2 content here.<|endoftext|>

那么经过分割后的text_segments为

[
    "Story 1 content here.",  # 文本段
    "<|endoftext|>",          # 分隔符（被保留）
    "Story 2 content here.",  # 文本段
    "<|endoftext|>"           # 分隔符（被保留）
]

然后再遍历这个text_segments,得到的结果格式类似于(不一定准确,只是形式上类似,具体还要看这个PAT对每个segment的分割规则):

{
    (b'S', b't', b'o', b'r', b'y'): 2,  # "Story" 出现2次
    (b' ',): 6,                          # 单个空格出现6次
    (b'1',): 1,                          # "1" 出现1次  
    (b'2',): 1,                          # "2" 出现1次
    (b'c', b'o', b'n', b't', b'e', b'n', b't'): 2,  # "content" 出现2次
    (b'h', b'e', b'r', b'e'): 2,         # "here" 出现2次
    (b'.',): 2,                          # "." 出现2次
}

这就完成了对每个chunk的内部切割,接下来需要一个并行代码把每个chunk映射到这个函数上

同时并行处理多个chunk

我们应当用多进程同时处理多个chunk,每个chunk会独立的返回这个chunk内容的frequency_dict,然后把他们concat到一起去就得到了整个语料的频率字典

def load_and_chunk_file(
    input_path: str,
    desired_num_chunks: int,
    split_special_token: list[str],
    Debug=False): 
    # bpe.py

    with open(input_path,"rb") as f: # 打开文件
        num_processes = 16 # 设定并行的进程数量

        split_token_bytes = split_special_token[0].encode("utf-8") # 拿到分隔符并且把他encode成bytes

        boundaries = find_chunk_boundaries(file=f,desired_num_chunks=desired_num_chunks,split_special_token=split_token_bytes)
        # 调用给出的find_chunk_boundaries函数, 这里返回的就是原始语料里面每个分隔符的位置,在这个项目里面就是每个<|endoftext|>的位置


    with multiprocessing.Pool(processes=num_processes) as pool: # 并行处理
        
        chunk_args=[(start,end,input_path,split_special_token) for start,end in zip(boundaries[:-1],boundaries[1:])]
        # 四个参数, 注意boundaries的数据类型是list[int], 从boundaries[0,1]开始滑动取得每一组start和end
        # input_path和split_special_token都是不变的

        results=pool.map(pretokenize_chunk,chunk_args)
        # 传入参数并且获得并行传回的结果


        total_frequencies = merge_frequencies(results)
        # 把结果"concat"起来

        return total_frequencies

注意这个concat其实不是真的concat, 因为不同的chunk里面有可能出现相同的词语,比如在chunk1和chunk2的frequency_dict里面也许都有

{(b'S', b't', b'o', b'r', b'y'): 2}

那显然concat之后得到的应该是:

{(b'S', b't', b'o', b'r', b'y'): 4}

所以需要对并行传回来的结果进行遍历,然后合并同类项

# bpe.py
def merge_frequencies(frequency_dict): # Calculate the frequencies from each chunks and sum them together

    total_frequencies = defaultdict(int) # 初始化结果

    for every_frequency_dict in frequency_dict: 
    # 遍历每个chunk的frequency_dict
        for pretoken_bytes, count in every_frequency_dict.items():

            total_frequencies[pretoken_bytes] += count
            # 保证不同chunk之间的同样的pretoken_bytes的计数不重不漏

    return total_frequencies

初始化Vocab和Merges

从前面那个例子可以看到, 训练过程本质上就是更新Vocab和Merges这两个结果的过程, 所以先对他们进行初始化

# bpe.py
def initialize_vocab_and_merges(special_tokens):
    vocab = {}

    for i in range(256):
        vocab[i] = bytes([i]) # 初始的一些默认bytes

    for special_token in special_tokens:
        vocab[len(vocab)] = special_token.encode("utf-8") # 我们自己附加的special_tokens

    return vocab,[]

随着训练进行, Vocab和Merges会越来越大

初始化pair计数

我们现在只得到了每个词语的计数, 但是如例子里面, 我们要对每个词语遍历拆出pair, 去对pair进行计数, 比如说输入是

{(b'S', b't', b'o', b'r', b'y'): 4}

那输出大概是

{'St':4, 'to':4, 'or':4, 'ry':4}

注意这里已经不分chunk了, 所以每次得到的结果就是全局的更新量了

#bpe.py
def get_initial_pair_frequencies(frequency_dict,Debug=False):

    pair_freq = defaultdict(int)
    pair_to_tokens = defaultdict(set)
    # 初始化

    for pretoken_bytes, count in frequency_dict.items():
        # 遍历词频字典, 获取每个"单词"和他的计数


        for i in range(len(pretoken_bytes)-1): # 遍历这个单词的所有相邻字符
            pair = ((pretoken_bytes[i],), (pretoken_bytes[i+1],)) # 取得相邻的pair
            pair_freq[pair] = pair_freq.get(pair,0) + count # 更新这个pair的计数
            pair_to_tokens[pair].add(pretoken_bytes) # 记住某一个pair出现在哪个单词里面


    return pair_freq, pair_to_tokens

找到最频繁出现的`pair`

假如有类似于

{'St':4, 'to':4, 'or':4, 'ry':4}

的pair的频率字典, 那么需要一个函数获取里面出现次数最多的那个pair

# bpe.py
def find_best_pair(pair_frequencies):
    if not pair_frequencies:
        return None

    best_pair = tuple()
    max_freq = -1

    for pair,freq in pair_frequencies.items(): # 遍历一遍即可, 追踪最大的freq的pair
        if freq > max_freq:
            max_freq = freq
            best_pair = pair

        elif freq == max_freq:
            best_pair = max(best_pair,pair)

    return best_pair

Merge操作

接下来这个是我觉得BPE里面最难的实现, 首先回顾一下我们现在有了什么:

frequency_dict: dict(词语:频率)

pair_frequencies: dict(字符对:频率)

pair_to_tokens: dict(字符,set(词语)), 这个来自get_initial_pair_frequencies函数, 标注哪些词语含有这个字符

best_pair: tuple(字符1, 字符2)

还有一个比较麻烦的地方在于这一步有性能要求, 如果太暴力的话可能会导致测试过不去, 一个简单的想法是遍历frequency_dict, 然后找到所有含有best_pair的词语对他们进行修改, 但是这样是无法通过测试的

想一个取巧的办法, 既然我们有pair_to_tokens, 我们可以先找出含有best_pair的词语然后对他们进行遍历, 这就不需要遍历所有的词语了

# bpe.py
def merge_pair(frequency_dict, pair_frequencies, pair_to_tokens, best_pair,Debug=False):

    byte1_tuple, byte2_tuple = best_pair # 拆开两个字符, 例如(b'S',)和(b't',)
    merged_byte = byte1_tuple[0] + byte2_tuple[0] # 转换数据结构

    affected_tokens = pair_to_tokens.get(best_pair,set()).copy()
    # 找到那些含有这个pair的词语

    if best_pair in pair_to_tokens:
        # 在被merge之后, 所有的词语当中应该不再含有这个pair, 此等价于这个pair不再属于任何词语, 所以从pair_to_tokens当中删除这个pair
        del pair_to_tokens[best_pair]

    for pretoken in affected_tokens: # 遍历含有这个pair的词语
        count = frequency_dict[pretoken] # 旧的词语计数, 注意合并后词语计数是会变化的, 比如旧的词语是(b'S', b't', b'o', b'r', b'y')
        # 如果St被合并之后, 这个旧的词语应该是不存在了, 产生了一个新的词语:(b'St', b'o', b'r', b'y'), 这个新的词语(其实不见得是新的), 可能以前有
        # 的计数要加上老的词语的计数才行

        frequency_dict[pretoken] -= count # 旧的词语计数减少
        if frequency_dict[pretoken] <= 0:
            del frequency_dict[pretoken] # 删除旧的词语

        
        new_pretoken_list = [] # 准备生成新词语
        i = 0
        while i < len(pretoken): # 用类似滑动窗口的形式
            if (i < len(pretoken) - 1 and 
                (pretoken[i],) == byte1_tuple and 
                (pretoken[i+1],) == byte2_tuple):
                new_pretoken_list.append(merged_byte) # 如果检测到best_pair那一位就直接把best_pair一起append
                i += 2
            else:
                new_pretoken_list.append(pretoken[i]) # 否则只append一个byte
                i += 1

        new_pretoken_tuple = tuple(new_pretoken_list) # 生成新词语
        frequency_dict[new_pretoken_tuple] += count # 加上老词语的计数

还要修改pair_frequencies和pair_to_tokens, 因为老的pair已经没有了并且由于产生了新的词语,会产生新的pair, 接着上面的函数代码:

# bpe.py
        if new_pretoken_tuple != pretoken: # 如果新词语不等于老词语(真的有合并发生)

            for i in range(len(pretoken) - 1):
                old_pair = ((pretoken[i],), (pretoken[i+1],))
                pair_frequencies[old_pair] -= count
                if pair_frequencies[old_pair] <= 0: # 修改pair_frequencies
                    del pair_frequencies[old_pair]

                if old_pair in pair_to_tokens:
                    pair_to_tokens[old_pair].discard(pretoken)
                    if not pair_to_tokens[old_pair]: # 修改pair_to_tokens
                        del pair_to_tokens[old_pair]
            

            for i in range(len(new_pretoken_tuple) - 1): # 新词语本身会引入新的pair, 所以更新pair_frequencies和pair_to_tokens
                new_pair = ((new_pretoken_tuple[i],), (new_pretoken_tuple[i+1],))
                pair_frequencies[new_pair] += count

                pair_to_tokens[new_pair].add(new_pretoken_tuple)

    return frequency_dict, pair_frequencies, pair_to_tokens

训练主函数

现在我们来写训练循环, 架构非常简单, 分为以下几个步骤

数据准备：

加载训练语料库文件
按分隔符（如 <|endoftext|>）将文件分割成多个数据块

预处理：

对每个数据块进行预分词，统计所有预分词的频率
初始化词频表 frequency_dict

字节对统计：

计算所有相邻字节对的出现频率
建立字节对频率表和映射关系

while 继续训练条件: 步骤1: 找到最优合并对 best_pair = 找到频率最高的字节对(pair_frequencies)

步骤2: 执行合并
merge_pair(frequency_dict, pair_frequencies, pair_to_tokens, best_pair)

步骤3: 更新模型状态
更新合并记录(merges列表)
更新词汇表(vocab集合)

BPE训练流程
├── 初始化阶段
│   ├── 加载并分块文件
│   ├── 初始化词频表  
│   └── 初始化字节对频率
├── 训练迭代 (循环开始)
│   ├── 选择最佳字节对
│   ├── 合并操作
│   │   ├── 更新词频统计
│   │   ├── 重建受影响的预分词
│   │   └── 更新字节对映射
│   ├── 记录合并操作
│   └── 扩展词汇表
└── 训练结束

# bpe.py
def train_bpe(input_path:str, vocab_size:int, special_tokens:list[str], Debug=False):

    frequency_dict = load_and_chunk_file(input_path, desired_num_chunks=4, split_special_token=special_tokens)
    # 生成原始词频表

    

    vocab, merges = initialize_vocab_and_merges(special_tokens)
    # 生成原始词汇表和空的Merges

    pair_frequencies, pair_to_tokens = get_initial_pair_frequencies(frequency_dict)
    # 生成初始的字符对的频率表
    

    while len(vocab) < vocab_size:
        # 循环条件: 未达到预设的vocab_size

        best_pair = find_best_pair(pair_frequencies) # 获取best_pair

        if not best_pair: 
            break

        frequency_dict, pair_frequencies, pair_to_tokens = merge_pair(frequency_dict, pair_frequencies, pair_to_tokens, best_pair) 
        # merge这个best_pair

        merges.append((best_pair[0][0],best_pair[1][0])) # 更新merges

        new_token = best_pair[0][0] + best_pair[1][0] # 更新vocab
        vocab[len(vocab)] = new_token

    return vocab,merges

评测接口

非常容易, 只要修改adapters.py的框架代码就行了:

# adapters.py

from cs336_basics import bpe

def run_train_bpe(
    input_path: str | os.PathLike,
    vocab_size: int,
    special_tokens: list[str],
    **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:

    input_path_str = str(input_path)
    vocab,merges = bpe.train_bpe(input_path=input_path_str,vocab_size=vocab_size,special_tokens=special_tokens,Debug=False)
    
    return vocab, merges

这里不需要手动指定路径, 测试自己存了一些demo文件去进行训练并和ref进行比对来判断正误, 启动测试的代码是:

root@autodl-container-8d994fbd73-e5baa69e:~/autodl-tmp/Stanford_CS336/assignment1-basics# uv run pytest tests/test_train_bpe.py

稍微注意一下启动时候的路径, 测试通过如下:

(base) root@autodl-container-8d994fbd73-e5baa69e:~/autodl-tmp/Stanford_CS336/assignment1-basics# uv run pytest tests/test_train_bpe.py 
============================================================================ test session starts =============================================================================
platform linux -- Python 3.12.3, pytest-8.4.1, pluggy-1.6.0
rootdir: /root/autodl-tmp/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 3 items                                                                                                                                                                                                                                   

tests/test_train_bpe.py::test_train_bpe_speed PASSED
tests/test_train_bpe.py::test_train_bpe PASSED
tests/test_train_bpe.py::test_train_bpe_special_tokens PASSED

================================================================================================================= 3 passed in 3.04s =================================================================================================================

在TinyStories和OpenWebText数据集上训练BPE分词器

这里只需要自己写一个脚本调用之前实现的训练过程就行了, 唯一要注意的就是Merges和Vocab持久化时候的格式问题

Vocab的输出格式是:

vocab = {0: b'hello', 1: b'world', ...}

需要转化为

vocab_unicode = {'hello': 0, 'world': 1, ...}

# train_bpe_tinystories.py

from cs336_basics.bpe import train_bpe
from loguru import logger
from tests.common import gpt2_bytes_to_unicode
import json
from pathlib import Path

def save_vocab_merge(vocab, merges, output_path='./../../TinyStories_Result'):
    output_dir = Path(output_path)
    output_dir.mkdir(exist_ok=True)

    byte_to_unicode = gpt2_bytes_to_unicode()

    # Vocab:{id:bytes} -> {unicode_string:id}
    vocab_unicode = {}

    for token_id,token_bytes in vocab.items():
        unicode_chars = [byte_to_unicode[b] for b in token_bytes]
        unicode_string = ''.join(unicode_chars)
        vocab_unicode[unicode_string] = token_id

    vocab_path = output_dir/"vocab.json"

    with open(vocab_path, 'w', encoding='utf-8') as f:
        json.dump(vocab_unicode, f, ensure_ascii=False, indent=2)

    merge_path = output_dir/"merges.txt"
    
    # Merge:[tuple(bytes,bytes)]
    with open(merge_path,'w',encoding='utf-8') as f:
        for merge_pair in merges:
            token1_bytes, token2_bytes = merge_pair
            token1_unicode = ''.join(byte_to_unicode[b] for b in token1_bytes)
            token2_unicode = ''.join(byte_to_unicode[b] for b in token2_bytes)
            f.write(f"{token1_unicode} {token2_unicode}\n")

    return vocab_path,merge_path

实际上就是把原始格式转化成人类可读的Unicode格式就可以了

训练接口函数

# train_bpe_tinystories.py
def train_bpe_tinystories():
    
    input_path = "./../../data/TinyStoriesV2-GPT4-train.txt"
    special_tokens = ["<|endoftext|>"]
    vocab_size = 10000

    vocab, merges = train_bpe(input_path=input_path,vocab_size=vocab_size,special_tokens=special_tokens,Debug=False)



    vocab_path, merges_path = save_vocab_merge(vocab, merges)



if __name__ == "__main__":
    train_bpe_tinystories()

另一个数据集只需要修改一下路径就好了, 不再赘述

# train_bpe_expts_owt.py
from cs336_basics.bpe import train_bpe
from loguru import logger
from tests.common import gpt2_bytes_to_unicode
import json
from pathlib import Path

# uv run python train_bpe_expts_owt.py 

def save_vocab_merge(vocab, merges, output_path='./../../OpenWebText_Result'):
    output_dir = Path(output_path)
    output_dir.mkdir(exist_ok=True)

    byte_to_unicode = gpt2_bytes_to_unicode()

    # Vocab:{id:bytes} -> {unicode_string:id}
    vocab_unicode = {}

    for token_id,token_bytes in vocab.items():
        unicode_chars = [byte_to_unicode[b] for b in token_bytes]
        unicode_string = ''.join(unicode_chars)
        vocab_unicode[unicode_string] = token_id

    vocab_path = output_dir/"vocab.json"

    with open(vocab_path, 'w', encoding='utf-8') as f:
        json.dump(vocab_unicode, f, ensure_ascii=False, indent=2)

    merge_path = output_dir/"merges.txt"
    
    # Merge:[tuple(bytes,bytes)]
    with open(merge_path,'w',encoding='utf-8') as f:
        for merge_pair in merges:
            token1_bytes, token2_bytes = merge_pair
            token1_unicode = ''.join(byte_to_unicode[b] for b in token1_bytes)
            token2_unicode = ''.join(byte_to_unicode[b] for b in token2_bytes)
            f.write(f"{token1_unicode} {token2_unicode}\n")

    return vocab_path,merge_path

def train_bpe_expts_owt():
    
    input_path = "./../../data/owt_train.txt"
    special_tokens = ["<|endoftext|>"]
    vocab_size = 32000

    vocab, merges = train_bpe(input_path=input_path,vocab_size=vocab_size,special_tokens=special_tokens,Debug=False)




    vocab_path, merges_path = save_vocab_merge(vocab, merges)



if __name__ == "__main__":
    train_bpe_expts_owt()

最后得到的merges.txt和vocab.json类似于

Ġ t
h e
Ġ a
Ġ s
Ġ w
n d
Ġt he
e d
Ġ b
Ġt o

{
  "Ā": 0,
  "ā": 1,
  "Ă": 2,
  "ă": 3,
  "Ą": 4,
  "ą": 5,
  "Ć": 6,
  "ć": 7,
  "Ĉ": 8,
  "ĉ": 9,
  "Ċ": 10,
}

分词器: 编码与解码

当Vocab和Merges被训练完成后, 我们可以用他们然后对给定的语料进行解码/编码

编码

本质上就是把语料先预分词, 然后应用已有的Merges, 最后再编码成整数序列, 举例如下:

Vocab = {0: b' ', 1: b'a', 2:b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b'at'}
Merges = [(b't', b'h'), (b' ', b'c'), (b' ', 'a'), (b'th', b'e'), (b' a', b't')]

对'the cat ate'进行编码的步骤如下:

首先预分词用空格来进行token划分,得到

['the', ' cat', ' ate']

对于第一个token:'the', 他的表示方式是:

[b't', b'h', b'e']

应用两次Merges中的合并就变成:

[b'the']

然后转化成整数序列就是[9], 其余两个token也一样, 最后得到的整数序列是:

[9, 7, 1, 5, 10, 3]

解码

解码过程相对简单, 对于输入的整数序列, 只需要一个个的查找在Vocab中对应的词并且concat起来就好了

`Tokenizer`类的实现(遵循实验文档的接口架构)

# tokenizer.py
class tokenizer():

    def __init__(self,vocab:dict[int,bytes],merges:list[tuple[bytes,bytes]],special_tokens:list[str]=None):
        self.vocab = vocab
        self.merges = merges
        self.special_tokens = special_tokens

        self.vocab_reverse = {token_bytes:token_id for token_id,token_bytes in self.vocab.items()} # 反向映射方便查找
        self.PAT = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")

先实现类方法apply_bpe_merges, 传入一个token, 查找Merges看是否有可以应用在这个token上的合并操作, 若有就应用后返回

# tokenizer.py
    def apply_bpe_merges(self,token_parts)->list[bytes]:

        current_parts = token_parts.copy()

        for merge_pair in self.merges: # 遍历合并表
            byte1,byte2 = merge_pair # 待合并的字节对

            i=0
            while i < len(current_parts) - 1: # 遍历token看是否有和合并的字节对一样的字节对
                if current_parts[i]==byte1 and current_parts[i+1]==byte2:
                    merged = byte1 + byte2
                    current_parts[i] = merged
                    del current_parts[i+1] # 执行合并操作并且删除原来的第二个字节

                else:
                    i += 1

        return current_parts

实现一个一般的编码函数, 单纯的通过PAT去匹配, 然后对各PAT分割出的部分应用Merge, 最后再通过self.vocab_reverse转化成整数序列完成编码

# tokenizer.py
    def encode_normal_part(self,text:str)->list[int]:


        pre_tokens = [] # 分割成的很多token
        for match in re.finditer(self.PAT,text):
            pre_tokens.append(match.group()) # 根据PAT来分割

        result = []

        for pre_token in pre_tokens: # 处理每个预分词
            pre_token_bytes = pre_token.encode("utf-8") # 生成字节序列
            token_parts = [bytes([b]) for b in pre_token_bytes] # 生成单个字节的列表, 例如"Hello" → b'Hello' → [b'H', b'e', b'l', b'l', b'o']

            merged_parts = self.apply_bpe_merges(token_parts) # 应用Merge

            for part in merged_parts: # 遍历Merge后的字节序列, 查表把每个字节序列转化为token id, 同时忽略那些不再词汇表里的字节序列
                token_id = self.vocab_reverse.get(part,None)
                if token_id is not None:
                    result.append(token_id)
        return result


        # 1. 预分词：["Hello"]

        # 2. 字节化：[b'H', b'e', b'l', b'l', b'o']

        # 3. BPE合并（假设有 'l','l' 合并）：[b'H', b'e', b'll', b'o']

        # 4. ID查找：
        # - b'H' → 假设ID为 72
        # - b'e' → 假设ID为 101  
        # - b'll' → 假设ID为 200
        # - b'o' → 假设ID为 111

        # 5. 输出：[72, 101, 200, 111]

考虑到有可能有self.special_tokens, 所以要在上面这个函数上实现一个更一般的encode函数


    def encode(self,text:str)->list[int]:

        if self.special_tokens: # 特殊token模式
            sorted_specials = sorted(self.special_tokens, key=len, reverse=True)
            special_pattern = "|".join(re.escape(token) for token in sorted_specials)
            # 在特殊token的地方切开
            # [普通文本1, 特殊token1, 普通文本2, 特殊token2, ...]

            parts = re.split(f"({special_pattern})", text)

            result = []

            for part in parts: 
                if not part:
                    continue

                elif part in self.special_tokens: # 特殊token部分不需要走Merges逻辑, 直接查表就行
                    special_bytes = part.encode("utf-8")
                    token_id = self.vocab_reverse.get(special_bytes,None)
                    if token_id is not None:
                        result.append(token_id)

                else: # 走之前的逻辑, 先应用Merges然后再查表
                    part_result = self.encode_normal_part(part)
                    result.extend(part_result)
            return result

        else: # 无特殊token, 直接走原有逻辑
            return self.encode_normal_part(text)

文档要求我们以流式处理和惰性求值的方式来encode, 非常简单, 实现如下:

# tokenizer.py
    def encode_iterable(self,iterable:Iterable[str])->Iterable[int]:
        for text in iterable:
            token_ids = self.encode(text)

            for token_id in token_ids:
                yield token_id

解码部分更容易, 就是一个查表然后join的函数而已

# tokenizer.py
    def decode(self,ids:list[int])->str:
        byte_sequences = []
        for token_id in ids:
            if token_id in self.vocab.keys():
                byte_sequences.append(self.vocab[token_id])


        combined_bytes = b''.join(byte_sequences)

        text = combined_bytes.decode('utf-8',errors='replace')

        return text

实验文档还要求我们实现一个类的构造器, 其实很容易, 只要读取那些被持久化了的Merges和Vocab,然后把他转化成持久化前的数据类型就行了, 相当于是前面训练中那个save_vocab_merge的逆向

# tokenizer.py
    @classmethod
    def from_files(cls,vocab_filepath:str,merges_filepath:str,special_tokens:list[str]=None):
        with open(file=vocab_filepath,mode='r',encoding='utf-8') as f:
            vocab_unicode = json.load(f)

        vocab = {}
        for unicode_str,token_id in vocab_unicode.items():
            vocab[token_id] = unicode_str.encode('utf-8')


        merges = []
        with open(file=merges_filepath,mode='r',encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    token1_str, token2_str = line.strip().split()

                    merges.append((token1_str.encode('utf-8'),token2_str.encode('utf-8')))

        return cls(vocab,merges,special_tokens)

评测接口

和之前一样, 修改adapters.py即可

# adapters.py
from cs336_basics import tokenizer

def get_tokenizer(
    vocab: dict[int, bytes],
    merges: list[tuple[bytes, bytes]],
    special_tokens: list[str] | None = None,
) -> Any:

    return tokenizer.tokenizer(vocab=vocab,merges=merges,special_tokens=special_tokens)

启动测试:

uv run pytest tests/test_tokenizer.py

测试结果:

(base) zyli@lab:~/Stanford_CS336/assignment1-basics$ uv run pytest tests/test_tokenizer.py
======================================================================== test session starts =========================================================================
platform linux -- Python 3.13.5, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/zyli/Stanford_CS336/assignment1-basics
configfile: pyproject.toml
plugins: jaxtyping-0.3.2
collected 25 items                                                                                                                                                   

tests/test_tokenizer.py::test_roundtrip_empty PASSED
tests/test_tokenizer.py::test_empty_matches_tiktoken PASSED
tests/test_tokenizer.py::test_roundtrip_single_character PASSED
tests/test_tokenizer.py::test_single_character_matches_tiktoken PASSED
tests/test_tokenizer.py::test_roundtrip_single_unicode_character PASSED
tests/test_tokenizer.py::test_single_unicode_character_matches_tiktoken PASSED
tests/test_tokenizer.py::test_roundtrip_ascii_string PASSED
tests/test_tokenizer.py::test_ascii_string_matches_tiktoken PASSED
tests/test_tokenizer.py::test_roundtrip_unicode_string PASSED
tests/test_tokenizer.py::test_unicode_string_matches_tiktoken PASSED
tests/test_tokenizer.py::test_roundtrip_unicode_string_with_special_tokens PASSED
tests/test_tokenizer.py::test_unicode_string_with_special_tokens_matches_tiktoken PASSED
tests/test_tokenizer.py::test_overlapping_special_tokens PASSED
tests/test_tokenizer.py::test_address_roundtrip PASSED
tests/test_tokenizer.py::test_address_matches_tiktoken PASSED
tests/test_tokenizer.py::test_german_roundtrip PASSED
tests/test_tokenizer.py::test_german_matches_tiktoken PASSED
tests/test_tokenizer.py::test_tinystories_sample_roundtrip PASSED
tests/test_tokenizer.py::test_tinystories_matches_tiktoken PASSED
tests/test_tokenizer.py::test_encode_special_token_trailing_newlines PASSED
tests/test_tokenizer.py::test_encode_special_token_double_newline_non_whitespace PASSED
tests/test_tokenizer.py::test_encode_iterable_tinystories_sample_roundtrip PASSED
tests/test_tokenizer.py::test_encode_iterable_tinystories_matches_tiktoken PASSED
tests/test_tokenizer.py::test_encode_iterable_memory_usage PASSED
tests/test_tokenizer.py::test_encode_memory_usage XFAIL (Tokenizer.encode is expected to take more memory than allotted (1MB).)

========================================================================== warnings summary ==========================================================================
tests/adapters.py:352
  /home/zyli/Stanford_CS336/assignment1-basics/tests/adapters.py:352: SyntaxWarning: invalid escape sequence '\T'
    rope_theta (float): The RoPE $\Theta$ parameter.

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================================================= 24 passed, 1 xfailed, 1 warning in 1696.82s (0:28:16) ========================================================

这个测试耗时在现有的复杂度下非常长, 而且也有一个内存限制测试是预计不通过的(测试文件里面写了这个测试本来就预计不通过), 在实验文档里有写可以通过Cpp或Rust实现来显著提升速度

To test your BPE training function against our provided tests, you will first need to implement the
test adapter at [adapters.run_train_bpe]. Then, run uv run pytest tests/test_train_bpe.py.
Your implementation should be able to pass all tests. Optionally (this could be a large time-investment),
you can implement the key parts of your training method using some systems language, for instance
C++ (consider cppyy for this) or Rust (using PyO3). If you do this, be aware of which operations
require copying vs reading directly from Python memory, and make sure to leave build instructions, or
make sure it builds using only pyproject.toml. Also note that the GPT-2 regex is not well-supported
in most regex engines and will be too slow in most that do. We have verified that Oniguruma is
reasonably fast and supports negative lookahead, but the regex package in Python is, if anything,
even faster.

编码数据集

只需要利用tokenizer类以及已有的Merges Vocab即可

# tokenizer_experiment.py

import enum
from cs336_basics import tokenizer
from typing import IO, Any, BinaryIO
from tests import test_tokenizer
import random
import numpy as np
import multiprocessing as mp
from functools import partial

TinyStories_Vocab_Path = './../TinyStories_Result/vocab.json'
TinyStories_Merges_Path = './../TinyStories_Result/merges.txt'

OpenWebText_Vocab_Path = './../OpenWebText_Result/vocab.json'
OpenWebText_Merges_Path = './../OpenWebText_Result/merges.txt'

TinyStories_Datapath = './../data/TinyStoriesV2-GPT4-train.txt'
OpenWebText_Datapath = './../data/owt_train.txt'

TinyStories_Valid_Datapath = './../data/TinyStoriesV2-GPT4-valid.txt'
OpenWebText_Valid_Datapath = './../data/owt_valid.txt'


def sample_documents_from_file(filepath,num_samples=10):
    documents = []

    with open(filepath,'r',encoding='utf-8') as f:
        content = f.read()

    parts = content.split('<|endoftext|>')

    for part in parts:
        if part.strip():
            documents.append(part+'<|endoftext|>')

    
    if len(documents) <= num_samples:
        return documents

    return random.sample(documents,num_samples)

def all_documents_from_file(filepath):
    documents = []

    with open(filepath,'r',encoding='utf-8') as f:
        content = f.read()

    parts = content.split('<|endoftext|>')

    for part in parts:
        if part.strip():
            documents.append(part+'<|endoftext|>')

    return documents

def calculate_compression_ratio(text,tokenizer):
    original_bytes = len(text.encode('utf-8'))

    tokens = tokenizer.encode(text)
    num_tokens = len(tokens)

    compression_ratio = original_bytes / num_tokens if num_tokens > 0 else 0
    
    return compression_ratio

def encode_text(text,tokenizer):

    tokens = tokenizer.encode(text)

def encode_entire_file(filepath,tokenizer):

    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 一次性编码整个文件内容
    tokens = tokenizer.encode(content)
    return tokens

def encode_documents_batch(docs_batch,tokenizer):
    tokens = []
    for doc in docs_batch:
        doc_tokens = tokenizer.encode(doc)
        tokens.extend(doc_tokens)
    return tokens

def encode_documents_parallel(documents, tokenizer, num_processes=None):

    if num_processes is None:
        num_processes = min(mp.cpu_count(), 96)  # 使用最多32个进程
    
    # 将文档分成批次
    batch_size = max(1, len(documents) // num_processes)
    doc_batches = [documents[i:i + batch_size] for i in range(0, len(documents), batch_size)]
    
    # 创建编码函数（固定tokenizer参数）
    encode_func = partial(encode_documents_batch, tokenizer=tokenizer)
    
    print(f"Using {len(doc_batches)} processes to encode {len(documents)} documents...")
    
    # 使用进程池并行处理
    with mp.Pool(processes=len(doc_batches)) as pool:
        results = pool.map(encode_func, doc_batches)
    
    # 合并所有批次的结果
    all_tokens = []
    for batch_tokens in results:
        all_tokens.extend(batch_tokens)
    
    return all_tokens

if __name__ == "__main__":



    print("\nLoading TinyStories tokenizer...")
    tinystories_tokenizer = test_tokenizer.get_tokenizer_from_vocab_merges_path(
        vocab_path=TinyStories_Vocab_Path,
        merges_path=TinyStories_Merges_Path,
        special_tokens=["<|endoftext|>"]
    )

    print("Loading OpenWebText tokenizer...")
    openwebtext_tokenizer = test_tokenizer.get_tokenizer_from_vocab_merges_path(
        vocab_path=OpenWebText_Vocab_Path,
        merges_path=OpenWebText_Merges_Path,
        special_tokens=["<|endoftext|>"]
    )


    print("\nEncoding all TinyStories Dataset")
    
    print("Encoding TinyStories valid dataset...")
    TinyStories_Valid_docs = all_documents_from_file(TinyStories_Valid_Datapath)
    TinyStories_Valid_Encode = encode_documents_parallel(TinyStories_Valid_docs, tinystories_tokenizer)


    print("Encoding TinyStories train dataset...")
    TinyStories_Train_docs = all_documents_from_file(TinyStories_Datapath)
    TinyStories_Train_Encode = encode_documents_parallel(TinyStories_Train_docs, tinystories_tokenizer)
    

    print("\nEncoding all OpenWebText Dataset")
    

    print("Encoding OpenWebText valid dataset...")
    OpenWebText_Valid_docs = all_documents_from_file(OpenWebText_Valid_Datapath)
    OpenWebText_Valid_Encode = encode_documents_parallel(OpenWebText_Valid_docs, openwebtext_tokenizer)

    
    print("Encoding OpenWebText train dataset...")
    OpenWebText_Train_docs = all_documents_from_file(OpenWebText_Datapath)
    OpenWebText_Train_Encode = encode_documents_parallel(OpenWebText_Train_docs, openwebtext_tokenizer)
    


    print("\nSaving encoded datasets as uint16 NumPy arrays...")
    
    np.save('./../TinyStories_Result/train_tokens.npy', np.array(TinyStories_Train_Encode, dtype=np.uint16))
    np.save('./../TinyStories_Result/valid_tokens.npy', np.array(TinyStories_Valid_Encode, dtype=np.uint16))
    np.save('./../OpenWebText_Result/train_tokens.npy', np.array(OpenWebText_Train_Encode, dtype=np.uint16))
    np.save('./../OpenWebText_Result/valid_tokens.npy', np.array(OpenWebText_Valid_Encode, dtype=np.uint16))

    print("All datasets encoded and saved successfully!")
    print(f"TinyStories train tokens: {len(TinyStories_Train_Encode)}")
    print(f"TinyStories valid tokens: {len(TinyStories_Valid_Encode)}")
    print(f"OpenWebText train tokens: {len(OpenWebText_Train_Encode)}")
    print(f"OpenWebText valid tokens: {len(OpenWebText_Valid_Encode)}")

只需要注意最后以np.uint16格式保存即可