mirror of https://github.com/hpcaitech/ColossalAI
143 lines
2.9 KiB
Plaintext
143 lines
2.9 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import torch\n",
|
|
"import torch.nn as nn"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 23,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Linear(in_features=10, out_features=5, bias=False) 50\n",
|
|
"Linear(in_features=5, out_features=10, bias=False) 50\n",
|
|
"Linear(in_features=10, out_features=10, bias=False) 100\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"class Toy(nn.Module):\n",
|
|
" \n",
|
|
" def __init__(self):\n",
|
|
" super(Toy, self).__init__()\n",
|
|
" self.fc1 = nn.Linear(10,5, bias=False)\n",
|
|
" self.m3 = nn.Sequential(nn.Linear(5, 10, bias=False), nn.Linear(10,10, bias=False))\n",
|
|
"\n",
|
|
"t = Toy()\n",
|
|
"for mod in t.modules():\n",
|
|
" for p in mod.parameters(recurse=False):\n",
|
|
" print(mod, p.numel())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 24,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"torch.Size([5, 10]) 50\n",
|
|
"torch.Size([10, 5]) 50\n",
|
|
"torch.Size([10, 10]) 100\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"for p in t.parameters():\n",
|
|
" print(p.shape, p.numel())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 27,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"'224'"
|
|
]
|
|
},
|
|
"execution_count": 27,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"conf_str = torch.__config__.parallel_info()\n",
|
|
"inter_str = conf_str.split(\"hardware_concurrency() : \")[1]\n",
|
|
"max_concurrency = inter_str.split(\"\\n\")[0]\n",
|
|
"max_concurrency"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"0 0\n",
|
|
"0 1\n",
|
|
"0 2\n",
|
|
"1 0\n",
|
|
"1 1\n",
|
|
"1 2\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"for i in range(3):\n",
|
|
" for j in range(3):\n",
|
|
" print(i, j)\n",
|
|
" if i == 1 and j == 2:break\n",
|
|
" else:\n",
|
|
" continue\n",
|
|
" break"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "colossalai-py310",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.10.14"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|