-
Notifications
You must be signed in to change notification settings - Fork 24
/
semseg-ptv2m2-0-base.py
160 lines (151 loc) · 6.61 KB
/
semseg-ptv2m2-0-base.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
_base_ = ['../_base_/default_runtime.py',
'../_base_/tests/segmentation.py']
# misc custom setting
batch_size = 12 # bs: total bs in all gpus
mix_prob = 0.8
empty_cache = False
enable_amp = True
# model settings
model = dict(
type="ptv2m2",
in_channels=9,
num_classes=20,
patch_embed_depth=1,
patch_embed_channels=48,
patch_embed_groups=6,
patch_embed_neighbours=8,
enc_depths=(2, 2, 6, 2),
enc_channels=(96, 192, 384, 512),
enc_groups=(12, 24, 48, 64),
enc_neighbours=(16, 16, 16, 16),
dec_depths=(1, 1, 1, 1),
dec_channels=(48, 96, 192, 384),
dec_groups=(6, 12, 24, 48),
dec_neighbours=(16, 16, 16, 16),
grid_sizes=(0.06, 0.15, 0.375, 0.9375), # x3, x2.5, x2.5, x2.5
attn_qkv_bias=True,
pe_multiplier=False,
pe_bias=True,
attn_drop_rate=0.,
drop_path_rate=0.3,
enable_checkpoint=False,
unpool_backend="map", # map / interp
)
# scheduler settings
epoch = 900
optimizer = dict(type='AdamW', lr=0.005, weight_decay=0.02)
scheduler = dict(type='OneCycleLR',
max_lr=optimizer["lr"],
pct_start=0.05,
anneal_strategy="cos",
div_factor=10.0,
final_div_factor=1000.0)
# dataset settings
dataset_type = "ScanNetDataset"
data_root = "data/scannet"
data = dict(
num_classes=20,
ignore_label=255,
names=["wall", "floor", "cabinet", "bed", "chair",
"sofa", "table", "door", "window", "bookshelf",
"picture", "counter", "desk", "curtain", "refridgerator",
"shower curtain", "toilet", "sink", "bathtub", "otherfurniture"],
train=dict(
type=dataset_type,
split="train",
data_root=data_root,
transform=[
dict(type="CenterShift", apply_z=True),
dict(type="RandomDropout", dropout_ratio=0.2, dropout_application_ratio=0.2),
# dict(type="RandomRotateTargetAngle", angle=(1/2, 1, 3/2), center=[0, 0, 0], axis='z', p=0.75),
dict(type="RandomRotate", angle=[-1, 1], axis='z', center=[0, 0, 0], p=0.5),
dict(type="RandomRotate", angle=[-1/64, 1/64], axis='x', p=0.5),
dict(type="RandomRotate", angle=[-1/64, 1/64], axis='y', p=0.5),
dict(type="RandomScale", scale=[0.9, 1.1]),
# dict(type="RandomShift", shift=[0.2, 0.2, 0.2]),
dict(type="RandomFlip", p=0.5),
dict(type="RandomJitter", sigma=0.005, clip=0.02),
dict(type="ElasticDistortion", distortion_params=[[0.2, 0.4], [0.8, 1.6]]),
dict(type="ChromaticAutoContrast", p=0.2, blend_factor=None),
dict(type="ChromaticTranslation", p=0.95, ratio=0.05),
dict(type="ChromaticJitter", p=0.95, std=0.05),
# dict(type="HueSaturationTranslation", hue_max=0.2, saturation_max=0.2),
# dict(type="RandomColorDrop", p=0.2, color_augment=0.0),
dict(type="Voxelize", voxel_size=0.02, hash_type='fnv', mode='train', return_min_coord=True),
dict(type="SphereCrop", point_max=100000, mode='random'),
dict(type="CenterShift", apply_z=False),
dict(type="NormalizeColor"),
dict(type="ShufflePoint"),
dict(type="ToTensor"),
dict(type="Collect", keys=("coord", "label"), feat_keys=("coord", "normal", "color"))
],
test_mode=False,
),
val=dict(
type=dataset_type,
split="val",
data_root=data_root,
transform=[
dict(type="CenterShift", apply_z=True),
dict(type="Voxelize", voxel_size=0.02, hash_type='fnv', mode='train', return_min_coord=True),
# dict(type="SphereCrop", point_max=1000000, mode='center'),
dict(type="CenterShift", apply_z=False),
dict(type="NormalizeColor"),
dict(type="ToTensor"),
dict(type="Collect", keys=("coord", "label"), feat_keys=("coord", "normal", "color"))
],
test_mode=False,
),
test=dict(
type=dataset_type,
split="val",
data_root=data_root,
transform=[
dict(type="CenterShift", apply_z=True),
dict(type="NormalizeColor"),
],
test_mode=True,
test_cfg=dict(
voxelize=dict(type="Voxelize",
voxel_size=0.02,
hash_type="fnv",
mode="test",
keys=("coord", "normal", "color")
),
crop=None,
post_transform=[
dict(type="CenterShift", apply_z=False),
dict(type="ToTensor"),
dict(type="Collect", keys=("coord", "index"), feat_keys=("coord", "normal", "color"))
],
aug_transform=[
[dict(type="RandomRotateTargetAngle", angle=[0], axis='z', center=[0, 0, 0], p=1)],
[dict(type="RandomRotateTargetAngle", angle=[1/2], axis='z', center=[0, 0, 0], p=1)],
[dict(type="RandomRotateTargetAngle", angle=[1], axis='z', center=[0, 0, 0], p=1)],
[dict(type="RandomRotateTargetAngle", angle=[3/2], axis='z', center=[0, 0, 0], p=1)],
[dict(type="RandomRotateTargetAngle", angle=[0], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[0.95, 0.95])],
[dict(type="RandomRotateTargetAngle", angle=[1 / 2], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[0.95, 0.95])],
[dict(type="RandomRotateTargetAngle", angle=[1], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[0.95, 0.95])],
[dict(type="RandomRotateTargetAngle", angle=[3 / 2], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[0.95, 0.95])],
[dict(type="RandomRotateTargetAngle", angle=[0], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[1.05, 1.05])],
[dict(type="RandomRotateTargetAngle", angle=[1 / 2], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[1.05, 1.05])],
[dict(type="RandomRotateTargetAngle", angle=[1], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[1.05, 1.05])],
[dict(type="RandomRotateTargetAngle", angle=[3 / 2], axis='z', center=[0, 0, 0], p=1),
dict(type="RandomScale", scale=[1.05, 1.05])],
[dict(type="RandomFlip", p=1)]
]
)
),
)
criteria = [
dict(type="CrossEntropyLoss",
loss_weight=1.0,
ignore_index=data["ignore_label"])
]